XFESystem

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

namespace XFESystem.Tools.JintCleanupPatch;

public static partial class Program
{
    private const string EngineWeakMarker = "XFESystem.JintOwnedEngineWeakCollectionsV1";
    private const string OriginalWeakTable = "System.Runtime.CompilerServices.ConditionalWeakTable`2";
    private const string PrivateDictionary = "System.Collections.Generic.Dictionary`2<Jint.Native.PrivateName,Jint.Native.Object.PrivateElement>";
    private const string PrivateFactory = "XfeCreatePrivateElementDictionary";

    private static MethodDefinition MainEngineConstructor(ModuleDefinition module)
    {
        var engine = module.GetType("Jint.Engine");
        var constructors = engine.Methods.Where(method => method.IsConstructor && !method.IsStatic).ToArray();
        var main = constructors.Single(method => method.Parameters.Count == 2 &&
            method.Parameters[0].ParameterType.FullName == "Jint.Options" &&
            method.Parameters[1].ParameterType.FullName == "System.Action`2<Jint.Engine,Jint.Options>");
        if (constructors.Length != 5 || constructors.Where(method => method != main).Any(method =>
            method.Body.Instructions.Count(instruction => instruction.OpCode == OpCodes.Call &&
                instruction.Operand is MethodReference target && target.FullName == main.FullName) != 1))
            throw new InvalidOperationException("Pinned Engine constructor forwarding changed.");
        return main;
    }

