Refactor: Optimize DiscoBallManager to create and cache at start of round

CHANGELOG.md

@@ -4,6 +4,7 @@
 
 - Fixed more missing flickering behaviours for some animators controllers.
 - Fixed some powered lights not fully turning off or flickering when there are multiple Light components per container.
+- Improved performance by pre-loading certain assets at the start of round instead of at a timing-critical frame update.
 
 ## MuzikaGromche 13.37.1337 - Photosensitivity Warning Edition
 

MuzikaGromche/DiscoBallManager.cs

@@ -1,4 +1,5 @@
 using DunGen;
+using HarmonyLib;
 using System;
 using System.Collections.Generic;
 using System.IO;
@@ -11,53 +12,18 @@ namespace MuzikaGromche
     public static class DiscoBallManager
     {
         // A struct holding a disco ball container object and the name of a tile for which it was designed.
-        private readonly record struct Data(string TileName, GameObject DiscoBallContainer)
+        private readonly record struct TilePatch(string TileName, GameObject DiscoBallContainer)
         {
             // We are specifically looking for cloned tiles, not the original prototypes.
             public readonly string TileCloneName = $"{TileName}(Clone)";
         }
 
-        private static readonly List<Data> Containers = [];
-        private static readonly List<GameObject> InstantiatedContainers = [];
+        private static TilePatch[] Patches = [];
 
-        public static void Initialize()
-        {
-            const string BundleFileName = "muzikagromche_discoball";
-            string bundlePath = Path.Combine(Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location), BundleFileName);
-            var assetBundle = AssetBundle.LoadFromFile(bundlePath)
-                ?? throw new NullReferenceException("Failed to load bundle");
+        private static readonly List<GameObject> CachedDiscoBalls = [];
+        private static readonly List<Animator> CachedDiscoBallAnimators = [];
 
-            foreach ((string prefabPath, string tileName) in new[] {
-                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerManor.prefab", "ManorStartRoomSmall"),
-                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerManorOLD.prefab", "ManorStartRoom"),
-                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerFactory.prefab", "StartRoom"),
-                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerMineShaft.prefab", "MineshaftStartTile"),
-                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerLargeForkTileB.prefab", "LargeForkTileB"),
-                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerBirthdayRoomTile.prefab", "BirthdayRoomTile"),
-            })
-            {
-                var container = assetBundle.LoadAsset<GameObject>(prefabPath);
-                Containers.Add(new(tileName, container));
-            }
-        }
-
-        public static void Enable()
-        {
-            // Just in case
-            Disable();
-
-            var query = from tile in Resources.FindObjectsOfTypeAll<Tile>()
-                        join container in Containers
-                        on tile.gameObject.name equals container.TileCloneName
-                        select (tile, container);
-
-            foreach (var (tile, container) in query)
-            {
-                Enable(tile, container);
-            }
-        }
-
-        private static readonly string[] animatorNames = [
+        private static readonly string[] AnimatorContainersNames = [
             "DiscoBallProp/AnimContainer",
             "DiscoBallProp1/AnimContainer",
             "DiscoBallProp2/AnimContainer",
@@ -66,29 +32,134 @@ namespace MuzikaGromche
             "DiscoBallProp5/AnimContainer",
         ];
 
-        private static void Enable(Tile tile, Data container)
+        public static void Load()
         {
-            Debug.Log($"{nameof(MuzikaGromche)} {nameof(DiscoBallManager)} Enabling at '{tile.gameObject.name}'");
-            var discoBall = UnityEngine.Object.Instantiate(container.DiscoBallContainer, tile.transform);
-            InstantiatedContainers.Add(discoBall);
+            const string BundleFileName = "muzikagromche_discoball";
+            string bundlePath = Path.Combine(Path.GetDirectoryName(Assembly.GetExecutingAssembly().Location), BundleFileName);
+            var assetBundle = AssetBundle.LoadFromFile(bundlePath)
+                ?? throw new NullReferenceException("Failed to load bundle");
 
-            foreach (var animatorName in animatorNames)
+            (string PrefabPath, string TileName)[] patchDescriptors =
+            [
+                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerManor.prefab", "ManorStartRoomSmall"),
+                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerManorOLD.prefab", "ManorStartRoom"),
+                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerFactory.prefab", "StartRoom"),
+                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerMineShaft.prefab", "MineshaftStartTile"),
+                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerLargeForkTileB.prefab", "LargeForkTileB"),
+                ("Assets/LethalCompany/Mods/MuzikaGromche/DiscoBallContainerBirthdayRoomTile.prefab", "BirthdayRoomTile"),
+            ];
+
+            Patches = [.. patchDescriptors.Select(d =>
+                new TilePatch(d.TileName, assetBundle.LoadAsset<GameObject>(d.PrefabPath))
+            )];
+        }
+
+        internal static void Patch(Tile tile)
         {
-                if (discoBall.transform.Find(animatorName)?.gameObject is GameObject animator)
+            var query = from patch in Patches
+                        where tile.gameObject.name == patch.TileCloneName
+                        select patch;
+
+            // Should be just one, but FirstOrDefault() isn't usable with structs
+            foreach (var patch in query)
             {
-                    animator.GetComponent<Animator>().SetBool("on", true);
+                Patch(tile, patch);
             }
         }
+
+        static void Patch(Tile tile, TilePatch patch)
+        {
+            var discoBall = UnityEngine.Object.Instantiate(patch.DiscoBallContainer, tile.transform);
+            if (discoBall == null)
+            {
+                return;
+            }
+
+            foreach (var animator in FindDiscoBallAnimators(discoBall))
+            {
+                CachedDiscoBallAnimators.Add(animator);
+            }
+            CachedDiscoBalls.Add(discoBall);
+            discoBall.SetActive(false);
+
+            Debug.Log($"{nameof(MuzikaGromche)} {nameof(DiscoBallManager)} Patched tile '{tile.gameObject.name}'");
+        }
+
+        static IEnumerable<Animator> FindDiscoBallAnimators(GameObject discoBall)
+        {
+            foreach (var animatorContainerName in AnimatorContainersNames)
+            {
+                var transform = discoBall.transform.Find(animatorContainerName);
+                if (transform == null)
+                {
+                    // Not all prefabs have all possible animators, and it's OK
+                    continue;
+                }
+
+                var animator = transform.gameObject?.GetComponent<Animator>();
+                if (animator == null)
+                {
+                    // This would be weird
+                    continue;
+                }
+
+                yield return animator;
+            }
+        }
+
+        public static void Toggle(bool on)
+        {
+            Debug.Log($"{nameof(MuzikaGromche)} {nameof(DiscoBallManager)} Toggle {(on ? "ON" : "OFF")} {CachedDiscoBallAnimators.Count} animators");
+
+            foreach (var discoBall in CachedDiscoBalls)
+            {
+                discoBall.SetActive(true);
+            }
+            foreach (var animator in CachedDiscoBallAnimators)
+            {
+                animator?.SetBool("on", on);
+            }
+        }
+
+        public static void Enable()
+        {
+            Toggle(true);
         }
 
         public static void Disable()
         {
-            foreach (var discoBall in InstantiatedContainers)
-            {
-                Debug.Log($"{nameof(MuzikaGromche)} {nameof(DiscoBallManager)}: Disabling {discoBall.name}");
-                UnityEngine.Object.Destroy(discoBall);
+            Toggle(false);
         }
-            InstantiatedContainers.Clear();
+
+        internal static void Clear()
+        {
+            Debug.Log($"{nameof(MuzikaGromche)} {nameof(DiscoBallManager)} Clearing {CachedDiscoBalls.Count} disco balls & {CachedDiscoBallAnimators.Count} animators");
+            CachedDiscoBallAnimators.Clear();
+            CachedDiscoBalls.Clear();
+        }
+    }
+
+    [HarmonyPatch(typeof(Tile))]
+    static class DiscoBallTilePatch
+    {
+        [HarmonyPatch(nameof(Tile.AddTriggerVolume))]
+        [HarmonyPostfix]
+        static void OnAddTriggerVolume(Tile __instance)
+        {
+            DiscoBallManager.Patch(__instance);
+        }
+    }
+
+    [HarmonyPatch(typeof(RoundManager))]
+    static class DiscoBallDespawnPatch
+    {
+        [HarmonyPatch(nameof(RoundManager.DespawnPropsAtEndOfRound))]
+        [HarmonyPatch(nameof(RoundManager.OnDestroy))]
+        [HarmonyPrefix]
+        static void OnDestroy(RoundManager __instance)
+        {
+            var _ = __instance;
+            DiscoBallManager.Clear();
         }
     }
 }

MuzikaGromche/Plugin.cs

@@ -544,13 +544,15 @@ namespace MuzikaGromche
                     track.LoadedLoop = DownloadHandlerAudioClip.GetContent(requests[i * 2 + 1]);
                 }
                 Config = new Config(base.Config);
-                DiscoBallManager.Initialize();
+                DiscoBallManager.Load();
                 PoweredLightsAnimators.Load();
                 var harmony = new Harmony(PluginInfo.PLUGIN_NAME);
                 harmony.PatchAll(typeof(JesterPatch));
                 harmony.PatchAll(typeof(EnemyAIPatch));
                 harmony.PatchAll(typeof(PoweredLightsAnimatorsPatch));
                 harmony.PatchAll(typeof(AllPoweredLightsPatch));
+                harmony.PatchAll(typeof(DiscoBallTilePatch));
+                harmony.PatchAll(typeof(DiscoBallDespawnPatch));
             }
             else
             {