railroader-setons-special-sauce/mcp-tools/DeepMethodProfiler.cs

using System;
using System.Collections;
using System.Collections.Generic;
using System.Diagnostics;
using System.Reflection;
using System.Runtime.CompilerServices;
using HarmonyLib;
using S3.Modules.Profiler;
using UnityEngine;

namespace S3.Mcp.Tools;

/// <summary>
/// Temporary, capture-scoped timings around game MonoBehaviour lifecycle methods.
/// This deliberately lives in the reloadable MCP tools pack so profiling coverage
/// can be refined without restarting the game.
/// </summary>
internal static class DeepMethodProfiler
{
    static readonly string[] LifecycleMethods =
    {
        "FixedUpdate",
        "Update",
        "LateUpdate",
        "OnPreCull",
        "OnPreRender",
        "OnPostRender",
        "OnWillRenderObject",
    };

    static readonly string[] ExcludedAssemblyPrefixes =
    {
        "System",
        "Microsoft",
        "Mono.",
        "mscorlib",
        "netstandard",
        "0Harmony",
        "Newtonsoft.",
        "UnityModManager",
        "S3",
    };

    static readonly Dictionary<MethodBase, string> ProbeIds = new();
    static Harmony? _harmony;

    public static bool Active => _harmony != null;
    public static int PatchedMethods => ProbeIds.Count;

    public static bool Start(out string details)
    {
        if (_harmony != null)
        {
            details = $"already active ({ProbeIds.Count} methods)";
            return true;
        }

        string harmonyId =
            "S3.mcp.deep." + typeof(DeepMethodProfiler).Assembly.GetName().Name;
        var harmony = new Harmony(harmonyId);
        var prefix = new HarmonyMethod(
            typeof(DeepMethodProfiler), nameof(Prefix));
        var postfix = new HarmonyMethod(
            typeof(DeepMethodProfiler), nameof(Postfix));
        int patched = 0;
        int lifecyclePatched = 0;
        int stateMachinePatched = 0;
        int failed = 0;

        foreach (Assembly assembly in AppDomain.CurrentDomain.GetAssemblies())
        {
            if (!IsCandidateAssembly(assembly)) continue;
            foreach (Type type in GetTypesSafely(assembly))
            {
                if (type == null || type.IsAbstract) continue;
                if (typeof(MonoBehaviour).IsAssignableFrom(type))
                {
                    for (int i = 0; i < LifecycleMethods.Length; i++)
                    {
                        MethodInfo? method = FindParameterlessMethod(
                            type, LifecycleMethods[i]);
                        if (!CanPatch(method)) continue;
                        if (TryPatch(
                                harmony, method!, prefix, postfix,
                                $"deep.{method!.Name}:{DisplayTypeName(type)}"))
                        {
                            patched++;
                            lifecyclePatched++;
                        }
                        else failed++;
                    }
                }

                bool coroutine = typeof(IEnumerator).IsAssignableFrom(type);
                bool asyncStateMachine =
                    typeof(IAsyncStateMachine).IsAssignableFrom(type);
                if (coroutine || asyncStateMachine)
                {
                    MethodInfo? moveNext = FindParameterlessMethod(type, "MoveNext");
                    if (!CanPatch(moveNext)) continue;
                    string kind = coroutine ? "Coroutine" : "Async";
                    if (TryPatch(
                            harmony, moveNext!, prefix, postfix,
                            $"deep.{kind}:{DisplayTypeName(type)}"))
                    {
                        patched++;
                        stateMachinePatched++;
                    }
                    else failed++;
                }
            }
        }

        if (patched == 0)
        {
            harmony.UnpatchAll(harmonyId);
            ProbeIds.Clear();
            details = $"no methods patched (failures={failed})";
            return false;
        }

        _harmony = harmony;
        details =
            $"patched={patched} lifecycle={lifecyclePatched} " +
            $"stateMachines={stateMachinePatched} failures={failed}";
        return true;
    }

    public static void Stop()
    {
        Harmony? harmony = _harmony;
        _harmony = null;
        if (harmony != null)
        {
            try { harmony.UnpatchAll(harmony.Id); }
            catch { }
        }
        ProbeIds.Clear();
    }

    static bool IsCandidateAssembly(Assembly assembly)
    {
        string name;
        try { name = assembly.GetName().Name ?? ""; }
        catch { return false; }
        if (name.Equals("S3", StringComparison.OrdinalIgnoreCase))
            return true;

        for (int i = 0; i < ExcludedAssemblyPrefixes.Length; i++)
        {
            if (name.StartsWith(
                    ExcludedAssemblyPrefixes[i],
                    StringComparison.OrdinalIgnoreCase))
                return false;
        }
        return true;
    }

    static Type[] GetTypesSafely(Assembly assembly)
    {
        try { return assembly.GetTypes(); }
        catch (ReflectionTypeLoadException ex)
        {
            var result = new List<Type>();
            Type?[] types = ex.Types;
            for (int i = 0; i < types.Length; i++)
                if (types[i] != null) result.Add(types[i]!);
            return result.ToArray();
        }
        catch { return Array.Empty<Type>(); }
    }

    static MethodInfo? FindParameterlessMethod(Type type, string name)
    {
        MethodInfo[] methods;
        try
        {
            methods = type.GetMethods(
                BindingFlags.Instance |
                BindingFlags.Public |
                BindingFlags.NonPublic |
                BindingFlags.DeclaredOnly);
        }
        catch { return null; }

        for (int i = 0; i < methods.Length; i++)
            if (methods[i].Name == name &&
                methods[i].GetParameters().Length == 0)
                return methods[i];
        return null;
    }

    static bool CanPatch(MethodInfo? method)
    {
        if (method == null || method.IsAbstract ||
            method.ContainsGenericParameters)
            return false;
        try { return method.GetMethodBody() != null; }
        catch { return false; }
    }

    static bool TryPatch(
        Harmony harmony,
        MethodInfo method,
        HarmonyMethod prefix,
        HarmonyMethod postfix,
        string id)
    {
        try
        {
            ProbeIds[method] = id;
            harmony.Patch(method, prefix, postfix);
            return true;
        }
        catch
        {
            ProbeIds.Remove(method);
            return false;
        }
    }

    static string DisplayTypeName(Type type)
    {
        string name = type.FullName ?? type.Name;
        Type? declaring = type.DeclaringType;
        if (declaring == null) return name;
        return (declaring.FullName ?? declaring.Name) + "." + type.Name;
    }

    static void Prefix(out long __state)
    {
        __state = Stopwatch.GetTimestamp();
    }

    static void Postfix(MethodBase __originalMethod, long __state)
    {
        if (__state == 0 || !SparseHitchSampler.Active) return;
        if (!ProbeIds.TryGetValue(__originalMethod, out string? id)) return;
        SparseHitchSampler.Record(id, Stopwatch.GetTimestamp() - __state);
    }
}