XFESystem

using Mono.Cecil;
using Mono.Cecil.Cil;
using XFESystem.Collections;

namespace XFESystem.Tools.JintCleanupPatch;

public static partial class Program
{
    private const string WeakMarker = "XFESystem.JintOwnedWeakCollectionsV1";
    private static void ValidateWeakCollections(ModuleDefinition module)
    {
        var engine=module.GetType("Jint.Engine");
        string owner=typeof(IWeakCollectionOwner).FullName!,scope=typeof(WeakCollectionScope).FullName!,table=typeof(OwnedWeakTable<,>).FullName!;
        var reference=module.AssemblyReferences.SingleOrDefault(item=>item.Name==typeof(WeakCollectionScope).Assembly.GetName().Name);
        if(reference is null||reference.Version!=typeof(WeakCollectionScope).Assembly.GetName().Version||
            engine.Interfaces.Count(item=>item.InterfaceType.FullName==owner)!=1||
            engine.Fields.Count(item=>item.Name=="_xfeWeakCollections"&&item.FieldType.FullName==scope)!=1||
            engine.Methods.Count(item=>item.Name=="get_WeakCollections"&&item.ReturnType.FullName==scope&&item.Overrides.Count==1)!=1)
            throw new InvalidOperationException("The per-engine weak storage owner binding is incomplete.");
        foreach(string name in new[]{"Jint.Native.JsWeakMap","Jint.Native.JsWeakSet"})
        {
            var type=module.GetType(name);var field=type?.Fields.SingleOrDefault(item=>item.Name=="_table");
            if(field?.FieldType is not GenericInstanceType generic||generic.ElementType.FullName!=table||
                generic.GenericArguments.Count!=2||generic.GenericArguments.Any(item=>item.FullName!="Jint.Native.JsValue"))
                throw new InvalidOperationException("The weak collection storage replacement is incomplete.");
            var calls=type!.Methods.Where(item=>item.HasBody).SelectMany(item=>item.Body.Instructions)
                .Select(item=>item.Operand).OfType<MethodReference>().ToArray();
            if(calls.Any(item=>item.DeclaringType.FullName.StartsWith("System.Runtime.CompilerServices.ConditionalWeakTable`2",StringComparison.Ordinal))||
                calls.Count(item=>item.DeclaringType.FullName==generic.FullName)!=(name.EndsWith("JsWeakMap",StringComparison.Ordinal)?11:4)||
                calls.Count(item=>item.DeclaringType.FullName==engine.FullName&&item.Name=="get_WeakCollections")!=1)
                throw new InvalidOperationException("The weak collection call-site replacement is incomplete.");
        }
    }
    private static void PatchWeakCollections(ModuleDefinition module)
    {
        var engine = module.GetType("Jint.Engine") ?? throw new InvalidOperationException("Pinned Engine is absent.");
        var scopeType = module.ImportReference(typeof(WeakCollectionScope));
        var ownerType = module.ImportReference(typeof(IWeakCollectionOwner));
        var tableType = module.ImportReference(typeof(OwnedWeakTable<,>));
        var scopeField = new FieldDefinition("_xfeWeakCollections", FieldAttributes.Private, scopeType);
        engine.Fields.Add(scopeField); engine.Interfaces.Add(new(ownerType));
        var getter = new MethodDefinition("get_WeakCollections", MethodAttributes.Public | MethodAttributes.Virtual |
            MethodAttributes.Final | MethodAttributes.NewSlot | MethodAttributes.HideBySig | MethodAttributes.SpecialName, scopeType);
        getter.Overrides.Add(module.ImportReference(typeof(IWeakCollectionOwner).GetProperty(nameof(IWeakCollectionOwner.WeakCollections))!.GetMethod!));
        var il = getter.Body.GetILProcessor(); var present = Instruction.Create(OpCodes.Ldarg_0);
        il.Append(Instruction.Create(OpCodes.Ldarg_0)); il.Append(Instruction.Create(OpCodes.Ldfld, scopeField));
        il.Append(Instruction.Create(OpCodes.Brtrue_S, present)); il.Append(Instruction.Create(OpCodes.Ldarg_0));
        il.Append(Instruction.Create(OpCodes.Newobj, module.ImportReference(typeof(WeakCollectionScope).GetConstructor(Type.EmptyTypes)!)));
        il.Append(Instruction.Create(OpCodes.Stfld, scopeField)); il.Append(present);
        il.Append(Instruction.Create(OpCodes.Ldfld, scopeField)); il.Append(Instruction.Create(OpCodes.Ret));
        engine.Methods.Add(getter); engine.Properties.Add(new PropertyDefinition("WeakCollections", PropertyAttributes.None, scopeType) { GetMethod = getter });

        foreach (string name in new[] { "Jint.Native.JsWeakMap", "Jint.Native.JsWeakSet" })
        {
            var type = module.GetType(name) ?? throw new InvalidOperationException("Pinned weak collection is absent: " + name);
            var field = type.Fields.Single(item => item.Name == "_table");
            if (field.FieldType is not GenericInstanceType old || old.ElementType.FullName != "System.Runtime.CompilerServices.ConditionalWeakTable`2" ||
                old.GenericArguments.Count != 2 || old.GenericArguments.Any(item => item.FullName != "Jint.Native.JsValue"))
                throw new InvalidOperationException("Pinned weak collection storage shape changed.");
            var replacement = new GenericInstanceType(tableType);
            foreach (var argument in old.GenericArguments) replacement.GenericArguments.Add(argument);
            field.FieldType = replacement;
            int constructors = 0, calls = 0;
            foreach (var method in type.Methods.Where(item => item.HasBody))
            {
                foreach (var instruction in method.Body.Instructions.ToArray())
                {
                    if (instruction.Operand is not MethodReference target || target.DeclaringType.FullName != old.FullName) continue;
                    if (target.Name is not (".ctor" or "TryGetValue" or "Remove" or "Add" or "AddOrUpdate"))
                        throw new InvalidOperationException("Unsupported original weak storage call: " + target.FullName);
                    if (target.Name == ".ctor")
                    {
                        if (method.Name != ".ctor" || method.Parameters.Count != 1 || method.Parameters[0].ParameterType.FullName != engine.FullName || target.Parameters.Count != 0)
                            throw new InvalidOperationException("Pinned weak constructor changed.");
                        var processor = method.Body.GetILProcessor();
                        processor.InsertBefore(instruction, Instruction.Create(OpCodes.Ldarg_1));
                        processor.InsertBefore(instruction, Instruction.Create(OpCodes.Callvirt, getter));
                        ++constructors;
                    }
                    var reference = new MethodReference(target.Name, target.ReturnType, replacement) { HasThis = true };
                    if (target.Name == ".ctor") reference.Parameters.Add(new(scopeType));
                    else foreach (var parameter in target.Parameters) reference.Parameters.Add(new(parameter.ParameterType));
                    instruction.Operand = reference; ++calls;
                }
            }
            if (constructors != 1 || calls != (name.EndsWith("JsWeakMap", StringComparison.Ordinal) ? 11 : 4))
                throw new InvalidOperationException($"Pinned weak storage call count changed: {name}: {constructors}/{calls}.");
        }
        module.Types.Add(new TypeDefinition("XFESystem", "JintOwnedWeakCollectionsV1", TypeAttributes.NotPublic | TypeAttributes.Abstract | TypeAttributes.Sealed, module.TypeSystem.Object));
        ValidateWeakCollections(module);
    }
}