railroader-setons-special-sauce/src/Modules/Profiler/HitchProbePatches.cs

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Reflection;
using HarmonyLib;
using S3.Core;

namespace S3.Modules.Profiler;

/// <summary>
/// Resolves vanilla targets at runtime so game-version method drift degrades telemetry rather
/// than preventing the Profiler module from loading.
/// </summary>
public static class HitchProbePatches
{
    readonly struct Target
    {
        public readonly string Type;
        public readonly string Method;
        public readonly string Id;

        public Target(string type, string method, string id)
        {
            Type = type;
            Method = method;
            Id = id;
        }
    }

    static readonly Target[] Targets =
    {
        new("RollingStock.CarCuller", "Update", "cars.culler_update"),
        new("RollingStock.CarCuller", "ProcessPending", "cars.process_pending"),
        new("RollingStock.CarCuller", "OnCarCullingGroupStateChanged", "cars.cull_transition"),
        new("Model.Car", "ModelLoadRetain", "cars.model_retain"),
        new("Model.Car", "HandleModelsLoaded", "cars.model_loaded"),
        new("Model.Car", "UnloadModels", "cars.model_unload"),
        new("Model.Car", "SetCullerDistanceBand", "cars.distance_band"),
        new("Model.Car", "PositionWheelBoundsFront", "cars.position_wheels"),
        new("Helpers.SceneryAssetInstance", "SetLoaded", "scenery.set_loaded"),
        new("Helpers.SceneryAssetInstance", "DidLoadModel", "scenery.did_load_model"),
        new("Helpers.SceneryAssetInstance", "CullingSphereStateChanged", "scenery.cull_transition"),
        new("Helpers.Culling.CullingManager", "Update", "culling.update"),
        new("Helpers.Culling.CullingManager", "FixedUpdate", "culling.fixed_update"),
        new("Cameras.StrategyCameraController", "Update", "camera.strategy_update"),
        new("Cameras.StrategyCameraController", "UpdateCameraPosition", "camera.update_position"),
        new("Cameras.StrategyCameraController", "FindGround", "camera.find_ground"),
        new("TrainController", "FixedUpdate", "train.fixed_update"),
        new("TrainController", "CarDidPosition", "train.car_did_position"),
        new("WorldStreamer2.Streamer", "Update", "streamer.update"),
        new("WorldStreamer2.Streamer", "CheckPositionTiles", "streamer.check_tiles"),
        new("WorldStreamer2.Streamer", "LoadLevelAsyncManage", "streamer.load_pump"),
        new("WorldStreamer2.Streamer", "SceneLoading", "streamer.scene_loading"),
        new("WorldStreamer2.Streamer", "SceneUnloading", "streamer.scene_unloading"),
        new("WorldStreamer2.StreamerLoadingManager", "Update", "streamer.manager_update"),
        new("WorldStreamer2.TerrainCullingSystem", "Update", "terrain.culling_update"),
        new("WorldStreamer2.TerrainCullingSystem", "CheckVisibility", "terrain.check_visibility"),
        new("WorldStreamer2.PhysicCullingSystem", "Update", "physics_culling.update"),
        new("WorldStreamer2.PhysicCullingSystem", "CheckVisibility", "physics_culling.check_visibility"),
    };

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

    public static void Install(Harmony harmony)
    {
        ProbeIds.Clear();
        var prefix = new HarmonyMethod(typeof(HitchProbePatches), nameof(Prefix));
        var postfix = new HarmonyMethod(typeof(HitchProbePatches), nameof(Postfix));
        int patched = 0;
        int missed = 0;

        foreach (Target target in Targets)
        {
            Type? type = FindType(target.Type);
            if (type == null)
            {
                Log.Warn($"[profiler] Hitch probe type missing: {target.Type}");
                missed++;
                continue;
            }

            bool found = false;
            foreach (MethodInfo method in AccessTools.GetDeclaredMethods(type))
            {
                if (method.Name != target.Method || method.IsAbstract || method.ContainsGenericParameters)
                    continue;
                found = true;
                if (ProbeIds.ContainsKey(method)) continue;
                try
                {
                    ProbeIds[method] = target.Id;
                    harmony.Patch(method, prefix, postfix);
                    patched++;
                }
                catch (Exception e)
                {
                    ProbeIds.Remove(method);
                    Log.Warn($"[profiler] Hitch probe failed: {type.FullName}.{method.Name}: {e.Message}");
                    missed++;
                }
            }

            if (!found)
            {
                Log.Warn($"[profiler] Hitch probe method missing: {target.Type}.{target.Method}");
                missed++;
            }
        }

        Log.Info($"[profiler] Hitch probes installed: {patched}; unavailable: {missed}.");
    }

    public static void Prefix(MethodBase __originalMethod, out long __state)
    {
        __state = HitchSampler.Active && ProbeIds.ContainsKey(__originalMethod)
            ? Stopwatch.GetTimestamp()
            : 0;
    }

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

    static Type? FindType(string fullName)
    {
        foreach (Assembly assembly in AppDomain.CurrentDomain.GetAssemblies())
        {
            try
            {
                Type? type = assembly.GetType(fullName, false);
                if (type != null) return type;
            }
            catch { }
        }
        return null;
    }
}