XFESystem

using System.Reflection;
using System.Runtime.Loader;
using Mono.Cecil;
using Mono.Cecil.Cil;

internal static class PeriodicSchedulerLogPatchChecks
{
    public static void Run(AssemblyDefinition before, AssemblyDefinition after, Action<string, bool> check)
    {
        var original = Target(before);
        var patched = Target(after);
        var oldInstructions = original.Body.Instructions;
        var newInstructions = patched.Body.Instructions;
        check("PeriodicSchedulerExactlyLogAndTerminalTimesliceInstructionsChanged", oldInstructions.Count == newInstructions.Count &&
            original.Body.CodeSize + 8 == patched.Body.CodeSize && original.Body.MaxStackSize == patched.Body.MaxStackSize &&
            original.Body.InitLocals == patched.Body.InitLocals &&
            original.Body.Variables.Select(local => local.VariableType.FullName).SequenceEqual(patched.Body.Variables.Select(local => local.VariableType.FullName)) &&
            original.Body.ExceptionHandlers.Count == 0 && patched.Body.ExceptionHandlers.Count == 0 &&
            oldInstructions.Where((instruction, index) => index != PeriodicSchedulerLogPatch.ChangedInstructionIndex && index != PeriodicSchedulerLogPatch.TimesliceInstructionIndex)
                .All(instruction => InstructionShape(instruction, original) == InstructionShape(newInstructions[oldInstructions.IndexOf(instruction)], patched)) &&
            oldInstructions[PeriodicSchedulerLogPatch.ChangedInstructionIndex].GetSize() == newInstructions[PeriodicSchedulerLogPatch.ChangedInstructionIndex].GetSize());
        var originalReachable = Reachable(original);
        var patchedReachable = Reachable(patched);
        check("PeriodicSchedulerOriginalSerialPrintsReachablePositiveControl", originalReachable.Count(IsSerialWrite) == 5);
        check("PeriodicSchedulerAllSerialPrintsUnreachable", patchedReachable.All(instruction => !IsSerialWrite(instruction)));
        check("PeriodicSchedulerTickCounterIncrementRemainsReachable", newInstructions.Take(4).All(patchedReachable.Contains) &&
            newInstructions[3].OpCode == OpCodes.Stsfld && newInstructions[3].Operand is FieldReference counter && counter.Name == "s_onTickLogCount");
        check("PeriodicSchedulerRuntimePassBoostAndGlobalAccountingRemainReachable", new[]
            { "set_TotalRuntime", "set_Pass", "set_IsBoosted", "UpdateGlobalPass" }.All(name =>
                patchedReachable.Any(instruction => instruction.Operand is MethodReference call && call.Name == name)));
        check("PeriodicSchedulerOtherMethodsUnchanged", original.DeclaringType.Methods.Where(method => method.Name != "OnTick").All(method =>
            Il(method) == Il(patched.DeclaringType.Methods.Single(candidate => candidate.FullName == method.FullName))));

        string prior = Il(patched);
        check("PeriodicSchedulerRepeatedApplyVerifiesWithoutMutation", PeriodicSchedulerLogPatch.Apply(after) && prior == Il(patched));
        using (var roundtrip = Clone(after))
            check("PeriodicSchedulerRoundtripIsIdempotent", PeriodicSchedulerLogPatch.Apply(roundtrip));
        using (var logOnly = Clone(after))
        {
            var terminal = Target(logOnly).Body.Instructions[PeriodicSchedulerLogPatch.TimesliceInstructionIndex];
            terminal.OpCode = OpCodes.Ldloc_2; terminal.Operand = null;
            check("PeriodicSchedulerVerifiedLogOnlyMigrationCompletesTimeslice", !PeriodicSchedulerLogPatch.Apply(logOnly) &&
                PeriodicSchedulerLogPatch.Apply(logOnly) && Il(Target(logOnly)) == Il(patched));
        }

        CheckRejected("PeriodicSchedulerChangedModuloRejected", before, assembly =>
            Target(assembly).Body.Instructions.Single(instruction => instruction.OpCode == OpCodes.Ldc_I4_S && Equals(instruction.Operand, (sbyte)100)).Operand = (sbyte)101);
        CheckRejected("PeriodicSchedulerChangedAccountingRejected", before, assembly =>
            Target(assembly).Body.Instructions[1].OpCode = OpCodes.Ldc_I4_2);
        CheckRejected("PeriodicSchedulerChangedAccountingDivisorRejected", before, assembly =>
            Target(assembly).Body.Instructions.Single(instruction => instruction.OpCode == OpCodes.Ldc_I4 && Equals(instruction.Operand, 10_000_000)).Operand = 1_000_000);
        CheckRejected("PeriodicSchedulerChangedTailThresholdRejected", after, assembly =>
            Target(assembly).Body.Instructions[PeriodicSchedulerLogPatch.TimesliceInstructionIndex].Operand = 2_000_000L);
        CheckRejected("PeriodicSchedulerChangedTailComparisonRejected", after, assembly =>
            Target(assembly).Body.Instructions[104].OpCode = OpCodes.Clt);
        CheckRejected("PeriodicSchedulerChangedDebugMessageRejected", before, assembly =>
            Target(assembly).Body.Instructions.Single(instruction => instruction.OpCode == OpCodes.Ldstr && Equals(instruction.Operand, " runQ=")).Operand = " changed=");
        CheckRejected("PeriodicSchedulerUnknownVersionRejected", before, assembly => assembly.Name.Version = new Version(3, 0, 95, 0));
        CheckRejected("PeriodicSchedulerChangedPatchedDestinationRejected", after, assembly =>
            Target(assembly).Body.Instructions[PeriodicSchedulerLogPatch.ChangedInstructionIndex].Operand = Target(assembly).Body.Instructions[83]);
        CheckRejected("PeriodicSchedulerPartialOrUnknownPatchRejected", after, assembly =>
        {
            var entry = Target(assembly).Body.Instructions[PeriodicSchedulerLogPatch.ChangedInstructionIndex];
            entry.OpCode = OpCodes.Nop; entry.Operand = null;
        });

        // Run the real original/patched callback IL on CoreCLR. Only referenced
        // data types and method tokens are replaced with these pure host doubles;
        // no scheduler, interrupt or hardware operations run on the host.
        using var originalFixture = ExecutableFixture(original);
        using var patchedFixture = ExecutableFixture(patched);
        using var originalRunner = new Runner(originalFixture);
        using var patchedRunner = new Runner(patchedFixture);
        Case[] cases =
        [
            new("BelowForcedTimeslice", 99, 999_999),
            new("ForcedTimesliceElapsed", 99, 1_000_000),
            new("QuantumElapsed", 99, 10_000_000),
            new("LowerPassPreempts", 99, 1_000_000, NextPass: 0),
            new("HigherPassKeepsCurrent", 99, 1_000_000, NextPass: 1_000_000),
            new("BoostDecay", 99, 1_000_000, Boosted: true),
            new("NullCurrent", 99, 1_000_000, HasCurrent: false),
            new("NullCpuData", 99, 1_000_000, HasCpu: false),
            new("ExitedThread", 99, 1_000_000, HasThreadData: false, NextPass: 0),
            new("CounterOverflow", uint.MaxValue, 1_000_000)
        ];
        var observations = cases.Select(test => (Test: test, Original: originalRunner.Run(test), Patched: patchedRunner.Run(test))).ToArray();
        foreach (var observation in observations)
            check("PeriodicSchedulerActualIl_" + observation.Test.Name, observation.Patched.SerialWrites == 0 &&
                observation.Patched with { Preempt = observation.Original.Preempt } == observation.Original with { SerialWrites = 0 } &&
                observation.Patched.Counter == unchecked(observation.Test.Counter + 1));
        check("PeriodicSchedulerExecutablePositiveControlActuallyPrinted", observations[1].Original.SerialWrites == 5 &&
            observations[1].Original.TotalRuntime == 1_000_100 && observations[1].Original.Pass == 602 &&
            !observations[1].Original.Preempt && observations[1].Original.GlobalUpdates == 1);
        check("PeriodicSchedulerOneMsForcesRotationButSubMsDoesNot", !observations[0].Patched.Preempt && observations[1].Patched.Preempt &&
            !observations[1].Original.Preempt && observations[4].Patched.Preempt && !observations[4].Original.Preempt);
        check("PeriodicSchedulerPreemptionAndBoostExecutableControl", observations[2].Patched.Preempt && observations[3].Patched.Preempt &&
            !observations[5].Patched.Boosted && !observations[6].Patched.Preempt &&
            !observations[7].Patched.Preempt && observations[8].Patched.Preempt);
        var oldFair = originalRunner.RunTwoRunnableThreads(0);
        var newFair = patchedRunner.RunTwoRunnableThreads(0);
        var oldGap = originalRunner.RunTwoRunnableThreads(1_000_000_000_000);
        var newGap = patchedRunner.RunTwoRunnableThreads(1_000_000_000_000);
        check("PeriodicSchedulerTwoRunnableThreadsBothProgressWithRealOneMsAccounting", newFair.FirstTicks == 50 && newFair.SecondTicks == 50 &&
            newFair.FirstRuntime == 50_000_000 && newFair.SecondRuntime == 50_000_000 && newFair.MaximumSameThreadRun == 1);
        check("PeriodicSchedulerLargeWallGapPassDifferenceCannotMonopolizeRunnablePair", oldGap.FirstTicks == 100 && oldGap.SecondTicks == 0 &&
            newGap.FirstTicks == 50 && newGap.SecondTicks == 50 && newGap.MaximumSameThreadRun == 1 && newGap.FirstRuntime + newGap.SecondRuntime == 100_000_000);
        check("PeriodicSchedulerOriginalFairPairPositiveControlProgresses", oldFair.FirstTicks > 0 && oldFair.SecondTicks > 0);

        void CheckRejected(string name, AssemblyDefinition source, Action<AssemblyDefinition> alter)
        {
            using var changed = Clone(source);
            alter(changed);
            string unchanged = Il(Target(changed));
            bool rejected;
            try { PeriodicSchedulerLogPatch.Apply(changed); rejected = false; }
            catch (InvalidOperationException) { rejected = true; }
            check(name, rejected && unchanged == Il(Target(changed)));
        }
    }