    private static void ValidateEngineWeakCollections(ModuleDefinition module)
    {
        ValidateWeakCollections(module);
        var engine = module.GetType("Jint.Engine");
        string owned = typeof(OwnedWeakTable<,>).FullName!;
        foreach (string name in new[] { "_emptyShapes", "_privateElementStore" })
        {
            var field = engine.Fields.Single(item => item.Name == name);
            string value = name == "_emptyShapes" ? "Jint.Native.Object.Shape" : PrivateDictionary;
            if (field.FieldType is not GenericInstanceType generic || generic.ElementType.FullName != owned ||
                generic.GenericArguments.Count != 2 || generic.GenericArguments[0].FullName != "Jint.Native.Object.ObjectInstance" ||
                generic.GenericArguments[1].FullName != value)
                throw new InvalidOperationException("Engine weak cache field replacement is incomplete: " + name);
            var calls = Types(module.Types).SelectMany(type => type.Methods).Where(method => method.HasBody)
                .SelectMany(method => method.Body.Instructions).Select(instruction => instruction.Operand)
                .OfType<MethodReference>().Where(target => target.DeclaringType.FullName == generic.FullName).ToArray();
            int constructors = name == "_emptyShapes" ? 1 : 2;
            if (calls.Count(target => target.Name == ".ctor" && target.Parameters.Count == 1 &&
                target.Parameters[0].ParameterType.FullName == typeof(WeakCollectionScope).FullName) != constructors ||
                calls.Length != (name == "_emptyShapes" ? 2 : 5) ||
                calls.Count(target => target.Name == (name == "_emptyShapes" ? "GetValue" : "GetOrCreateValue") && target.Parameters.Count == 2) != constructors)
                throw new InvalidOperationException("Engine weak cache call-site replacement is incomplete: " + name);
            foreach (var method in Types(module.Types).SelectMany(type => type.Methods).Where(method => method.HasBody))
            {
                if (method.Body.Variables.Any(local => local.VariableType is GenericInstanceType oldLocal &&
                    oldLocal.ElementType.FullName == OriginalWeakTable && oldLocal.GenericArguments.Count == 2 &&
                    oldLocal.GenericArguments[0].FullName == "Jint.Native.Object.ObjectInstance" && oldLocal.GenericArguments[1].FullName == value))
                    throw new InvalidOperationException("An original Engine weak cache local remains: " + name);
                foreach (var instruction in method.Body.Instructions.Where(instruction => instruction.OpCode == OpCodes.Newobj &&
                    instruction.Operand is MethodReference target && target.DeclaringType.FullName == generic.FullName))
                {
                    var ownerCall = instruction.Previous;
                    if (ownerCall?.OpCode != OpCodes.Callvirt || ownerCall.Operand is not MethodReference owner ||
                        owner.DeclaringType.FullName != engine.FullName || owner.Name != "get_WeakCollections" ||
                        (name == "_emptyShapes" ? ownerCall.Previous?.OpCode != OpCodes.Ldarg_0 :
                            ownerCall.Previous?.OpCode != OpCodes.Ldfld || ownerCall.Previous.Operand is not FieldReference ownerField ||
                            ownerField.DeclaringType.FullName != "Jint.Native.Object.ObjectInstance" || ownerField.Name != "_engine" ||
                            ownerCall.Previous.Previous?.OpCode != OpCodes.Ldarg_0))
                        throw new InvalidOperationException("The Engine weak cache owner is not the allocating engine: " + name);
                }
            }
            if (Types(module.Types).SelectMany(type => type.Methods).Where(method => method.HasBody)
                .SelectMany(method => method.Body.Instructions).Select(instruction => instruction.Operand)
                .OfType<MethodReference>().Any(target => target.DeclaringType is GenericInstanceType old &&
                    old.ElementType.FullName == OriginalWeakTable && old.GenericArguments.Count == 2 &&
                    old.GenericArguments[0].FullName == "Jint.Native.Object.ObjectInstance" && old.GenericArguments[1].FullName == value))
                throw new InvalidOperationException("An original Engine weak cache call remains: " + name);
        }
        var factory = engine.Methods.SingleOrDefault(method => method.Name == PrivateFactory);
        if (factory is null || !factory.IsStatic || !factory.IsAssembly || factory.ReturnType.FullName != PrivateDictionary ||
            factory.Parameters.Count != 0 || factory.Body.Instructions.Count != 2 ||
            factory.Body.Instructions[0].OpCode != OpCodes.Newobj ||
            factory.Body.Instructions[0].Operand is not MethodReference constructor || constructor.Parameters.Count != 0 ||
            constructor.DeclaringType.FullName != PrivateDictionary || factory.Body.Instructions[1].OpCode != OpCodes.Ret)
            throw new InvalidOperationException("The private-element dictionary factory is incomplete.");
        var objectType = module.GetType("Jint.Native.Object.ObjectInstance");
        foreach (string name in new[] { "PrivateMethodOrAccessorAdd", "PrivateFieldAdd" })
        {
            var method = objectType.Methods.Single(item => item.Name == name);
            if (method.Body.Instructions.Count(instruction => instruction.OpCode == OpCodes.Ldftn &&
                instruction.Operand is MethodReference target && target.FullName == factory.FullName) != 1 ||
                method.Body.Instructions.Count(instruction => instruction.Operand is MethodReference target &&
                    target.DeclaringType.FullName == engine.FullName && target.Name == "get_WeakCollections") != 1)
                throw new InvalidOperationException("Private-element factory/owner binding is incomplete.");
        }
        var main = MainEngineConstructor(module);
        var clauses = main.Body.ExceptionHandlers;
        if (clauses.Count != 1 || clauses[0].HandlerType != ExceptionHandlerType.Catch || clauses[0].CatchType.FullName != "System.Exception")
            throw new InvalidOperationException("The Engine constructor ownership handler is incomplete.");
        var handler = clauses[0];
        var instructions = main.Body.Instructions;
        int start = instructions.IndexOf(handler.HandlerStart), end = instructions.IndexOf(handler.HandlerEnd);
        var body = instructions.Skip(start).Take(end - start).ToArray();
        var baseCall = instructions.Single(instruction => instruction.OpCode == OpCodes.Call &&
            instruction.Operand is MethodReference target && target.DeclaringType.FullName == "System.Object" && target.Name == ".ctor");
        if (handler.TryStart != baseCall.Next || handler.TryEnd != handler.HandlerStart || handler.HandlerEnd != instructions[^1] ||
            body.Length != 8 || body[0].OpCode != OpCodes.Pop || body[^1].OpCode != OpCodes.Rethrow ||
            body[1].OpCode != OpCodes.Ldarg_0 || body[2].OpCode != OpCodes.Ldfld ||
            body[4].OpCode != OpCodes.Ldarg_0 || body[5].OpCode != OpCodes.Ldfld ||
            instructions[^2].OpCode != OpCodes.Rethrow || instructions[^1].OpCode != OpCodes.Ret ||
            instructions.Count(instruction => instruction.OpCode == OpCodes.Leave && instruction.Operand == instructions[^1]) != 1 ||
            body.Count(instruction => instruction.Operand is FieldReference field && field.Name == "_xfeWeakCollections") != 2 ||
            body.Count(instruction => instruction.OpCode == OpCodes.Callvirt && instruction.Operand is MethodReference target &&
                target.DeclaringType.FullName == typeof(WeakCollectionScope).FullName && target.Name == "Release") != 1 ||
            body.Any(instruction => instruction.Operand is MethodReference target && target.Name == "get_WeakCollections") ||
            body[3].OpCode != OpCodes.Brfalse || body[3].Operand != body[^1])
            throw new InvalidOperationException("The Engine constructor ownership cleanup changed.");
    }

