XFESystem

using System.Security.Cryptography;
using System.Text;
using Mono.Cecil;
using Mono.Cecil.Cil;

/// <summary>
/// Skips periodic debug output and explicitly uses a 1 ms forced timeslice.
/// The real elapsed-time accounting and 10 ms Stride weight divisor stay intact.
/// </summary>
internal static class PeriodicSchedulerLogPatch
{
    internal const string SchedulerType = "Cosmos.Kernel.Core.Scheduler.Stride.StrideScheduler";
    internal const int ChangedInstructionIndex = 63;
    internal const int ResumeInstructionIndex = 82;
    internal const int TimesliceInstructionIndex = 103;
    internal const long ForcedTimesliceNanoseconds = 1_000_000;
    internal const string OriginalMethodSha256 = "e1a70ae8c5e8c5bbd17e0255a98f0d40fd1529a3d32cf33cd0fa475dea3d396c";

    public static bool Apply(AssemblyDefinition assembly)
    {
        if (assembly.Name.Name != "Cosmos.Kernel.Core" || assembly.Name.Version != new Version(3, 0, 94, 0))
            throw new InvalidOperationException("Periodic scheduler logging patch requires Cosmos.Kernel.Core 3.0.94.");
        var type = assembly.MainModule.Types.Single(candidate => candidate.FullName == SchedulerType);
        var method = type.Methods.Single(candidate => candidate.Name == "OnTick");
        var counter = type.Fields.Single(candidate => candidate.Name == "s_onTickLogCount");
        if (method.IsStatic || method.ReturnType.MetadataType != MetadataType.Boolean ||
            method.Parameters.Count != 3 || method.Parameters[0].ParameterType.FullName != "Cosmos.Kernel.Core.Scheduler.PerCpuState" ||
            method.Parameters[1].ParameterType.FullName != "Cosmos.Kernel.Core.Scheduler.SchedulerThread" ||
            method.Parameters[2].ParameterType.MetadataType != MetadataType.UInt64 || !method.HasBody ||
            !counter.IsPrivate || !counter.IsStatic || counter.FieldType.MetadataType != MetadataType.UInt32 ||
            method.Body.Instructions.Sum(instruction => instruction.GetSize()) is not (269 or 277) || method.Body.Instructions.Count != 108)
            throw new InvalidOperationException("The pinned scheduler timer callback contract changed.");

        var instructions = method.Body.Instructions;
        var change = instructions[ChangedInstructionIndex];
        var resume = instructions[ResumeInstructionIndex];
        bool alreadyPatched = change.OpCode == OpCodes.Br && ReferenceEquals(change.Operand, resume);
        var timeslice = instructions[TimesliceInstructionIndex];
        bool timeslicePatched = timeslice.OpCode == OpCodes.Ldc_I8 && Equals(timeslice.Operand, ForcedTimesliceNanoseconds);
        if (!timeslicePatched && timeslice.OpCode != OpCodes.Ldloc_2)
            throw new InvalidOperationException("Unknown forced timeslice expression; refusing to change scheduler arithmetic.");
        if (instructions.Sum(instruction => instruction.GetSize()) != (timeslicePatched ? 277 : 269) || instructions[102].OpCode != OpCodes.Ldarg_3 ||
            instructions[104].OpCode != OpCodes.Clt_Un || instructions[105].OpCode != OpCodes.Ldc_I4_0 ||
            instructions[106].OpCode != OpCodes.Ceq || instructions[107].OpCode != OpCodes.Ret)
            throw new InvalidOperationException("The exact terminal elapsed-time comparison changed.");
        if (!alreadyPatched && (change.OpCode != OpCodes.Ldsfld || change.Operand is not FieldReference field || field.FullName != counter.FullName))
            throw new InvalidOperationException("Unknown periodic scheduler log entry; refusing to patch it.");
        // Bind the whole method, including every original branch destination and
        // local type. Normalization accepts exactly the one five-byte replacement.
        if (HashCanonical(method, counter) != OriginalMethodSha256)
            throw new InvalidOperationException("The pinned scheduler OnTick IL differs outside the periodic logging patch.");
        if (alreadyPatched && timeslicePatched) return true;

        // The log-entry ldsfld and br both occupy five bytes. Apart from
        // the explicit terminal threshold below, preserve all instructions,
        // branch targets, locals, exception regions and metadata entries.
        change.OpCode = OpCodes.Br;
        change.Operand = resume;
        // Only the final preemption comparison changes. The local quantum
        // remains 10 ms everywhere used for Pass/Stride runtime accounting.
        // A per-interrupt 1 ms elapsed value can never reach the old terminal
        // 10 ms test: it did not accumulate a timeslice and could starve a
        // ready UI behind a higher virtual pass after IRQ-off work.
        timeslice.OpCode = OpCodes.Ldc_I8;
        timeslice.Operand = ForcedTimesliceNanoseconds;
        return false;
    }

    private static string HashCanonical(MethodDefinition method, FieldReference counter)
    {
        var body = method.Body;
        // Restore the two precisely validated source instructions in the
        // fingerprint, including the eight-byte tail growth, so every other
        // instruction/local/branch remains bound to the fixed official hash.
        string canonical = 269 + "|" + body.MaxStackSize + "|" + body.InitLocals + "|" + body.ExceptionHandlers.Count + "|" +
            string.Join(";", body.Variables.Select(variable => variable.VariableType.FullName)) + "|" +
            string.Join("\n", body.Instructions.Select((instruction, index) => index == ChangedInstructionIndex
                ? Code.Ldsfld + ":" + counter.FullName
                : index == TimesliceInstructionIndex ? Code.Ldloc_2 + ":"
                : instruction.OpCode.Code + ":" + (instruction.Operand switch
                {
                    Instruction target => "#" + body.Instructions.IndexOf(target),
                    VariableDefinition local => "$" + local.Index,
                    MemberReference member => member.FullName,
                    null => "",
                    var operand => operand.ToString()
                })));
        return Convert.ToHexStringLower(SHA256.HashData(Encoding.UTF8.GetBytes(canonical)));
    }
}