Revert back to old IMonoExtension system

src/BAUPlugin.cs

@@ -8,7 +8,6 @@ using BetterAmongUs.Data;
 using BetterAmongUs.Data.Config;
 using BetterAmongUs.Data.Json;
 using BetterAmongUs.Enums;
-using BetterAmongUs.Interfaces;
 using BetterAmongUs.Managers;
 using BetterAmongUs.Modules;
 using BetterAmongUs.Modules.OptionItems;
@@ -158,7 +157,6 @@ internal class BAUPlugin : BasePlugin
         AudioOverrideManager.Initialize();
         Translator.Initialize();
         Harmony.PatchAll();
-        IMonoExtensionPatcher.PatchAll();
         GameSettingsPatch.SetupSettings(true);
         BAUModdedSupportEvents.OnBAUOptionsLoadedEvent.InvokeAll([.. OptionItem.AllOptions.Cast<object>()]);
         AutoRegisterAttribute.Initialize();

src/Interfaces/IMonoExtension.cs

@@ -1,7 +1,6 @@
-using HarmonyLib;
-using System.Reflection;
-using System.Runtime.CompilerServices;
-using System.Runtime.Serialization;
+using BepInEx.Unity.IL2CPP.Utils;
+using Il2CppInterop.Runtime;
+using System.Collections;
 using UnityEngine;
 
 namespace BetterAmongUs.Interfaces;
@@ -30,9 +29,60 @@ internal interface IMonoExtension
     void OnDestroy();
 
     /// <summary>
-    /// Static dictionary mapping MonoBehaviour instances to their active IMonoExtension instances.
+    /// Static dictionary mapping base MonoBehaviour types to their extension pairs.
     /// </summary>
-    private static readonly ConditionalWeakTable<MonoBehaviour, List<IMonoExtension>> MonoToMonoExtensionLookup = [];
+    private static readonly Dictionary<Type, List<ExtensionPair>> _extensionsByBaseType = [];
+
+    /// <summary>
+    /// Represents a pairing between a base MonoBehavior and its extension.
+    /// </summary>
+    private struct ExtensionPair
+    {
+        /// <summary>
+        /// The base MonoBehavior.
+        /// </summary>
+        internal MonoBehaviour? Base;
+
+        /// <summary>
+        /// The extension attached to the base.
+        /// </summary>
+        internal IMonoExtension? Extension;
+    }
+
+    /// <summary>
+    /// Removes all entries from the lookup dictionary where the base MonoBehaviour or extension has been destroyed.
+    /// </summary>
+    private static void CleanupLookups()
+    {
+        foreach (var kvp in _extensionsByBaseType.ToArray())
+        {
+            var extensions = kvp.Value;
+
+            for (int i = extensions.Count - 1; i >= 0; i--)
+            {
+                var pair = extensions[i];
+
+                bool baseIsDead = pair.Base == null || pair.Base.IsDestroyedOrNull();
+                bool extensionIsDead = pair.Extension == null ||
+                                       (pair.Extension as MonoBehaviour)?.IsDestroyedOrNull() != false;
+
+                if (baseIsDead || extensionIsDead)
+                {
+                    if (pair.Extension != null && !extensionIsDead)
+                    {
+                        pair.Extension.OnDestroy();
+                    }
+
+                    extensions.RemoveAt(i);
+                }
+            }
+
+            if (extensions.Count == 0)
+            {
+                _extensionsByBaseType.Remove(kvp.Key);
+            }
+        }
+    }
 
     /// <summary>
     /// Determines whether a MonoBehaviour has an extension of the specified type attached to it.
