XFESystem

using Mono.Cecil;
using Mono.Cecil.Cil;

namespace XFESystem.Tools.JintCleanupPatch;

public static partial class Program
{
    private const string RegexSearchMarker = "XFESystem.JintFailedRegexSuffixV1";
    private const string UncachedRegexSearch = "XfeExecuteWithTimeoutUncached";
    private static readonly string[] RegexSearchFields = ["_xfeNoMatchInput", "_xfeNoMatchEngine", "_xfeNoMatchStart", "_xfeNoMatchHits"];

    // An unsuccessful non-sticky search proves that the same immutable input
    // has no match at any later start. Keep that fact on the JsRegExp instance,
    // never on a shared AST or in a process-wide table. No JS results/captures
    // are cached. Engine identity invalidates RegExp.prototype.compile changes.
    private static void PatchRegexSearch(ModuleDefinition module)
    {
        if (module.GetType(RegexSearchMarker) is not null) { ValidateRegexSearch(module); return; }
        var regex = module.GetType("Jint.Native.JsRegExp");
        var prototype = module.GetType("Jint.Native.RegExp.RegExpPrototype");
        var original = prototype.Methods.Single(method => method.Name == "ExecuteWithTimeout");
        if (!original.IsStatic || !original.IsPrivate || original.Parameters.Count != 4 ||
            original.Parameters[0].ParameterType != regex || original.Parameters[2].ParameterType.MetadataType != MetadataType.String ||
            original.Parameters[3].ParameterType.MetadataType != MetadataType.Int32 ||
            original.ReturnType.FullName != "Jint.Runtime.RegExp.RegExpMatchResult" ||
            original.Body.Instructions.Count(instruction => instruction.Operand is MethodReference call &&
                call.Name == "Execute" && call.DeclaringType.FullName == "Jint.Runtime.RegExp.JintRegExpEngine") != 2 ||
            RegexSearchFields.Any(name => regex.Fields.Any(field => field.Name == name)))
            throw new InvalidOperationException("Pinned custom regex timeout/search shape changed.");
        var callers = Types(module.Types).SelectMany(type => type.Methods).Where(method => method.HasBody)
            .SelectMany(method => method.Body.Instructions).Where(instruction => instruction.Operand is MethodReference call && call.FullName == original.FullName).ToArray();
        if (callers.Length != 5) throw new InvalidOperationException("Pinned custom regex search call sites changed.");
        var fields = new[]
        {
            new FieldDefinition(RegexSearchFields[0], FieldAttributes.Assembly, module.TypeSystem.String),
            new FieldDefinition(RegexSearchFields[1], FieldAttributes.Assembly, original.Parameters[1].ParameterType),
            new FieldDefinition(RegexSearchFields[2], FieldAttributes.Assembly, module.TypeSystem.Int32),
            new FieldDefinition(RegexSearchFields[3], FieldAttributes.Assembly, module.TypeSystem.Int64)
        };
        foreach (var field in fields) regex.Fields.Add(field);
        var wrapper = new MethodDefinition(original.Name, original.Attributes, original.ReturnType);
        foreach (var parameter in original.Parameters) wrapper.Parameters.Add(new(parameter.Name, parameter.Attributes, parameter.ParameterType));
        original.Name = UncachedRegexSearch;
        prototype.Methods.Add(wrapper);
        foreach (var caller in callers) caller.Operand = wrapper;
        wrapper.Body.InitLocals = true;
        var result = new VariableDefinition(wrapper.ReturnType); wrapper.Body.Variables.Add(result);
        var match = module.GetType(wrapper.ReturnType.FullName);
        var success = match.Methods.Single(method => method.Name == "get_Success");
        var noMatch = match.Methods.Single(method => method.Name == "get_NoMatch");
        var sticky = regex.Methods.Single(method => method.Name == "get_Sticky");
        var length = module.ImportReference(typeof(string).GetProperty(nameof(string.Length))!.GetMethod!);
        var uncached = Instruction.Create(OpCodes.Ldarg_0);
        var done = Instruction.Create(OpCodes.Ldloc, result);
        var il = wrapper.Body.GetILProcessor();
        // Sticky failures say nothing about later positions. Bound retained input
        // to 128 KiB, matching the editor's code limit; larger inputs use Jint.
        il.Emit(OpCodes.Ldarg_0); il.Emit(OpCodes.Callvirt, sticky); il.Emit(OpCodes.Brtrue, uncached);
        il.Emit(OpCodes.Ldarg_2); il.Emit(OpCodes.Callvirt, length); il.Emit(OpCodes.Ldc_I4, 131072); il.Emit(OpCodes.Bgt, uncached);
        il.Emit(OpCodes.Ldarg_0); il.Emit(OpCodes.Ldfld, fields[0]); il.Emit(OpCodes.Ldarg_2); il.Emit(OpCodes.Ceq); il.Emit(OpCodes.Brfalse, uncached);
        il.Emit(OpCodes.Ldarg_0); il.Emit(OpCodes.Ldfld, fields[1]); il.Emit(OpCodes.Ldarg_1); il.Emit(OpCodes.Ceq); il.Emit(OpCodes.Brfalse, uncached);
        il.Emit(OpCodes.Ldarg_3); il.Emit(OpCodes.Ldarg_0); il.Emit(OpCodes.Ldfld, fields[2]); il.Emit(OpCodes.Blt, uncached);
        il.Emit(OpCodes.Ldarg_0); il.Emit(OpCodes.Ldarg_0); il.Emit(OpCodes.Ldfld, fields[3]);
        il.Emit(OpCodes.Ldc_I4_1); il.Emit(OpCodes.Conv_I8); il.Emit(OpCodes.Add); il.Emit(OpCodes.Stfld, fields[3]);
        il.Emit(OpCodes.Call, noMatch); il.Emit(OpCodes.Ret);
        il.Append(uncached); il.Emit(OpCodes.Ldarg_1); il.Emit(OpCodes.Ldarg_2); il.Emit(OpCodes.Ldarg_3);
        il.Emit(OpCodes.Call, original); il.Emit(OpCodes.Stloc, result);
        il.Emit(OpCodes.Ldloca, result); il.Emit(OpCodes.Call, success); il.Emit(OpCodes.Brtrue, done);
        il.Emit(OpCodes.Ldarg_0); il.Emit(OpCodes.Callvirt, sticky); il.Emit(OpCodes.Brtrue, done);
        il.Emit(OpCodes.Ldarg_2); il.Emit(OpCodes.Callvirt, length); il.Emit(OpCodes.Ldc_I4, 131072); il.Emit(OpCodes.Bgt, done);
        il.Emit(OpCodes.Ldarg_0); il.Emit(OpCodes.Ldarg_2); il.Emit(OpCodes.Stfld, fields[0]);
        il.Emit(OpCodes.Ldarg_0); il.Emit(OpCodes.Ldarg_1); il.Emit(OpCodes.Stfld, fields[1]);
        il.Emit(OpCodes.Ldarg_0); il.Emit(OpCodes.Ldarg_3); il.Emit(OpCodes.Stfld, fields[2]);
        il.Append(done); il.Emit(OpCodes.Ret);
        module.Types.Add(new("XFESystem", "JintFailedRegexSuffixV1", TypeAttributes.NotPublic | TypeAttributes.Abstract | TypeAttributes.Sealed, module.TypeSystem.Object));
        ValidateRegexSearch(module);
    }

    private static void ValidateRegexSearch(ModuleDefinition module)
    {
        var regex = module.GetType("Jint.Native.JsRegExp");
        var prototype = module.GetType("Jint.Native.RegExp.RegExpPrototype");
        var wrapper = prototype.Methods.SingleOrDefault(method => method.Name == "ExecuteWithTimeout");
        var original = prototype.Methods.SingleOrDefault(method => method.Name == UncachedRegexSearch);
        if (module.GetType(RegexSearchMarker) is null || original is null || wrapper is null ||
            wrapper.Body.ExceptionHandlers.Count != 0 || wrapper.Body.Instructions.Count != 57 ||
            RegexSearchFields.Any(name => regex.Fields.Count(field => field.Name == name && field.IsAssembly) != 1) ||
            wrapper.Body.Instructions.Count(instruction => instruction.Operand is MethodReference call && call.FullName == original.FullName) != 1 ||
            wrapper.Body.Instructions.Count(instruction => instruction.OpCode == OpCodes.Ldc_I4 && instruction.Operand is int value && value == 131072) != 2)
            throw new InvalidOperationException("The failed-regex-suffix patch is incomplete.");
    }
}