XFESystem

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

/// <summary>
/// Single-CPU context virtualization for the pinned runtime's otherwise global
/// exception dispatcher. All storage is allocated before a thread is runnable;
/// the interrupt-side switch contains field copies and two tiny native calls.
/// </summary>
internal static class ThreadEhRuntimePatch
{
    public const string SwitchName = "XfeSwitchThreadEhV1";
    public const string FieldPrefix = "_xfeThreadEhV1_";
    private static readonly string[] StateNames = ["s_isHandlingException", "s_activeCatchFramePointers", "s_activeCatchHandlers", "s_activeCatchExInfos", "s_activeCatchCount", "s_xfeEhTraceTypeCountV1", "s_xfeEhTraceTypesV1"];

    public static bool Apply(AssemblyDefinition assembly)
    {
        var module = assembly.MainModule;
        var eh = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Runtime.ExceptionHandling.ExceptionHelper");
        var thread = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Scheduler.SchedulerThread");
        var cpu = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Scheduler.PerCpuState");
        var setter = cpu.Methods.Single(method => method.Name == "set_CurrentThread");
        var current = cpu.Fields.Single(field => field.Name == "<CurrentThread>k__BackingField");
        var constructor = thread.Methods.Single(method => method.IsConstructor && !method.IsStatic);
        var originals = StateNames.Select(name => eh.Fields.Single(field => field.Name == name)).ToArray();
        var existing = eh.Methods.SingleOrDefault(method => method.Name == SwitchName);
        if (existing is not null)
        {
            if (thread.Fields.Count(field => field.Name.StartsWith(FieldPrefix, StringComparison.Ordinal)) != StateNames.Length + 1 ||
                setter.Body.Instructions.Count(instruction => instruction.Operand is MethodReference method && method.FullName == existing.FullName) != 1 ||
                existing.Body.Instructions.Any(instruction => instruction.OpCode.Code is Code.Newobj or Code.Newarr or Code.Throw) ||
                originals.Any(field => field.IsInitOnly) ||
                constructor.Body.Instructions.Count(instruction => instruction.OpCode == OpCodes.Newarr) != 4 ||
                eh.Methods.Count(method => method.Name == "XfeThreadEhContractVersionV1" || method.Name == "XfeCurrentSchedulerThreadIdV1") != 2)
                throw new InvalidOperationException("Unexpected per-thread EH patch shape.");
            return true;
        }
        if (constructor.Body.Instructions.Count != 3 || constructor.Body.Instructions[^1].OpCode != OpCodes.Ret ||
            setter.Body.Instructions.Count != 4 || setter.Body.Instructions[2].OpCode != OpCodes.Stfld ||
            setter.Body.Instructions[2].Operand is not FieldReference stored || stored.FullName != current.FullName)
            throw new InvalidOperationException("Pinned scheduler constructor/current-thread setter changed; refusing to guess.");

        var state = originals.Select(field => new FieldDefinition(FieldPrefix + field.Name, FieldAttributes.Assembly, field.FieldType)).ToArray();
        foreach (var field in state) thread.Fields.Add(field);
        var nativeHead = new FieldDefinition(FieldPrefix + "NativeExInfoHead", FieldAttributes.Assembly, module.TypeSystem.UIntPtr);
        thread.Fields.Add(nativeHead);
        // The NativeAOT optimizer must not treat the currently installed arrays
        // as immutable globals: they change at each actual thread switch.
        foreach (var field in originals) field.Attributes &= ~FieldAttributes.InitOnly;

        var ctorIl = constructor.Body.GetILProcessor();
        var ctorReturn = constructor.Body.Instructions[^1];
        for (int index = 0; index < state.Length; index++)
            if (state[index].FieldType is ArrayType array)
            {
                ctorIl.InsertBefore(ctorReturn, ctorIl.Create(OpCodes.Ldarg_0));
                ctorIl.InsertBefore(ctorReturn, ctorIl.Create(OpCodes.Ldc_I4, StateNames[index] == "s_xfeEhTraceTypesV1" ? EhRuntimePatch.TraceLimit : EhRuntimePatch.Limit));
                ctorIl.InsertBefore(ctorReturn, ctorIl.Create(OpCodes.Newarr, array.ElementType));
                ctorIl.InsertBefore(ctorReturn, ctorIl.Create(OpCodes.Stfld, state[index]));
            }
        constructor.Body.MaxStackSize = 2;

        var getHead = Import("XfeGetNativeExInfoHeadV1", "XfeRuntimeGetExInfoHead", module.TypeSystem.UIntPtr);
        var setHead = Import("XfeSetNativeExInfoHeadV1", "XfeRuntimeSetExInfoHead", module.TypeSystem.Void);
        setHead.Parameters.Add(new("value", ParameterAttributes.None, module.TypeSystem.UIntPtr));
        var change = new MethodDefinition(SwitchName, MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.Void);
        change.Parameters.Add(new("previous", ParameterAttributes.None, thread));
        change.Parameters.Add(new("next", ParameterAttributes.None, thread));
        eh.Methods.Add(change);
        var il = change.Body.GetILProcessor();
        var end = il.Create(OpCodes.Ret);
        var realSwitch = il.Create(OpCodes.Ldarg_0);
        il.Append(il.Create(OpCodes.Ldarg_1)); il.Append(il.Create(OpCodes.Brfalse, end));
        il.Append(il.Create(OpCodes.Ldarg_0)); il.Append(il.Create(OpCodes.Ldarg_1)); il.Append(il.Create(OpCodes.Beq, end));
        il.Append(il.Create(OpCodes.Ldarg_0)); il.Append(il.Create(OpCodes.Brtrue, realSwitch));
        // The first null->idle assignment adopts the already running boot
        // thread's state instead of resetting a live boot exception chain.
        Save(1); il.Append(il.Create(OpCodes.Br, end));
        il.Append(realSwitch); il.Append(il.Create(OpCodes.Pop));
        Save(0);
        for (int index = 0; index < state.Length; index++)
        { il.Append(il.Create(OpCodes.Ldarg_1)); il.Append(il.Create(OpCodes.Ldfld, state[index])); il.Append(il.Create(OpCodes.Stsfld, originals[index])); }
        il.Append(il.Create(OpCodes.Ldarg_1)); il.Append(il.Create(OpCodes.Ldfld, nativeHead)); il.Append(il.Create(OpCodes.Call, setHead));
        il.Append(end); change.Body.MaxStackSize = 2;

        var setterIl = setter.Body.GetILProcessor(); var first = setter.Body.Instructions[0];
        setterIl.InsertBefore(first, setterIl.Create(OpCodes.Ldarg_0)); setterIl.InsertBefore(first, setterIl.Create(OpCodes.Ldfld, current));
        setterIl.InsertBefore(first, setterIl.Create(OpCodes.Ldarg_1)); setterIl.InsertBefore(first, setterIl.Create(OpCodes.Call, change));
        setter.Body.MaxStackSize = Math.Max(2, setter.Body.MaxStackSize);
        var version = new MethodDefinition("XfeThreadEhContractVersionV1", MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.Int32);
        var export = new CustomAttribute(module.ImportReference(typeof(System.Runtime.InteropServices.UnmanagedCallersOnlyAttribute).GetConstructor(Type.EmptyTypes)!));
        export.Fields.Add(new("EntryPoint", new CustomAttributeArgument(module.TypeSystem.String, "XfeRuntimeThreadEhContractVersion")));
        version.CustomAttributes.Add(export); eh.Methods.Add(version);
        var versionIl = version.Body.GetILProcessor(); var disabled = versionIl.Create(OpCodes.Ldc_I4_0);
        var scheduler = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Scheduler.SchedulerManager");
        versionIl.Append(versionIl.Create(OpCodes.Call, scheduler.Methods.Single(method => method.Name == "get_IsRunning")));
        versionIl.Append(versionIl.Create(OpCodes.Brfalse, disabled));
        versionIl.Append(versionIl.Create(OpCodes.Call, scheduler.Methods.Single(method => method.Name == "get_CpuCount")));
        versionIl.Append(versionIl.Create(OpCodes.Ldc_I4_1)); versionIl.Append(versionIl.Create(OpCodes.Bne_Un, disabled));
        versionIl.Append(versionIl.Create(OpCodes.Ldc_I4_1)); versionIl.Append(versionIl.Create(OpCodes.Ret));
        versionIl.Append(disabled); versionIl.Append(versionIl.Create(OpCodes.Ret)); version.Body.MaxStackSize = 2;
        var identity = new MethodDefinition("XfeCurrentSchedulerThreadIdV1", MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.Int32);
        var identityExport = new CustomAttribute(module.ImportReference(typeof(System.Runtime.InteropServices.UnmanagedCallersOnlyAttribute).GetConstructor(Type.EmptyTypes)!));
        identityExport.Fields.Add(new("EntryPoint", new CustomAttributeArgument(module.TypeSystem.String, "XfeRuntimeCurrentSchedulerThreadId")));
        identity.CustomAttributes.Add(identityExport); eh.Methods.Add(identity);
        var identityIl = identity.Body.GetILProcessor(); var haveCpu = identityIl.Create(OpCodes.Callvirt, cpu.Methods.Single(method => method.Name == "get_CurrentThread"));
        var haveThread = identityIl.Create(OpCodes.Callvirt, thread.Methods.Single(method => method.Name == "get_Id"));
        identityIl.Append(identityIl.Create(OpCodes.Call, scheduler.Methods.Single(method => method.Name == "get_CurrentCpuState")));
        identityIl.Append(identityIl.Create(OpCodes.Dup)); identityIl.Append(identityIl.Create(OpCodes.Brtrue, haveCpu));
        identityIl.Append(identityIl.Create(OpCodes.Pop)); identityIl.Append(identityIl.Create(OpCodes.Ldc_I4_0)); identityIl.Append(identityIl.Create(OpCodes.Ret));
        identityIl.Append(haveCpu); identityIl.Append(identityIl.Create(OpCodes.Dup)); identityIl.Append(identityIl.Create(OpCodes.Brtrue, haveThread));
        identityIl.Append(identityIl.Create(OpCodes.Pop)); identityIl.Append(identityIl.Create(OpCodes.Ldc_I4_0)); identityIl.Append(identityIl.Create(OpCodes.Ret));
        identityIl.Append(haveThread); identityIl.Append(identityIl.Create(OpCodes.Ret)); identity.Body.MaxStackSize = 2;
        return false;

        void Save(int argument)
        {
            for (int index = 0; index < state.Length; index++)
            { il.Append(il.Create(argument == 0 ? OpCodes.Ldarg_0 : OpCodes.Ldarg_1)); il.Append(il.Create(OpCodes.Ldsfld, originals[index])); il.Append(il.Create(OpCodes.Stfld, state[index])); }
            il.Append(il.Create(argument == 0 ? OpCodes.Ldarg_0 : OpCodes.Ldarg_1)); il.Append(il.Create(OpCodes.Call, getHead)); il.Append(il.Create(OpCodes.Stfld, nativeHead));
        }
        MethodDefinition Import(string name, string symbol, TypeReference returns)
        {
            var method = new MethodDefinition(name, MethodAttributes.Private | MethodAttributes.Static | MethodAttributes.HideBySig | MethodAttributes.PInvokeImpl, returns)
            { ImplAttributes = MethodImplAttributes.PreserveSig };
            var library = module.ModuleReferences.SingleOrDefault(reference => reference.Name == "libSystem.Native");
            if (library is null) { library = new ModuleReference("libSystem.Native"); module.ModuleReferences.Add(library); }
            method.PInvokeInfo = new PInvokeInfo(PInvokeAttributes.CallConvCdecl | PInvokeAttributes.NoMangle, symbol, library);
            method.CustomAttributes.Add(new CustomAttribute(module.ImportReference(typeof(System.Runtime.InteropServices.SuppressGCTransitionAttribute).GetConstructor(Type.EmptyTypes)!)));
            eh.Methods.Add(method); return method;
        }
    }
}