@@ -42,14 +92,16 @@ internal interface IMonoExtension
     /// <returns>true if an extension of type T exists on the MonoBehaviour; otherwise, false.</returns>
     internal static bool HasExtension<T>(MonoBehaviour monoBehaviour) where T : IMonoExtension
     {
-        if (monoBehaviour == null)
+        if (monoBehaviour == null || monoBehaviour.IsDestroyedOrNull())
             return false;
 
-        if (MonoToMonoExtensionLookup.TryGetValue(monoBehaviour, out var extensions))
+        CleanupLookups();
+
+        if (_extensionsByBaseType.TryGetValue(monoBehaviour.GetType(), out var extensions))
         {
-            foreach (var extension in extensions)
+            foreach (var pair in extensions)
             {
-                if (extension is T)
+                if (pair.Base == monoBehaviour && pair.Extension is T)
                     return true;
             }
         }
@@ -65,14 +117,16 @@ internal interface IMonoExtension
     /// <returns>The extension instance if found, otherwise default(T).</returns>
     internal static T? GetExtension<T>(MonoBehaviour monoBehaviour) where T : IMonoExtension
     {
-        if (monoBehaviour == null)
+        if (monoBehaviour == null || monoBehaviour.IsDestroyedOrNull())
             return default;
 
-        if (MonoToMonoExtensionLookup.TryGetValue(monoBehaviour, out var extensions))
+        CleanupLookups();
+
+        if (_extensionsByBaseType.TryGetValue(monoBehaviour.GetType(), out var extensions))
         {
-            foreach (var extension in extensions)
+            foreach (var pair in extensions)
             {
-                if (extension is T typedExtension)
+                if (pair.Base == monoBehaviour && pair.Extension is T typedExtension)
                     return typedExtension;
             }
         }
@@ -88,25 +142,71 @@ internal interface IMonoExtension
     /// <returns>The newly created extension instance, or null if creation failed.</returns>
     internal static T? AddExtension<T>(MonoBehaviour monoBehaviour) where T : MonoBehaviour, IMonoExtension
     {
-        if (monoBehaviour == null)
+        if (monoBehaviour == null || monoBehaviour.IsDestroyedOrNull())
             return null;
 
-        // Check if extension already exists
-        var existing = GetExtension<T>(monoBehaviour);
-        if (existing != null)
-            return existing;
+        CleanupLookups();
+        T? existingComponent = monoBehaviour.GetComponent<T>();
+        if (existingComponent != null)
+        {
+            var baseType = monoBehaviour.GetType();
+            if (_extensionsByBaseType.TryGetValue(baseType, out var extensions))
+            {
+                bool isRegistered = extensions.Any(pair => pair.Base == monoBehaviour && pair.Extension == existingComponent);
+
+                if (!isRegistered)
+                {
+                    extensions.Add(new ExtensionPair
+                    {
+                        Base = monoBehaviour,
+                        Extension = existingComponent
+                    });
+                    existingComponent.BaseMono = monoBehaviour;
+                    existingComponent.OnExtensionAwake(monoBehaviour);
+                }
+            }
+            else
+            {
+                var newExtensions = new List<ExtensionPair>
+                {
+                    new() {
+                        Base = monoBehaviour,
+                        Extension = existingComponent
+                    }
+                };
+                _extensionsByBaseType[baseType] = newExtensions;
+                existingComponent.BaseMono = monoBehaviour;
+                existingComponent.OnExtensionAwake(monoBehaviour);
+            }
+
+            return existingComponent;
+        }
+
+        var existingExtensions = monoBehaviour.GetComponentsInChildren(Il2CppType.From(typeof(T)), true);
+        if (existingExtensions.Length > 0)
+        {
+            var found = existingExtensions.FirstOrDefault() as T;
+            if (found != null)
+                return found;
+        }
 
         T? monoExtension = monoBehaviour.gameObject.AddComponent<T>();
         if (monoExtension != null)
         {
-            // Get or create the extension list for this MonoBehaviour
-            if (!MonoToMonoExtensionLookup.TryGetValue(monoBehaviour, out var extensions))
+            var baseType = monoBehaviour.GetType();
+
+            if (!_extensionsByBaseType.TryGetValue(baseType, out var extensions))
             {
                 extensions = [];
-                MonoToMonoExtensionLookup.Add(monoBehaviour, extensions);
+                _extensionsByBaseType[baseType] = extensions;
             }
 
-            extensions.Add(monoExtension);
+            extensions.Add(new ExtensionPair
+            {
+                Base = monoBehaviour,
+                Extension = monoExtension
+            });
+
             monoExtension.BaseMono = monoBehaviour;
             monoExtension.OnExtensionAwake(monoBehaviour);
             return monoExtension;
@@ -116,7 +216,7 @@ internal interface IMonoExtension
     }
 
     /// <summary>
-    /// Tries to removes an extension from its attached MonoBehaviour.
+    /// Tries to remove an extension from its attached MonoBehaviour.
     /// </summary>
     /// <param name="monoExtension">The extension instance to remove.</param>
     internal static void TryRemoveExtension(IMonoExtension monoExtension)
@@ -124,14 +224,55 @@ internal interface IMonoExtension
         if (monoExtension == null)
             return;
 
+        CleanupLookups();
+
         if (monoExtension.BaseMono != null)
         {
-            if (MonoToMonoExtensionLookup.TryGetValue(monoExtension.BaseMono, out var extensions))
+            var baseType = monoExtension.BaseMono.GetType();
+
+            if (_extensionsByBaseType.TryGetValue(baseType, out var extensions))
             {
-                extensions.Remove(monoExtension);
+                for (int i = extensions.Count - 1; i >= 0; i--)
+                {
+                    if (extensions[i].Extension == monoExtension)
+                    {
+                        extensions.RemoveAt(i);
+                        break;
+                    }
+                }
+
+                if (extensions.Count == 0)
+                {
+                    _extensionsByBaseType.Remove(baseType);
+                }
             }
         }
     }
+
+    /// <summary>
+    /// Runs a callback when a MonoBehavior extension becomes available.
+    /// </summary>
+    /// <typeparam name="T">The type of extension to wait for.</typeparam>
+    /// <param name="mono">The base MonoBehavior.</param>
+    /// <param name="getExtension">Function to retrieve the extension.</param>
+    /// <param name="callback">Callback to execute when extension is available.</param>
+    internal static void RunWhenNotNull<T>(MonoBehaviour mono, Func<T?> getExtension, Action<T> callback) where T : class, IMonoExtension
+    {
+        mono.StartCoroutine(CoWaitForExtension(getExtension, callback));
+    }
+
+    /// <summary>
+    /// Coroutine that waits for an extension to become available.
+    /// </summary>
+    private static IEnumerator CoWaitForExtension<T>(Func<T?> getExtension, Action<T> callback) where T : class, IMonoExtension
+    {
+        T? extension;
+        while ((extension = getExtension()) == null)
+        {
+            yield return null;
+        }
+        callback(extension);
+    }
 }
 
 /// <summary>
@@ -174,176 +315,23 @@ internal interface IMonoExtension<T> : IMonoExtension where T : MonoBehaviour
 }
 
 /// <summary>