    private static MethodDefinition Target(AssemblyDefinition assembly) => assembly.MainModule.Types.Single(type =>
        type.FullName == PeriodicSchedulerLogPatch.SchedulerType).Methods.Single(method => method.Name == "OnTick");
    private static string Il(MethodDefinition method) => !method.HasBody ? "" : string.Join("\n", method.Body.Instructions.Select(instruction => instruction.ToString()));
    private static string InstructionShape(Instruction instruction, MethodDefinition method) => instruction.OpCode.Code + ":" + (instruction.Operand switch
    {
        Instruction branch => "#" + method.Body.Instructions.IndexOf(branch),
        VariableDefinition local => "$" + local.Index,
        MemberReference member => member.FullName,
        null => "",
        var operand => operand.ToString()
    });
    private static bool IsSerialWrite(Instruction instruction) => instruction.Operand is MethodReference call &&
        call.DeclaringType.FullName == "Cosmos.Kernel.Core.IO.Serial";
    private static AssemblyDefinition Clone(AssemblyDefinition assembly)
    {
        // Cecil reads method bodies lazily; the returned assembly retains this
        // backing stream until its own disposal.
        var stream = new MemoryStream(); assembly.Write(stream); stream.Position = 0;
        return AssemblyDefinition.ReadAssembly(stream, new ReaderParameters { InMemory = true, ReadingMode = ReadingMode.Immediate });
    }
    private static HashSet<Instruction> Reachable(MethodDefinition method)
    {
        var found = new HashSet<Instruction>();
        var queue = new Queue<Instruction>(); queue.Enqueue(method.Body.Instructions[0]);
        while (queue.TryDequeue(out var instruction))
        {
            if (!found.Add(instruction)) continue;
            if (instruction.Operand is Instruction target) queue.Enqueue(target);
            if (instruction.Operand is Instruction[] targets) foreach (var branch in targets) queue.Enqueue(branch);
            if (instruction.OpCode.FlowControl is not (FlowControl.Branch or FlowControl.Return or FlowControl.Throw) && instruction.Next is { } next)
                queue.Enqueue(next);
        }
        return found;
    }

