XFESystem

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

/// <summary>Leaf reads of the pinned collector's completed-GC snapshot.</summary>
internal static class GcMemoryRuntimePatch
{
    private const string Prefix = "XfeLastGcMetricV1_";
    private static readonly string[] Fields = ["HeapSizeBytes", "FragmentedBytes", "TotalCommittedBytes", "CollectionIndex"];
    private static readonly string[] Exports = ["XfeRuntimeLastGcHeapSpanBytes", "XfeRuntimeLastGcFragmentedBytes", "XfeRuntimeLastGcCommittedBytes", "XfeRuntimeLastGcCollectionIndex"];

    public static bool Apply(AssemblyDefinition assembly)
    {
        var module = assembly.MainModule;
        var gc = module.Types.Single(type => type.FullName == "Cosmos.Kernel.Core.Memory.GarbageCollector.GarbageCollector");
        var snapshot = gc.Fields.Single(field => field.Name == "s_lastGCMemoryInfo");
        var info = snapshot.FieldType.Resolve();
        if (!snapshot.IsStatic || !info.IsValueType || Fields.Where((name, index) =>
                info.Fields.Single(field => field.Name == name).FieldType.MetadataType != (index == 3 ? MetadataType.Int32 : MetadataType.UInt64)).Any())
            throw new InvalidOperationException("The pinned completed-GC snapshot layout changed.");

        var existing = gc.Methods.Where(method => method.Name.StartsWith(Prefix, StringComparison.Ordinal)).ToArray();
        if (existing.Length != 0)
        {
            if (existing.Length != Fields.Length) throw new InvalidOperationException("Incomplete GC snapshot exports.");
            Verify();
            return true;
        }
        for (int index = 0; index < Fields.Length; index++)
        {
            var field = info.Fields.Single(candidate => candidate.Name == Fields[index]);
            var method = new MethodDefinition(Prefix + Fields[index], MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.UInt64);
            var export = new CustomAttribute(module.ImportReference(typeof(System.Runtime.InteropServices.UnmanagedCallersOnlyAttribute).GetConstructor(Type.EmptyTypes)!));
            export.Fields.Add(new("EntryPoint", new CustomAttributeArgument(module.TypeSystem.String, Exports[index])));
            method.CustomAttributes.Add(export);
            gc.Methods.Add(method);
            var il = method.Body.GetILProcessor();
            il.Append(il.Create(OpCodes.Ldsflda, snapshot));
            il.Append(il.Create(OpCodes.Ldfld, field));
            if (index == 3) il.Append(il.Create(OpCodes.Conv_U8));
            il.Append(il.Create(OpCodes.Ret));
            method.Body.MaxStackSize = 1;
        }
        var handleManager = gc.Fields.Single(field => field.Name == "s_gCHandleManager");
        var dependentStore = handleManager.FieldType.Resolve().Fields.Single(field => field.Name == "DependentHandleStore");
        var count = dependentStore.FieldType.Resolve().Methods.Single(method => method.Name == "get_Count");
        var countExport = new MethodDefinition("XfeDependentHandleCountV1", MethodAttributes.Assembly | MethodAttributes.Static | MethodAttributes.HideBySig, module.TypeSystem.UInt64);
        var countAttribute = new CustomAttribute(module.ImportReference(typeof(System.Runtime.InteropServices.UnmanagedCallersOnlyAttribute).GetConstructor(Type.EmptyTypes)!));
        countAttribute.Fields.Add(new("EntryPoint", new CustomAttributeArgument(module.TypeSystem.String, "XfeRuntimeDependentHandleCount")));
        countExport.CustomAttributes.Add(countAttribute); gc.Methods.Add(countExport);
        var countIl = countExport.Body.GetILProcessor();
        countIl.Append(countIl.Create(OpCodes.Ldsflda, handleManager));
        countIl.Append(countIl.Create(OpCodes.Ldflda, dependentStore));
        countIl.Append(countIl.Create(OpCodes.Call, count));
        countIl.Append(countIl.Create(OpCodes.Ret)); countExport.Body.MaxStackSize = 1;
        Verify();
        return false;

        void Verify()
        {
            var count = gc.Methods.SingleOrDefault(method => method.Name == "XfeDependentHandleCountV1");
            if (count is null || count.Parameters.Count != 0 || count.ReturnType.MetadataType != MetadataType.UInt64 || count.Body.Instructions.Count != 4 ||
                count.Body.Instructions[0].OpCode != OpCodes.Ldsflda || count.Body.Instructions[0].Operand is not FieldReference manager || manager.Name != "s_gCHandleManager" ||
                count.Body.Instructions[1].OpCode != OpCodes.Ldflda || count.Body.Instructions[1].Operand is not FieldReference store || store.Name != "DependentHandleStore" ||
                count.Body.Instructions[2].OpCode != OpCodes.Call || count.Body.Instructions[2].Operand is not MethodReference getter || getter.Name != "get_Count" ||
                getter.DeclaringType.FullName != "Cosmos.Kernel.Core.Memory.GarbageCollector.GCHandleSegmentStore" || count.Body.Instructions[3].OpCode != OpCodes.Ret)
                throw new InvalidOperationException("The dependent-handle count export is not read-only.");
            for (int index = 0; index < Fields.Length; index++)
            {
                var method = gc.Methods.Single(candidate => candidate.Name == Prefix + Fields[index]);
                var il = method.Body.Instructions;
                if (method.Parameters.Count != 0 || method.ReturnType.MetadataType != MetadataType.UInt64 || method.Body.ExceptionHandlers.Count != 0 ||
                    il.Count != (index == 3 ? 4 : 3) || il[0].OpCode != OpCodes.Ldsflda || il[0].Operand is not FieldReference source || source.FullName != snapshot.FullName ||
                    il[1].OpCode != OpCodes.Ldfld || il[1].Operand is not FieldReference value || value.FullName != info.Fields.Single(field => field.Name == Fields[index]).FullName ||
                    index == 3 && il[2].OpCode != OpCodes.Conv_U8 || il[^1].OpCode != OpCodes.Ret ||
                    !method.CustomAttributes.Any(attribute => attribute.AttributeType.Name == "UnmanagedCallersOnlyAttribute" && attribute.Fields.Any(field =>
                        field.Name == "EntryPoint" && Equals(field.Argument.Value, Exports[index]))))
                    throw new InvalidOperationException("A GC snapshot export is not a read-only leaf.");
            }
        }
    }
}