-/// Provides functionality to patch MonoBehaviour methods and apply extension patches.
+/// Extension methods for Unity's UnityEngine.Object types.
 /// </summary>
-internal interface IMonoExtensionPatcher
+public static class UnityObjectExtensions
 {
     /// <summary>
-    /// Gets the target patch information.
+    /// Determines whether a Unity object has been destroyed or is null.
     /// </summary>
-    TargetPatch Target { get; }
-
-    /// <summary>
-    /// Adds an extension patch to the specified MonoBehaviour.
-    /// </summary>
-    /// <param name="monoBehaviour">The MonoBehaviour to patch.</param>
-    void AddExtensionPatch(MonoBehaviour monoBehaviour);
-
-    /// <summary>
-    /// Represents a target method to be patched.
-    /// </summary>
-    /// <param name="Type">The type containing the method.</param>
-    /// <param name="MethodName">The name of the method to patch.</param>
-    internal sealed record TargetPatch(Type Type, string MethodName);
-
-    /// <summary>
-    /// Patches all IMonoExtensionPatcher implementations found in the mod assembly.
-    /// </summary>
-    internal static void PatchAll()
+    /// <param name="obj">The Unity object to check.</param>
+    /// <returns>
+    /// <c>true</c> if the object has been destroyed or is null; otherwise, <c>false</c>.
+    /// </returns>
+    /// <remarks>
+    /// This method uses Unity's overloaded == operator which returns true for
+    /// destroyed objects even though the C# reference still exists.
+    /// </remarks>
+    public static bool IsDestroyedOrNull(this UnityEngine.Object obj)
     {
-        var assembly = ModInfo.Assembly;
-        var patcherTypes = assembly.GetTypes()
-            .Where(t => typeof(IMonoExtensionPatcher).IsAssignableFrom(t) && !t.IsInterface && !t.IsAbstract)
-            .ToList();
-
-        // Group patchers by their target method to avoid duplicate patches
-        var patchersByTarget = patcherTypes
-            .Select(t => (Type: t, Instance: (IMonoExtensionPatcher)FormatterServices.GetUninitializedObject(t)))
-            .GroupBy(x => x.Instance.Target)
-            .ToList();
-
-        foreach (var group in patchersByTarget)
-        {
-            var target = group.Key;
-            var patcherTypesForTarget = group.Select(x => x.Type).ToList();
-
-            var originalMethod = target.Type.GetMethod(target.MethodName,
-                BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic);
-
-            if (originalMethod == null)
-                continue;
-
-            _monoBehaviourToPatcherMap[target.Type] = patcherTypesForTarget;
-
-            var postfixMethod = typeof(IMonoExtensionPatcher).GetMethod(nameof(Postfix),
-                BindingFlags.Static | BindingFlags.NonPublic);
-
-            BAUPlugin.Harmony.Patch(originalMethod, postfix: new HarmonyMethod(postfixMethod));
-        }
+        return obj == null;
     }
-
-    /// <summary>
-    /// Patches a specific IMonoExtensionPatcher implementation.
-    /// </summary>
-    /// <param name="harmony">The Harmony instance to use for patching.</param>
-    /// <param name="monoExtensionPatcherType">The patcher type to apply.</param>
-    private static void Patch(Harmony harmony, Type monoExtensionPatcherType)
-    {
-        var instance = (IMonoExtensionPatcher)FormatterServices.GetUninitializedObject(monoExtensionPatcherType);
-        var target = instance.Target;
-
-        var originalMethod = target.Type.GetMethod(target.MethodName,
-            BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic);
-
-        if (originalMethod == null)
-            return;
-
-        if (!_monoBehaviourToPatcherMap.TryGetValue(target.Type, out var patcherTypes))
-        {
-            patcherTypes = [];
-            _monoBehaviourToPatcherMap[target.Type] = patcherTypes;
-        }
-
-        if (!patcherTypes.Contains(monoExtensionPatcherType))
-        {
-            patcherTypes.Add(monoExtensionPatcherType);
-        }
-
-        var postfixMethod = typeof(IMonoExtensionPatcher).GetMethod(nameof(Postfix),
-            BindingFlags.Static | BindingFlags.NonPublic);
-
-        harmony.Patch(originalMethod, postfix: new HarmonyMethod(postfixMethod));
-    }
-
-    /// <summary>
-    /// Postfix method that applies extension patches to MonoBehaviour instances.
-    /// </summary>
-    /// <param name="__instance">The MonoBehaviour instance being patched.</param>
-    private static void Postfix(MonoBehaviour __instance)
-    {
-        var type = __instance.GetType();
-
-        // Try exact match first
-        if (!_monoBehaviourToPatcherMap.TryGetValue(type, out var patcherTypes))
-        {
-            // Check for assignable types (handles inheritance)
-            foreach (var kvp in _monoBehaviourToPatcherMap)
-            {
-                if (kvp.Key.IsAssignableFrom(type))
-                {
-                    patcherTypes = kvp.Value;
-                    _monoBehaviourToPatcherMap[type] = patcherTypes;
-                    break;
-                }
-            }
-
-            if (patcherTypes == null)
-                return;
-        }
-
-        foreach (var patcherType in patcherTypes)
-        {
-            var patcher = GetUninitializedMonoExtensionPatcher(patcherType, __instance);
-            patcher.AddExtensionPatch(__instance);
-        }
-    }
-
-    /// <summary>
-    /// Cache for uninitialized patcher instances.
-    /// </summary>
-    private static readonly Dictionary<Type, IMonoExtensionPatcher> _uninitializedMonoExtensionPatcherLookup = [];
-
-    /// <summary>
-    /// Maps MonoBehaviour types to their corresponding patcher types.
-    /// </summary>
-    private static readonly Dictionary<Type, List<Type>> _monoBehaviourToPatcherMap = [];
-
-    /// <summary>
-    /// Gets or creates an uninitialized patcher instance for the specified type.
-    /// </summary>
-    /// <param name="monoExtensionPatcherType">The patcher type.</param>
-    /// <param name="monoBehaviour">The MonoBehaviour instance.</param>
-    /// <returns>An uninitialized patcher instance.</returns>
-    private static IMonoExtensionPatcher GetUninitializedMonoExtensionPatcher(Type monoExtensionPatcherType, MonoBehaviour monoBehaviour)
-    {
-        if (!_uninitializedMonoExtensionPatcherLookup.TryGetValue(monoExtensionPatcherType, out var uninitialized))
-        {
-            uninitialized = (IMonoExtensionPatcher)FormatterServices.GetUninitializedObject(monoExtensionPatcherType);
-            _uninitializedMonoExtensionPatcherLookup[monoExtensionPatcherType] = uninitialized;
-        }
-
-        return uninitialized;
-    }
-}
-
-/// <summary>
-/// Provides generic functionality to patch MonoBehaviour methods and apply extension patches.
-/// </summary>
-/// <typeparam name="T">The specific MonoBehaviour type to patch.</typeparam>
-internal interface IMonoExtensionPatcher<T> : IMonoExtensionPatcher where T : MonoBehaviour
-{
-    /// <summary>
-    /// Adds an extension patch to the specified MonoBehaviour.
-    /// </summary>
-    /// <param name="monoBehaviour">The MonoBehaviour to patch.</param>
-    void IMonoExtensionPatcher.AddExtensionPatch(MonoBehaviour monoBehaviour)
-    {
-        AddExtensionPatch((T)monoBehaviour);
-    }
-
-    /// <summary>
-    /// Adds an extension patch to the specified MonoBehaviour of type T.
-    /// </summary>
-    /// <param name="monoBehaviour">The MonoBehaviour instance to patch.</param>
-    void AddExtensionPatch(T monoBehaviour);
 }

