XFESystem

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

internal static class EhRuntimePatch
{
    public const int Limit = 4096;
    public const int TraceLimit = 8;
    private const string HelperName = "XfeTraceFirstExceptionTypesV1";
    private const string CountName = "s_xfeEhTraceTypeCountV1";
    private const string TypesName = "s_xfeEhTraceTypesV1";
    private static readonly string[] TrackerNames = ["s_activeCatchFramePointers", "s_activeCatchHandlers", "s_activeCatchExInfos"];
    public static bool Apply(AssemblyDefinition assembly)
    {
        var module = assembly.MainModule;
        var type = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Runtime.ExceptionHandling.ExceptionHelper");
        var initializer = type.Methods.Single(method => method.IsConstructor && method.IsStatic && method.HasBody);
        var dispatch = type.Methods.Single(method => method.Name == "DispatchExceptionWithContext" && method.HasBody);
        var field = type.Fields.Single(field => field.Name == "s_activeCatchCount");
        var constant = type.Fields.Single(field => field.Name == "MaxActiveCatchDepth" && field.IsLiteral);
        var bound = dispatch.Body.Instructions.Where(instruction => instruction.OpCode == OpCodes.Ldsfld && instruction.Operand is FieldReference reference && reference.FullName == field.FullName &&
            instruction.Next?.OpCode.Code is Code.Ldc_I4 or Code.Ldc_I4_S && instruction.Next.Next.OpCode.Code is Code.Bge or Code.Bge_S).ToArray();
        if (bound.Length != 1) throw new InvalidOperationException("Pinned active catch comparison changed: " + string.Join("; ", dispatch.Body.Instructions.Where(instruction => instruction.Operand is FieldReference reference && reference.Name == field.Name).Select(instruction => instruction + " / " + instruction.Next + " / " + instruction.Next?.Next)));
        var allocations = new List<Instruction>();
        foreach (var name in TrackerNames)
        {
            var tracker = type.Fields.Single(field => field.Name == name && field.FieldType.FullName == "System.UIntPtr[]");
            var store = initializer.Body.Instructions.Single(instruction => instruction.OpCode == OpCodes.Stsfld && instruction.Operand is FieldReference reference && reference.FullName == tracker.FullName);
            if (store.Previous?.OpCode != OpCodes.Newarr || store.Previous.Operand is not TypeReference element || element.MetadataType != MetadataType.UIntPtr || store.Previous.Previous?.OpCode.Code is not (Code.Ldc_I4 or Code.Ldc_I4_S))
                throw new InvalidOperationException("Pinned active catch tracker allocation changed.");
            allocations.Add(store.Previous.Previous);
        }
        var failFast = type.Methods.Single(method => method.Name == "FailFast" && method.Parameters.Count == 2 && method.Parameters[0].ParameterType.MetadataType == MetadataType.String && method.Parameters[1].ParameterType.FullName == "System.Exception");
        var existing = type.Methods.Where(method => method.Name == HelperName).ToArray();
        if (existing.Length != 0)
        {
            if (existing.Length != 1 || !Equals(constant.Constant, Limit) || allocations.Any(instruction => Convert.ToInt32(instruction.Operand) != Limit) || Convert.ToInt32(bound[0].Next.Operand) != Limit)
                throw new InvalidOperationException("An unexpected active catch capacity patch already exists.");
            var failure = dispatch.Body.Instructions.SingleOrDefault(instruction => instruction.OpCode == OpCodes.Ldstr && Equals(instruction.Operand, "Nested catch tracking capacity exceeded"));
            if (failure?.Next.OpCode != OpCodes.Ldarg_0 || failure.Next.Next.Operand is not MethodReference call || call.FullName != failFast.FullName || failure.Next.Next.Next.OpCode != OpCodes.Ret ||
                dispatch.Body.Instructions.Count(instruction => instruction.Operand is MethodReference hook && hook.FullName == existing[0].FullName) != 1 ||
                type.Fields.Count(candidate => candidate.Name == CountName || candidate.Name == TypesName) != 2)
                throw new InvalidOperationException("The exceptional overflow/trace patch has an unexpected shape.");
            return true;
        }
        if (!Equals(constant.Constant, 64) || allocations.Any(instruction => Convert.ToInt32(instruction.Operand) != 64) || Convert.ToInt32(bound[0].Next.Operand) != 64)
            throw new InvalidOperationException("Pinned active catch limit changed; refusing to guess.");
        var warning = dispatch.Body.Instructions.Single(instruction => instruction.OpCode == OpCodes.Ldstr && Equals(instruction.Operand, "[EH] Too many nested catch funclets to track\n"));
        if (warning.Next.OpCode != OpCodes.Call || warning.Next.Operand is not MethodReference serial || serial.Name != "WriteString" || serial.DeclaringType.FullName != "Cosmos.Kernel.Core.IO.Serial")
            throw new InvalidOperationException("Pinned warning path changed.");
        var addressRead = dispatch.Body.Instructions.Single(instruction => instruction.OpCode == OpCodes.Ldfld && instruction.Operand is FieldReference reference && reference.Name == "HandlerAddress" && instruction.Next.OpCode == OpCodes.Stelem_I);
        if (addressRead.Previous.OpCode.Code is not (Code.Ldloc or Code.Ldloc_S) || addressRead.Previous.Operand is not VariableDefinition clauseVariable || clauseVariable.VariableType.FullName != "Cosmos.Kernel.Core.Runtime.ExceptionHandling.EHClause")
            throw new InvalidOperationException("Pinned caught-clause local changed.");
        var clause = module.Types.Single(candidate => candidate.FullName == clauseVariable.VariableType.FullName);
        var address = clause.Fields.Single(candidate => candidate.Name == "HandlerAddress");
        var trace = CreateTrace(module, type, initializer, field);
        var il = dispatch.Body.GetILProcessor();
        var first = il.Create(OpCodes.Ldarg_0);
        Redirect(dispatch.Body, bound[0], first);
        il.InsertBefore(bound[0], first); il.InsertBefore(bound[0], il.Create(OpCodes.Ldsfld, field));
        il.InsertBefore(bound[0], il.Create(OpCodes.Ldloca, clauseVariable)); il.InsertBefore(bound[0], il.Create(OpCodes.Ldfld, address));
        il.InsertBefore(bound[0], il.Create(OpCodes.Conv_U)); il.InsertBefore(bound[0], il.Create(OpCodes.Call, trace));
        constant.Constant = Limit; foreach (var instruction in allocations) { instruction.OpCode = OpCodes.Ldc_I4; instruction.Operand = Limit; } bound[0].Next.OpCode = OpCodes.Ldc_I4; bound[0].Next.Operand = Limit;
        warning.Operand = "Nested catch tracking capacity exceeded";
        var failureCall = warning.Next; il.InsertBefore(failureCall, il.Create(OpCodes.Ldarg_0)); failureCall.Operand = failFast;
        il.InsertAfter(failureCall, il.Create(OpCodes.Ret));
        Expand(dispatch.Body); dispatch.Body.MaxStackSize = Math.Max(8, dispatch.Body.MaxStackSize);
        return false;
    }

