ClientModLoader/Plugin.ClientModLoader/ModInjector.cs

using System.Collections.Immutable;
using System.Reflection.Emit;
using HarmonyLib;
using Plugin.ClientModLoader.Utils;
using Sandbox.Definitions;
using Sandbox.Engine.Networking;
using Sandbox.Game.World;
using VRage.Game;

namespace Plugin.ClientModLoader;

[HarmonyPatch]
internal static class ModInjector
{
    public static HashSet<ulong> Mods = [];
    private static readonly List<MyObjectBuilder_Checkpoint.ModItem> AdditionalFilledModItems = [];

    [HarmonyPatch(typeof(MyWorkshop), nameof(MyWorkshop.DownloadWorldModsBlockingInternal))]
    [HarmonyPrefix]
    private static void DownloadModsBlockingPrefix(ref List<MyObjectBuilder_Checkpoint.ModItem> mods, ref List<MyObjectBuilder_Checkpoint.ModItem> __state)
    {
        AdditionalFilledModItems.Clear();
        __state = mods;
        AppendToList(ref mods);
    }

    [HarmonyPatch(typeof(MyWorkshop), nameof(MyWorkshop.DownloadWorldModsBlockingInternal))]
    [HarmonyPostfix]
    private static void DownloadModsBlockingPostfix(MyWorkshop.ResultData __result, List<MyObjectBuilder_Checkpoint.ModItem> mods, List<MyObjectBuilder_Checkpoint.ModItem> __state)
    {
        if (__result.Result != VRage.GameServices.MyGameServiceCallResult.OK)
            return; //world will not load, and mod data isn't loaded

        var worldMods = __state.Select(b => b.PublishedFileId).ToImmutableHashSet();
        var resolvedMods = mods.ToImmutableDictionary(b => b.PublishedFileId);

        // list of selected mods which are resolved
        var requestedMods = mods.IntersectBy(Mods, b => b.PublishedFileId).ExceptBy(worldMods, b => b.PublishedFileId)
            .ToDictionary(b => b.PublishedFileId);

        // add dependencies of requested mods
        // but skip if those are also requested by world we're loading in
        foreach (var dependency in requestedMods.Values
                     .SelectMany(b => b.GetModData().Dependencies)
                     .Distinct().ToArray())
        {
            if (worldMods.Contains(dependency))
                continue;

            requestedMods.TryAdd(dependency, resolvedMods[dependency]);
        }

        // add resolved client mods and their exclusive dependencies
        AdditionalFilledModItems.AddRange(requestedMods.Values);

        // upsert world mods by resolved ones excluding our client ones and their exclusive dependencies
        // so world mods is only populated by dependencies of original world mods
        foreach (var mod in mods.ExceptBy(requestedMods.Keys, b => b.PublishedFileId))
        {
            var index = __state.FindIndex(b =>
                b.PublishedFileId == mod.PublishedFileId && b.PublishedServiceName == mod.PublishedServiceName);

            if (index != -1)
            {
                var stateMod = __state[index];
                stateMod.SetModData(mod.GetModData());
                __state[index] = stateMod;
            }
            else
                __state.Add(mod);
        }
    }

    [HarmonyPatch(typeof(MyWorkshop), nameof(MyWorkshop.DownloadModsAsync))]
    [HarmonyTranspiler]
    private static IEnumerable<CodeInstruction> DownloadModsAsyncTranspiler(IEnumerable<CodeInstruction> instructions)
    {
        var getCount = AccessTools.PropertyGetter(typeof(List<MyObjectBuilder_Checkpoint.ModItem>), nameof(List<MyObjectBuilder_Checkpoint.ModItem>.Count));
        return new CodeMatcher(instructions)
            .SearchForward(b => b.Calls(getCount))
            .Advance(1)
            .Insert(
                CodeInstruction.LoadField(typeof(ModInjector), nameof(Mods)),
                new(OpCodes.Callvirt, AccessTools.PropertyGetter(typeof(HashSet<ulong>), nameof(HashSet<ulong>.Count))),
                new(OpCodes.Add)
            )
            .Instructions();
    }

    [HarmonyPatch(typeof(MyDefinitionManager), nameof(MyDefinitionManager.LoadData))]
    [HarmonyPrefix]
    private static void LoadDefinitionsPrefix(ref List<MyObjectBuilder_Checkpoint.ModItem> mods) => AppendToList(ref mods);

    [HarmonyPatch(typeof(MyScriptManager), nameof(MyScriptManager.LoadData))]
    [HarmonyTranspiler]
    private static IEnumerable<CodeInstruction> LoadScriptsTranspiler(IEnumerable<CodeInstruction> instructions, ILGenerator generator)
    {
        var staticGetter = AccessTools.PropertyGetter(typeof(MySession), nameof(MySession.Static));
        var modsGetter = AccessTools.PropertyGetter(typeof(MySession), nameof(MySession.Mods));
        return new CodeMatcher(instructions, generator)
            .Start()
            .DeclareLocal(typeof(List<MyObjectBuilder_Checkpoint.ModItem>), out var modsLocal)
            .CreateLabel(out var start)
            .InsertAndAdvance(
                new(OpCodes.Call, staticGetter),
                new(OpCodes.Call, modsGetter),
                new(OpCodes.Stloc, modsLocal),
                new(OpCodes.Ldloc, modsLocal),
                new(OpCodes.Brfalse, start),
                new(OpCodes.Ldloca, modsLocal),
                CodeInstruction.Call(typeof(ModInjector), nameof(AppendToList))
            )
            .MatchStartForward(CodeMatch.Calls(staticGetter), CodeMatch.Calls(modsGetter))
            .Repeat(a => a
                .SetAndAdvance(OpCodes.Ldloc, modsLocal)
                .SetAndAdvance(OpCodes.Nop, null))
            .Instructions();
    }

    private static void AppendToList(ref List<MyObjectBuilder_Checkpoint.ModItem> mods)
    {
        // copy
        mods = [.. mods];

        if (AdditionalFilledModItems.Count > 0)
        {
            mods.AddRange(AdditionalFilledModItems);
            return;
        }

        foreach (var mod in Mods)
        {
            //avoid duplicates
            if (mods.Exists(m => m.PublishedFileId == mod))
                continue;

            mods.Add(new MyObjectBuilder_Checkpoint.ModItem(mod, "Steam"));
        }
    }
}