    private static AssemblyDefinition ExecutableFixture(MethodDefinition source)
    {
        var fixture = AssemblyDefinition.CreateAssembly(new AssemblyNameDefinition("SchedulerFixture_" + Guid.NewGuid().ToString("N"), new Version(1, 0)),
            "SchedulerFixture", ModuleKind.Dll);
        var module = fixture.MainModule;
        var type = new TypeDefinition("Fixture", "Callback", Mono.Cecil.TypeAttributes.Public | Mono.Cecil.TypeAttributes.Abstract | Mono.Cecil.TypeAttributes.Sealed,
            module.TypeSystem.Object); module.Types.Add(type);
        var counter = new FieldDefinition("Counter", Mono.Cecil.FieldAttributes.Public | Mono.Cecil.FieldAttributes.Static, module.TypeSystem.UInt32);
        type.Fields.Add(counter);
        var method = new MethodDefinition("OnTick", Mono.Cecil.MethodAttributes.Public | Mono.Cecil.MethodAttributes.Static, module.TypeSystem.Boolean);
        foreach (var parameter in new[] { typeof(PeriodicSchedulerFixture.Scheduler), typeof(PeriodicSchedulerFixture.CpuState), typeof(PeriodicSchedulerFixture.SchedulerThread), typeof(ulong) })
            method.Parameters.Add(new("arg" + method.Parameters.Count, Mono.Cecil.ParameterAttributes.None, module.ImportReference(parameter)));
        type.Methods.Add(method); method.Body.InitLocals = source.Body.InitLocals; method.Body.MaxStackSize = source.Body.MaxStackSize;
        var locals = source.Body.Variables.ToDictionary(local => local, local => new VariableDefinition(local.VariableType.FullName switch
        {
            "Cosmos.Kernel.Core.Scheduler.Stride.StrideCpuData" => module.ImportReference(typeof(PeriodicSchedulerFixture.CpuData)),
            "Cosmos.Kernel.Core.Scheduler.Stride.StrideThreadData" => module.ImportReference(typeof(PeriodicSchedulerFixture.ThreadData)),
            "System.UInt64" => module.TypeSystem.UInt64,
            _ => throw new InvalidOperationException("Unknown scheduler fixture local type.")
        }));
        foreach (var local in locals.Values) method.Body.Variables.Add(local);
        var instructions = source.Body.Instructions.ToDictionary(instruction => instruction, _ => Instruction.Create(OpCodes.Nop));
        foreach (var instruction in source.Body.Instructions)
        {
            var clone = instructions[instruction]; clone.OpCode = instruction.OpCode;
            clone.Operand = instruction.Operand switch
            {
                Instruction branch => instructions[branch],
                Instruction[] branches => branches.Select(branch => instructions[branch]).ToArray(),
                VariableDefinition local => locals[local],
                ParameterDefinition parameter => method.Parameters[parameter.Index + 1],
                FieldReference field when field.Name == "s_onTickLogCount" => counter,
                MethodReference call => FixtureCall(call),
                null => null,
                sbyte or byte or int or long or string => instruction.Operand,
                _ => throw new InvalidOperationException("Unknown scheduler fixture operand.")
            };
            method.Body.Instructions.Add(clone);
        }
        return fixture;

        MethodReference FixtureCall(MethodReference call)
        {
            Type hostType = call.DeclaringType.FullName switch
            {
                PeriodicSchedulerLogPatch.SchedulerType when call.Name is "CpuDataOf" or "ThreadDataOf" => typeof(PeriodicSchedulerFixture),
                PeriodicSchedulerLogPatch.SchedulerType => typeof(PeriodicSchedulerFixture.Scheduler),
                "Cosmos.Kernel.Core.Scheduler.SchedulerThread" => typeof(PeriodicSchedulerFixture.SchedulerThread),
                "Cosmos.Kernel.Core.Scheduler.Stride.StrideThreadData" => typeof(PeriodicSchedulerFixture.ThreadData),
                "Cosmos.Kernel.Core.Scheduler.Stride.StrideCpuData" => typeof(PeriodicSchedulerFixture.CpuData),
                "System.Collections.Generic.List`1<Cosmos.Kernel.Core.Scheduler.SchedulerThread>" => typeof(List<PeriodicSchedulerFixture.SchedulerThread>),
                "Cosmos.Kernel.Core.IO.Serial" => typeof(PeriodicSchedulerFixture),
                _ => throw new InvalidOperationException("Unknown scheduler fixture call type: " + call.FullName)
            };
            return module.ImportReference(hostType.GetMethod(call.Name, BindingFlags.Public | BindingFlags.Instance | BindingFlags.Static)
                ?? throw new InvalidOperationException("Missing scheduler fixture method: " + call.FullName));
        }
    }