    private static void PatchEngineWeakCollections(ModuleDefinition module)
    {
        var engine = module.GetType("Jint.Engine");
        var getter = engine.Methods.Single(method => method.Name == "get_WeakCollections");
        var main = MainEngineConstructor(module);
        if (main.Body.ExceptionHandlers.Count != 0 || main.Body.Instructions.Count(instruction => instruction.OpCode == OpCodes.Ret) != 1 ||
            main.Body.Instructions[^1].OpCode != OpCodes.Ret || engine.Methods.Any(method => method.Name == PrivateFactory))
            throw new InvalidOperationException("Pinned Engine constructor/factory shape changed.");
        var privateField = engine.Fields.Single(field => field.Name == "_privateElementStore");
        if (privateField.FieldType is not GenericInstanceType oldPrivate || oldPrivate.ElementType.FullName != OriginalWeakTable ||
            oldPrivate.GenericArguments.Count != 2 || oldPrivate.GenericArguments[1].FullName != PrivateDictionary)
            throw new InvalidOperationException("Pinned private-element storage shape changed.");
        var dictionary = oldPrivate.GenericArguments[1];
        var factory = new MethodDefinition(PrivateFactory, MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, dictionary);
        var dictionaryConstructor = new MethodReference(".ctor", module.TypeSystem.Void, dictionary) { HasThis = true };
        factory.Body.Instructions.Add(Instruction.Create(OpCodes.Newobj, dictionaryConstructor));
        factory.Body.Instructions.Add(Instruction.Create(OpCodes.Ret)); engine.Methods.Add(factory);
        var func = new GenericInstanceType(module.ImportReference(typeof(Func<>))); func.GenericArguments.Add(dictionary);
        var funcConstructor = new MethodReference(".ctor", module.TypeSystem.Void, func) { HasThis = true };
        funcConstructor.Parameters.Add(new(module.TypeSystem.Object)); funcConstructor.Parameters.Add(new(module.TypeSystem.IntPtr));
        var objectType = module.GetType("Jint.Native.Object.ObjectInstance");
        var objectEngine = objectType.Fields.Single(field => field.Name == "_engine" && field.FieldType.FullName == engine.FullName);

        foreach (string name in new[] { "_emptyShapes", "_privateElementStore" })
        {
            var field = engine.Fields.Single(item => item.Name == name);
            if (field.FieldType is not GenericInstanceType old || old.ElementType.FullName != OriginalWeakTable ||
                old.GenericArguments.Count != 2 || old.GenericArguments[0].FullName != objectType.FullName ||
                old.GenericArguments[1].FullName != (name == "_emptyShapes" ? "Jint.Native.Object.Shape" : PrivateDictionary))
                throw new InvalidOperationException("Pinned Engine weak cache shape changed: " + name);
            var replacement = new GenericInstanceType(module.ImportReference(typeof(OwnedWeakTable<,>)));
            foreach (var argument in old.GenericArguments) replacement.GenericArguments.Add(argument);
            int constructors = 0, calls = 0;
            foreach (var method in Types(module.Types).SelectMany(type => type.Methods).Where(method => method.HasBody))
            {
                bool edited = false;
                foreach (var local in method.Body.Variables.Where(local => local.VariableType.FullName == old.FullName)) { local.VariableType = replacement; edited = true; }
                foreach (var instruction in method.Body.Instructions.ToArray())
                {
                    if (instruction.Operand is FieldReference reference && reference.DeclaringType.FullName == engine.FullName && reference.Name == name)
                    { instruction.Operand = field; edited = true; }
                    if (instruction.Operand is not MethodReference target || target.DeclaringType.FullName != old.FullName) continue;
                    bool empty = name == "_emptyShapes";
                    bool allowed = empty ? method.DeclaringType == engine && method.Name == "GetEmptyShape" && target.Name is ".ctor" or "GetValue" :
                        method.DeclaringType == objectType && (method.Name is "PrivateMethodOrAccessorAdd" or "PrivateFieldAdd" && target.Name is ".ctor" or "GetOrCreateValue" ||
                            method.Name == "PrivateElementFind" && target.Name == "TryGetValue");
                    if (!allowed) throw new InvalidOperationException("Unexpected Engine weak storage call: " + method.FullName + " / " + target.FullName);
                    var processor = method.Body.GetILProcessor();
                    if (target.Name == ".ctor")
                    {
                        if (target.Parameters.Count != 0) throw new InvalidOperationException("Pinned Engine weak constructor changed.");
                        processor.InsertBefore(instruction, Instruction.Create(OpCodes.Ldarg_0));
                        if (!empty) processor.InsertBefore(instruction, Instruction.Create(OpCodes.Ldfld, objectEngine));
                        processor.InsertBefore(instruction, Instruction.Create(OpCodes.Callvirt, getter)); ++constructors;
                    }
                    var call = new MethodReference(target.Name, target.ReturnType, replacement) { HasThis = true };
                    foreach (var parameter in target.Parameters) call.Parameters.Add(new(parameter.ParameterType));
                    if (target.Name == ".ctor") call.Parameters.Add(new(module.ImportReference(typeof(WeakCollectionScope))));
                    if (target.Name == "GetOrCreateValue")
                    {
                        processor.InsertBefore(instruction, Instruction.Create(OpCodes.Ldnull));
                        processor.InsertBefore(instruction, Instruction.Create(OpCodes.Ldftn, factory));
                        processor.InsertBefore(instruction, Instruction.Create(OpCodes.Newobj, funcConstructor));
                        var openFunc = new GenericInstanceType(module.ImportReference(typeof(Func<>)));
                        openFunc.GenericArguments.Add(replacement.ElementType.GenericParameters[1]);
                        call.Parameters.Add(new(openFunc));
                    }
                    instruction.Operand = call; ++calls; edited = true;
                }
                if (edited) ExpandBranches(method.Body);
            }
            field.FieldType = replacement;
            if (constructors != (name == "_emptyShapes" ? 1 : 2) || calls != (name == "_emptyShapes" ? 2 : 5))
                throw new InvalidOperationException("Pinned Engine weak storage call count changed: " + name);
        }
        AddConstructorOwnershipHandler(module, main, engine.Fields.Single(field => field.Name == "_xfeWeakCollections"));
        module.Types.Add(new TypeDefinition("XFESystem", "JintOwnedEngineWeakCollectionsV1", TypeAttributes.NotPublic | TypeAttributes.Abstract | TypeAttributes.Sealed, module.TypeSystem.Object));
        ValidateEngineWeakCollections(module);
    }