    private static MethodDefinition CreateTrace(ModuleDefinition module, TypeDefinition owner, MethodDefinition initializer, FieldDefinition catchCount)
    {
        var count = new FieldDefinition(CountName, FieldAttributes.Private | FieldAttributes.Static, module.TypeSystem.Int32);
        var types = new FieldDefinition(TypesName, FieldAttributes.Private | FieldAttributes.Static | FieldAttributes.InitOnly, new ArrayType(module.TypeSystem.UIntPtr));
        owner.Fields.Add(count); owner.Fields.Add(types);
        var init = initializer.Body.GetILProcessor(); var returnInstruction = initializer.Body.Instructions.Last(instruction => instruction.OpCode == OpCodes.Ret);
        var first = init.Create(OpCodes.Ldc_I4, TraceLimit); Redirect(initializer.Body, returnInstruction, first);
        init.InsertBefore(returnInstruction, first); init.InsertBefore(returnInstruction, init.Create(OpCodes.Newarr, module.TypeSystem.UIntPtr)); init.InsertBefore(returnInstruction, init.Create(OpCodes.Stsfld, types));
        var exception = owner.Methods.Single(method => method.Name == "FailFast" && method.Parameters.Count == 2).Parameters[1].ParameterType;
        var trace = new MethodDefinition(HelperName, MethodAttributes.Private | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.Void);
        trace.Parameters.Add(new("exception", ParameterAttributes.None, exception)); trace.Parameters.Add(new("depth", ParameterAttributes.None, module.TypeSystem.Int32)); trace.Parameters.Add(new("handler", ParameterAttributes.None, module.TypeSystem.UIntPtr));
        trace.Body.InitLocals = true; owner.Methods.Add(trace);
        var methodTable = new VariableDefinition(module.TypeSystem.UIntPtr); var index = new VariableDefinition(module.TypeSystem.Int32); trace.Body.Variables.Add(methodTable); trace.Body.Variables.Add(index);
        var unsafeCast = owner.Methods.Single(method => method.Name == "InvokeCatchHandler").Body.Instructions.Select(instruction => instruction.Operand).OfType<GenericInstanceMethod>()
            .Single(method => method.Name == "As" && method.GenericArguments.Count == 2 && method.GenericArguments[0].FullName == "System.Exception" && method.GenericArguments[1].MetadataType == MetadataType.IntPtr);
        var serial = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.IO.Serial");
        var write = serial.Methods.Single(method => method.Name == "WriteString" && method.Parameters.Count == 1 && method.Parameters[0].ParameterType.MetadataType == MetadataType.String);
        var number = serial.Methods.Single(method => method.Name == "WriteNumber" && method.Parameters.Count == 2 && method.Parameters[0].ParameterType.MetadataType == MetadataType.UInt64);
        var hex = serial.Methods.Single(method => method.Name == "WriteHex" && method.Parameters.Count == 1 && method.Parameters[0].ParameterType.MetadataType == MetadataType.UInt64);
        var il = trace.Body.GetILProcessor(); var end = il.Create(OpCodes.Ret); var check = il.Create(OpCodes.Ldloc, index); var record = il.Create(OpCodes.Ldsfld, types);
        il.Append(il.Create(OpCodes.Ldsfld, count)); il.Append(il.Create(OpCodes.Ldc_I4, TraceLimit)); il.Append(il.Create(OpCodes.Bge, end));
        il.Append(il.Create(OpCodes.Ldc_I4_0)); il.Append(il.Create(OpCodes.Stloc, index));
        // Read the object header through the existing strongly typed Unsafe.As conversion.
        // No reflection, exception Message, formatting allocation or GC is invoked here.
        il.Append(il.Create(OpCodes.Ldarga, trace.Parameters[0])); il.Append(il.Create(OpCodes.Call, unsafeCast)); il.Append(il.Create(OpCodes.Ldind_I)); il.Append(il.Create(OpCodes.Ldind_I));
        il.Append(il.Create(OpCodes.Ldc_I4, -2)); il.Append(il.Create(OpCodes.Conv_I)); il.Append(il.Create(OpCodes.And)); il.Append(il.Create(OpCodes.Stloc, methodTable));
        il.Append(check); il.Append(il.Create(OpCodes.Ldsfld, count)); il.Append(il.Create(OpCodes.Bge, record));
        il.Append(il.Create(OpCodes.Ldsfld, types)); il.Append(il.Create(OpCodes.Ldloc, index)); il.Append(il.Create(OpCodes.Ldelem_I)); il.Append(il.Create(OpCodes.Ldloc, methodTable)); il.Append(il.Create(OpCodes.Beq, end));
        il.Append(il.Create(OpCodes.Ldloc, index)); il.Append(il.Create(OpCodes.Ldc_I4_1)); il.Append(il.Create(OpCodes.Add)); il.Append(il.Create(OpCodes.Stloc, index)); il.Append(il.Create(OpCodes.Br, check));
        il.Append(record); il.Append(il.Create(OpCodes.Ldsfld, count)); il.Append(il.Create(OpCodes.Ldloc, methodTable)); il.Append(il.Create(OpCodes.Stelem_I));
        il.Append(il.Create(OpCodes.Ldsfld, count)); il.Append(il.Create(OpCodes.Ldc_I4_1)); il.Append(il.Create(OpCodes.Add)); il.Append(il.Create(OpCodes.Stsfld, count));
        Text("XFE_EH_FIRST_TYPE|EEType|0x"); il.Append(il.Create(OpCodes.Ldloc, methodTable)); il.Append(il.Create(OpCodes.Conv_U8)); il.Append(il.Create(OpCodes.Call, hex));
        Text("|Depth|"); il.Append(il.Create(OpCodes.Ldarg_1)); il.Append(il.Create(OpCodes.Conv_U8)); il.Append(il.Create(OpCodes.Ldc_I4_0)); il.Append(il.Create(OpCodes.Call, number));
        Text("|Handler|0x"); il.Append(il.Create(OpCodes.Ldarg_2)); il.Append(il.Create(OpCodes.Conv_U8)); il.Append(il.Create(OpCodes.Call, hex));
        Text("|Category|");
        var doneCategory = il.Create(OpCodes.Ldstr, "\n");
        foreach (var category in new[] { typeof(TimeoutException), typeof(OutOfMemoryException), typeof(InvalidCastException), typeof(StackOverflowException) })
        {
            var next = il.Create(OpCodes.Nop); il.Append(il.Create(OpCodes.Ldarg_0)); il.Append(il.Create(OpCodes.Isinst, module.ImportReference(category))); il.Append(il.Create(OpCodes.Brfalse, next));
            Text(category.Name); il.Append(il.Create(OpCodes.Br, doneCategory)); il.Append(next);
        }
        Text("Other"); il.Append(doneCategory); il.Append(il.Create(OpCodes.Call, write)); il.Append(end);
        trace.Body.MaxStackSize = 8;
        return trace;
        void Text(string text) { il.Append(il.Create(OpCodes.Ldstr, text)); il.Append(il.Create(OpCodes.Call, write)); }
    }
    private static void Redirect(MethodBody body, Instruction old, Instruction replacement)
    {
        foreach (var instruction in body.Instructions)
        {
            if (ReferenceEquals(instruction.Operand, old)) instruction.Operand = replacement;
            else if (instruction.Operand is Instruction[] targets) for (var index = 0; index < targets.Length; ++index) if (targets[index] == old) targets[index] = replacement;
        }
    }
    private static void Expand(MethodBody body)
    {
        foreach (var instruction in body.Instructions) instruction.OpCode = instruction.OpCode.Code switch
        {
            Code.Br_S => OpCodes.Br, Code.Brfalse_S => OpCodes.Brfalse, Code.Brtrue_S => OpCodes.Brtrue, Code.Bge_S => OpCodes.Bge,
            Code.Bge_Un_S => OpCodes.Bge_Un, Code.Bgt_S => OpCodes.Bgt, Code.Bgt_Un_S => OpCodes.Bgt_Un, Code.Ble_S => OpCodes.Ble,
            Code.Ble_Un_S => OpCodes.Ble_Un, Code.Blt_S => OpCodes.Blt, Code.Blt_Un_S => OpCodes.Blt_Un, Code.Beq_S => OpCodes.Beq,
            Code.Bne_Un_S => OpCodes.Bne_Un, Code.Leave_S => OpCodes.Leave, _ => instruction.OpCode
        };
    }
}