    private sealed class Runner : IDisposable
    {
        private readonly AssemblyLoadContext context = new("SchedulerCallbackControl", isCollectible: true);
        private readonly MethodInfo callback;
        private readonly FieldInfo counter;
        public Runner(AssemblyDefinition fixture)
        {
            using var bytes = new MemoryStream(); fixture.Write(bytes); bytes.Position = 0;
            var type = context.LoadFromStream(bytes).GetType("Fixture.Callback", throwOnError: true)!;
            callback = type.GetMethod("OnTick")!; counter = type.GetField("Counter")!;
        }
        public Observation Run(Case test)
        {
            var scheduler = new PeriodicSchedulerFixture.Scheduler { Now = 10_000_000 };
            var data = new PeriodicSchedulerFixture.ThreadData { Pass = 500, Stride = 1024, IsBoosted = test.Boosted };
            var current = test.HasCurrent ? new PeriodicSchedulerFixture.SchedulerThread { Id = 7, TotalRuntime = 100, Data = test.HasThreadData ? data : null } : null;
            var cpuData = test.HasCpu ? new PeriodicSchedulerFixture.CpuData() : null;
            if (cpuData is not null && test.NextPass is { } pass)
                cpuData.RunQueue.Add(new() { Id = 8, Data = new() { Pass = pass } });
            var cpu = new PeriodicSchedulerFixture.CpuState { Data = cpuData };
            PeriodicSchedulerFixture.SerialWrites = 0; counter.SetValue(null, test.Counter);
            bool preempt = (bool)callback.Invoke(null, [scheduler, cpu, current, test.Elapsed])!;
            return new(preempt, (uint)counter.GetValue(null)!, current?.TotalRuntime ?? 0, data.Pass, data.IsBoosted, scheduler.GlobalUpdates,
                PeriodicSchedulerFixture.SerialWrites);
        }
        public FairnessObservation RunTwoRunnableThreads(long secondInitialPass)
        {
            var scheduler = new PeriodicSchedulerFixture.Scheduler { Now = 10_000_000 };
            var first = new PeriodicSchedulerFixture.SchedulerThread { Id = 7, Data = new() { Stride = 1024, Pass = 0 } };
            var second = new PeriodicSchedulerFixture.SchedulerThread { Id = 8, Data = new() { Stride = 1024, Pass = secondInitialPass } };
            var data = new PeriodicSchedulerFixture.CpuData(); var cpu = new PeriodicSchedulerFixture.CpuState { Data = data };
            var current = first; data.RunQueue.Add(second);
            int firstTicks = 0, secondTicks = 0, consecutive = 0, maximumConsecutive = 0; uint previous = 0;
            for (int tick = 0; tick < 100; tick++)
            {
                if (current == first) firstTicks++; else secondTicks++;
                consecutive = previous == current.Id ? consecutive + 1 : 1; previous = current.Id;
                maximumConsecutive = Math.Max(maximumConsecutive, consecutive);
                bool preempt = (bool)callback.Invoke(null, [scheduler, cpu, current, 1_000_000UL])!;
                scheduler.Now += 1_000_000;
                if (!preempt) continue;
                // With exactly two runnable threads, pinned PickNext removes
                // the sole queued entry before OnThreadYield inserts prev.
                var next = data.RunQueue[0]; data.RunQueue.RemoveAt(0); data.RunQueue.Add(current); current = next;
            }
            return new(firstTicks, secondTicks, first.TotalRuntime, second.TotalRuntime, maximumConsecutive);
        }
        public void Dispose() => context.Unload();
    }