    private static void AddConstructorOwnershipHandler(ModuleDefinition module, MethodDefinition constructor, FieldDefinition scope)
    {
        var body = constructor.Body; ExpandBranches(body);
        var baseCall = body.Instructions.Single(instruction => instruction.OpCode == OpCodes.Call &&
            instruction.Operand is MethodReference target && target.DeclaringType.FullName == "System.Object" && target.Name == ".ctor");
        var processor = body.GetILProcessor(); var done = Instruction.Create(OpCodes.Ret);
        var originalReturn = body.Instructions[^1]; originalReturn.OpCode = OpCodes.Leave; originalReturn.Operand = done;
        var start = Instruction.Create(OpCodes.Pop); var rethrow = Instruction.Create(OpCodes.Rethrow);
        processor.Append(start); processor.Append(Instruction.Create(OpCodes.Ldarg_0)); processor.Append(Instruction.Create(OpCodes.Ldfld, scope));
        processor.Append(Instruction.Create(OpCodes.Brfalse, rethrow)); processor.Append(Instruction.Create(OpCodes.Ldarg_0));
        processor.Append(Instruction.Create(OpCodes.Ldfld, scope));
        processor.Append(Instruction.Create(OpCodes.Callvirt, module.ImportReference(typeof(WeakCollectionScope).GetMethod(nameof(WeakCollectionScope.Release))!)));
        processor.Append(rethrow); processor.Append(done);
        body.ExceptionHandlers.Add(new(ExceptionHandlerType.Catch) { CatchType = module.ImportReference(typeof(Exception)),
            TryStart = baseCall.Next, TryEnd = start, HandlerStart = start, HandlerEnd = done });
        body.MaxStackSize = Math.Max(8, body.MaxStackSize);
    }
}