src/MonoScripts/Extended/ExtendedPlayerControl.cs

@@ -9,15 +9,8 @@ namespace BetterAmongUs.MonoScripts.Extended;
 /// Extends PlayerControl with additional functionality.
 /// </summary>
 [RegisterInIl2Cpp]
-internal sealed class ExtendedPlayerControl : MonoBehaviour, IMonoExtension<PlayerControl>, IMonoExtensionPatcher<PlayerControl>
+internal sealed class ExtendedPlayerControl : MonoBehaviour, IMonoExtension<PlayerControl>
 {
-    public IMonoExtensionPatcher.TargetPatch Target => new(typeof(PlayerControl), nameof(PlayerControl.Awake));
-
-    public void AddExtensionPatch(PlayerControl playerControl)
-    {
-        IMonoExtension.AddExtension<ExtendedPlayerControl>(playerControl);
-    }
-
     public void OnExtensionAwake(PlayerControl playerControl)
     {
         playerControl.gameObject.AddComponent<PlayerInfoDisplay>().Init(playerControl);

src/MonoScripts/Extended/ExtendedPlayerInfo.cs

@@ -1,5 +1,4 @@
 using AmongUs.GameOptions;
-using BepInEx.Unity.IL2CPP.Utils;
 using BetterAmongUs.Attributes;
 using BetterAmongUs.Generated;
 using BetterAmongUs.Interfaces;
@@ -12,7 +11,6 @@ using HarmonyLib;
 using Hazel;
 using Il2CppInterop.Runtime.Attributes;
 using InnerNet;
-using System.Collections;
 using UnityEngine;
 
 namespace BetterAmongUs.MonoScripts.Extended;
@@ -21,7 +19,7 @@ namespace BetterAmongUs.MonoScripts.Extended;
 /// Extended player information with additional data and anti-cheat features.
 /// </summary>
 [RegisterInIl2Cpp]
-internal sealed class ExtendedPlayerInfo : MonoBehaviour, IMonoExtension<NetworkedPlayerInfo>, IMonoExtensionPatcher<PlayerControl>
+internal sealed class ExtendedPlayerInfo : MonoBehaviour, IMonoExtension<NetworkedPlayerInfo>
 {
     internal ExtendedPlayerInfo()
     {
@@ -36,23 +34,6 @@ internal sealed class ExtendedPlayerInfo : MonoBehaviour, IMonoExtension<Network
         }
     }
 
-    public IMonoExtensionPatcher.TargetPatch Target => new(typeof(PlayerControl), nameof(PlayerControl.Awake));
-
-    public void AddExtensionPatch(PlayerControl playerControl)
-    {
-        playerControl.StartCoroutine(CoAddExtensionPatch(playerControl));
-    }
-
-    private static IEnumerator CoAddExtensionPatch(PlayerControl playerControl)
-    {
-        while (playerControl.Data == null)
-        {
-            yield return null;
-        }
-
-        IMonoExtension.AddExtension<ExtendedPlayerInfo>(playerControl.Data);
-    }
-
     public NetworkedPlayerInfo? BaseMono { get; set; }
 
     public void OnExtensionAwake(NetworkedPlayerInfo networkedPlayerInfo)

src/Patches/Client/MonoExtensionPatch.cs

@@ -0,0 +1,29 @@
+using BepInEx.Unity.IL2CPP.Utils;
+using BetterAmongUs.Interfaces;
+using BetterAmongUs.MonoScripts.Extended;
+using HarmonyLib;
+using System.Collections;
+
+namespace BetterAmongUs.Patches.Client;
+
+[HarmonyPatch]
+internal static class MonoExtensionPatch
+{
+    [HarmonyPatch(typeof(PlayerControl), nameof(PlayerControl.Awake))]
+    [HarmonyPostfix]
+    private static void PlayerControl_Awake_Postfix(PlayerControl __instance)
+    {
+        IMonoExtension.AddExtension<ExtendedPlayerControl>(__instance);
+        __instance.StartCoroutine(CoAddExtensionPatch(__instance));
+    }
+
+    private static IEnumerator CoAddExtensionPatch(PlayerControl playerControl)
+    {
+        while (playerControl.Data == null)
+        {
+            yield return null;
+        }
+
+        IMonoExtension.AddExtension<ExtendedPlayerInfo>(playerControl.Data);
+    }
+}