    private sealed record Case(string Name, uint Counter, ulong Elapsed, long? NextPass = null, bool Boosted = false,
        bool HasCurrent = true, bool HasCpu = true, bool HasThreadData = true);
    private sealed record Observation(bool Preempt, uint Counter, ulong TotalRuntime, long Pass, bool Boosted, int GlobalUpdates, int SerialWrites);
    private sealed record FairnessObservation(int FirstTicks, int SecondTicks, ulong FirstRuntime, ulong SecondRuntime, int MaximumSameThreadRun);
}

public static class PeriodicSchedulerFixture
{
    public static int SerialWrites;
    public static void WriteString(string text) => SerialWrites++;
    public static void WriteNumber(uint number, bool hex) => SerialWrites++;
    public static CpuData? CpuDataOf(CpuState cpu) => cpu.Data;
    public static ThreadData? ThreadDataOf(SchedulerThread thread) => thread.Data;
    public sealed class Scheduler
    {
        public ulong Now;
        public int GlobalUpdates;
        public ulong GetTimestamp() => Now;
        public void UpdateGlobalPass(CpuData cpu) => GlobalUpdates++;
    }
    public sealed class CpuState { public CpuData? Data; }
    public sealed class CpuData { public List<SchedulerThread> RunQueue { get; } = []; }
    public sealed class SchedulerThread
    {
        public uint Id { get; set; }
        public ulong TotalRuntime { get; set; }
        public ThreadData? Data;
    }
    public sealed class ThreadData
    {
        public long Pass { get; set; }
        public ulong Stride { get; set; }
        public bool IsBoosted { get; set; }
        public ulong LastWakeup { get; set; }
    }
}