diff --git a/src/BAUPlugin.cs b/src/BAUPlugin.cs index 95b784c..16476a6 100644 --- a/src/BAUPlugin.cs +++ b/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()]); AutoRegisterAttribute.Initialize(); diff --git a/src/Interfaces/IMonoExtension.cs b/src/Interfaces/IMonoExtension.cs index 374e1db..2103230 100644 --- a/src/Interfaces/IMonoExtension.cs +++ b/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(); /// - /// Static dictionary mapping MonoBehaviour instances to their active IMonoExtension instances. + /// Static dictionary mapping base MonoBehaviour types to their extension pairs. /// - private static readonly ConditionalWeakTable> MonoToMonoExtensionLookup = []; + private static readonly Dictionary> _extensionsByBaseType = []; + + /// + /// Represents a pairing between a base MonoBehavior and its extension. + /// + private struct ExtensionPair + { + /// + /// The base MonoBehavior. + /// + internal MonoBehaviour? Base; + + /// + /// The extension attached to the base. + /// + internal IMonoExtension? Extension; + } + + /// + /// Removes all entries from the lookup dictionary where the base MonoBehaviour or extension has been destroyed. + /// + 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); + } + } + } /// /// Determines whether a MonoBehaviour has an extension of the specified type attached to it. @@ -42,14 +92,16 @@ internal interface IMonoExtension /// true if an extension of type T exists on the MonoBehaviour; otherwise, false. internal static bool HasExtension(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 /// The extension instance if found, otherwise default(T). internal static T? GetExtension(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 /// The newly created extension instance, or null if creation failed. internal static T? AddExtension(MonoBehaviour monoBehaviour) where T : MonoBehaviour, IMonoExtension { - if (monoBehaviour == null) + if (monoBehaviour == null || monoBehaviour.IsDestroyedOrNull()) return null; - // Check if extension already exists - var existing = GetExtension(monoBehaviour); - if (existing != null) - return existing; + CleanupLookups(); + T? existingComponent = monoBehaviour.GetComponent(); + 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 + { + 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(); 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 } /// - /// Tries to removes an extension from its attached MonoBehaviour. + /// Tries to remove an extension from its attached MonoBehaviour. /// /// The extension instance to remove. 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); + } } } } + + /// + /// Runs a callback when a MonoBehavior extension becomes available. + /// + /// The type of extension to wait for. + /// The base MonoBehavior. + /// Function to retrieve the extension. + /// Callback to execute when extension is available. + internal static void RunWhenNotNull(MonoBehaviour mono, Func getExtension, Action callback) where T : class, IMonoExtension + { + mono.StartCoroutine(CoWaitForExtension(getExtension, callback)); + } + + /// + /// Coroutine that waits for an extension to become available. + /// + private static IEnumerator CoWaitForExtension(Func getExtension, Action callback) where T : class, IMonoExtension + { + T? extension; + while ((extension = getExtension()) == null) + { + yield return null; + } + callback(extension); + } } /// @@ -174,176 +315,23 @@ internal interface IMonoExtension : IMonoExtension where T : MonoBehaviour } /// -/// Provides functionality to patch MonoBehaviour methods and apply extension patches. +/// Extension methods for Unity's UnityEngine.Object types. /// -internal interface IMonoExtensionPatcher +public static class UnityObjectExtensions { /// - /// Gets the target patch information. + /// Determines whether a Unity object has been destroyed or is null. /// - TargetPatch Target { get; } - - /// - /// Adds an extension patch to the specified MonoBehaviour. - /// - /// The MonoBehaviour to patch. - void AddExtensionPatch(MonoBehaviour monoBehaviour); - - /// - /// Represents a target method to be patched. - /// - /// The type containing the method. - /// The name of the method to patch. - internal sealed record TargetPatch(Type Type, string MethodName); - - /// - /// Patches all IMonoExtensionPatcher implementations found in the mod assembly. - /// - internal static void PatchAll() + /// The Unity object to check. + /// + /// true if the object has been destroyed or is null; otherwise, false. + /// + /// + /// This method uses Unity's overloaded == operator which returns true for + /// destroyed objects even though the C# reference still exists. + /// + 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; } - - /// - /// Patches a specific IMonoExtensionPatcher implementation. - /// - /// The Harmony instance to use for patching. - /// The patcher type to apply. - 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)); - } - - /// - /// Postfix method that applies extension patches to MonoBehaviour instances. - /// - /// The MonoBehaviour instance being patched. - 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); - } - } - - /// - /// Cache for uninitialized patcher instances. - /// - private static readonly Dictionary _uninitializedMonoExtensionPatcherLookup = []; - - /// - /// Maps MonoBehaviour types to their corresponding patcher types. - /// - private static readonly Dictionary> _monoBehaviourToPatcherMap = []; - - /// - /// Gets or creates an uninitialized patcher instance for the specified type. - /// - /// The patcher type. - /// The MonoBehaviour instance. - /// An uninitialized patcher instance. - 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; - } -} - -/// -/// Provides generic functionality to patch MonoBehaviour methods and apply extension patches. -/// -/// The specific MonoBehaviour type to patch. -internal interface IMonoExtensionPatcher : IMonoExtensionPatcher where T : MonoBehaviour -{ - /// - /// Adds an extension patch to the specified MonoBehaviour. - /// - /// The MonoBehaviour to patch. - void IMonoExtensionPatcher.AddExtensionPatch(MonoBehaviour monoBehaviour) - { - AddExtensionPatch((T)monoBehaviour); - } - - /// - /// Adds an extension patch to the specified MonoBehaviour of type T. - /// - /// The MonoBehaviour instance to patch. - void AddExtensionPatch(T monoBehaviour); } \ No newline at end of file diff --git a/src/MonoScripts/Extended/ExtendedPlayerControl.cs b/src/MonoScripts/Extended/ExtendedPlayerControl.cs index 5ab6a83..40b7933 100644 --- a/src/MonoScripts/Extended/ExtendedPlayerControl.cs +++ b/src/MonoScripts/Extended/ExtendedPlayerControl.cs @@ -9,15 +9,8 @@ namespace BetterAmongUs.MonoScripts.Extended; /// Extends PlayerControl with additional functionality. /// [RegisterInIl2Cpp] -internal sealed class ExtendedPlayerControl : MonoBehaviour, IMonoExtension, IMonoExtensionPatcher +internal sealed class ExtendedPlayerControl : MonoBehaviour, IMonoExtension { - public IMonoExtensionPatcher.TargetPatch Target => new(typeof(PlayerControl), nameof(PlayerControl.Awake)); - - public void AddExtensionPatch(PlayerControl playerControl) - { - IMonoExtension.AddExtension(playerControl); - } - public void OnExtensionAwake(PlayerControl playerControl) { playerControl.gameObject.AddComponent().Init(playerControl); diff --git a/src/MonoScripts/Extended/ExtendedPlayerInfo.cs b/src/MonoScripts/Extended/ExtendedPlayerInfo.cs index f5bf624..4979eb0 100644 --- a/src/MonoScripts/Extended/ExtendedPlayerInfo.cs +++ b/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. /// [RegisterInIl2Cpp] -internal sealed class ExtendedPlayerInfo : MonoBehaviour, IMonoExtension, IMonoExtensionPatcher +internal sealed class ExtendedPlayerInfo : MonoBehaviour, IMonoExtension { internal ExtendedPlayerInfo() { @@ -36,23 +34,6 @@ internal sealed class ExtendedPlayerInfo : MonoBehaviour, IMonoExtension 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(playerControl.Data); - } - public NetworkedPlayerInfo? BaseMono { get; set; } public void OnExtensionAwake(NetworkedPlayerInfo networkedPlayerInfo) diff --git a/src/Patches/Client/MonoExtensionPatch.cs b/src/Patches/Client/MonoExtensionPatch.cs new file mode 100644 index 0000000..1f0ef5a --- /dev/null +++ b/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(__instance); + __instance.StartCoroutine(CoAddExtensionPatch(__instance)); + } + + private static IEnumerator CoAddExtensionPatch(PlayerControl playerControl) + { + while (playerControl.Data == null) + { + yield return null; + } + + IMonoExtension.AddExtension(playerControl.Data); + } +}