Refactor ModuleShieldGeneratorManager with generic multiplied object management

Projects/TerraTech/TerraTech/ModuleShieldGeneratorManager.cs

@@ -4,9 +4,7 @@ using BepInEx.Configuration;
 using HarmonyLib;
 
 namespace TerraTech {
-    [HarmonyPatch]
+    public class MultipliedField<TObject, TValue> {
-    public class ModuleShieldGeneratorManager {
-        private class MultipliedField<TObject, TValue> {
         private readonly string fieldName;
         private readonly ConfigEntry<TValue> multiplier;
         private readonly Traverse parentTraverse;
@@ -18,30 +16,17 @@ namespace TerraTech {
             this.parentTraverse = parentTraverse;
         }
 
-            public TValue GetValue() {
+        public TValue GetValue() => (TValue)parentTraverse.Field(fieldName).GetValue();
-                return (TValue)parentTraverse.Field(fieldName).GetValue();
+        public void SetValue(TValue value) => parentTraverse.Field(fieldName).SetValue(value);
-            }
+        public void CaptureOriginal() => originalValue = GetValue();
-
+        public void Restore() => SetValue(originalValue);
-            public void SetValue(TValue value) {
-                parentTraverse.Field(fieldName).SetValue(value);
-            }
-
-            public void CaptureOriginal() {
-                originalValue = GetValue();
-            }
 
         public void Apply() {
-                // This assumes TValue implements multiplication with multiplier.Value
-                // For float, this would be originalValue * multiplier.Value
             dynamic originalDynamic = originalValue;
             dynamic multiplierDynamic = multiplier.Value;
             SetValue((TValue)(originalDynamic * multiplierDynamic));
         }
 
-            public void Restore() {
-                SetValue(originalValue);
-            }
-
         public void LogValue(string prefix) {
             if (!Main.debug.Value)
                 return;
@@ -49,81 +34,100 @@ namespace TerraTech {
         }
     }
 
-        private class MultipliedObject<T> {
+    public class MultipliedObject<T> {
         private readonly Traverse objectTraverse;
         private readonly List<MultipliedField<T, float>> fields = new List<MultipliedField<T, float>>();
 
-            public MultipliedObject(T instance) objectTraverse = Traverse.Create(instance);
+        public MultipliedObject(T instance) => objectTraverse = Traverse.Create(instance);
-            public void AddField(string fieldName, ConfigEntry<float> multiplier) fields.Add(
+        public void AddField(string fieldName, ConfigEntry<float> multiplier) =>
-                new MultipliedField<T, float>(fieldName, multiplier, objectTraverse));
+            fields.Add(new MultipliedField<T, float>(fieldName, multiplier, objectTraverse));
-            public void CaptureFrom() foreach (var field in fields) field.CaptureOriginal();
+        public void CaptureFrom() => fields.ForEach(f => f.CaptureOriginal());
-            public void ApplyTo() foreach (var field in fields) field.Apply();
+        public void ApplyTo() => fields.ForEach(f => f.Apply());
-            public void RestoreTo() foreach (var field in fields) field.Restore();
+        public void RestoreTo() => fields.ForEach(f => f.Restore());
-
         public void LogValues(string prefix) {
             if (!Main.debug.Value)
                 return;
-                foreach (var field in fields) field.LogValue(prefix);
+            fields.ForEach(f => f.LogValue(prefix));
         }
     }
 
-        private static Dictionary<ModuleShieldGenerator, MultipliedObject<ModuleShieldGenerator>> originalProperties =
+    public class MultipliedObjectManager<T>
-            new Dictionary<ModuleShieldGenerator, MultipliedObject<ModuleShieldGenerator>>();
+        where T : class {
+        private readonly Dictionary<T, MultipliedObject<T>> managedObjects = new Dictionary<T, MultipliedObject<T>>();
+        private readonly Action<MultipliedObject<T>> configureObject;
+
+        public MultipliedObjectManager(Action<MultipliedObject<T>> configureObject) {
+            this.configureObject = configureObject;
+        }
+
+        public void OnObjectAttached(T instance) {
+            if (!managedObjects.ContainsKey(instance)) {
+                if (Main.debug.Value)
+                    Console.WriteLine($"{typeof(T).Name}.OnAttached");
+
+                var multipliedObject = new MultipliedObject<T>(instance);
+                configureObject(multipliedObject);
+                multipliedObject.CaptureFrom();
+                managedObjects.Add(instance, multipliedObject);
+                multipliedObject.LogValues("Patching");
+
+                ApplyTo(instance);
+                multipliedObject.LogValues("Patched");
+            }
+        }
+
+        public void OnObjectDetached(T instance) {
+            if (managedObjects.TryGetValue(instance, out var multipliedObject)) {
+                if (Main.debug.Value) {
+                    Console.WriteLine($"{typeof(T).Name}.OnDetaching");
+                    multipliedObject.LogValues("Restoring");
+                }
+
+                RestoreTo(instance);
+                multipliedObject.LogValues("Restored");
+                managedObjects.Remove(instance);
+            }
+        }
+
+        public void ApplyAll() {
+            if (Main.debug.Value)
+                Console.WriteLine($"Modifying {managedObjects.Count} {typeof(T).Name}");
+            foreach (var instance in managedObjects.Keys) {
+                RestoreTo(instance);
+                ApplyTo(instance);
+            }
+        }
+
+        private void ApplyTo(T instance) {
+            if (managedObjects.TryGetValue(instance, out var obj))
+                obj.ApplyTo();
+        }
+
+        private void RestoreTo(T instance) {
+            if (managedObjects.TryGetValue(instance, out var obj))
+                obj.RestoreTo();
+        }
+    }
+
+    [HarmonyPatch]
+    public class ModuleShieldGeneratorManager {
+        private static readonly MultipliedObjectManager<ModuleShieldGenerator> manager =
+            new MultipliedObjectManager<ModuleShieldGenerator>(ConfigureShieldGenerator);
+
+        private static void ConfigureShieldGenerator(MultipliedObject<ModuleShieldGenerator> obj) {
+            obj.AddField("m_HealingHeartbeatInterval", Main.shieldHeartbeatIntervalMultiplier);
+            obj.AddField("m_Radius", Main.shieldRadiusMultiplier);
+            obj.AddField("m_PowerUpDelay", Main.powerUpDelayMultiplier);
+        }
 
         [HarmonyPrefix]
         [HarmonyPatch(typeof(ModuleShieldGenerator), "OnAttached")]
-        static void PostfixCreate(ModuleShieldGenerator __instance) {
+        static void PostfixCreate(ModuleShieldGenerator __instance) => manager.OnObjectAttached(__instance);
-            if (Main.debug.Value)
-                Console.WriteLine("ModuleShieldGenerator.OnAttached");
-
-            if (!originalProperties.ContainsKey(__instance)) {
-                var props = new MultipliedObject<ModuleShieldGenerator>(__instance);
-                props.AddField("m_HealingHeartbeatInterval", Main.shieldHeartbeatIntervalMultiplier);
-                props.AddField("m_Radius", Main.shieldRadiusMultiplier);
-                props.AddField("m_PowerUpDelay", Main.powerUpDelayMultiplier);
-
-                props.CaptureFrom();
-                originalProperties.Add(__instance, props);
-                props.LogValues("Patching");
-            }
-
-            DoPatchSingle(__instance);
-
-            if (Main.debug.Value) {
-                originalProperties[__instance].LogValues("Patched");
-            }
-        }
 
         [HarmonyPrefix]
         [HarmonyPatch(typeof(ModuleShieldGenerator), "OnDetaching")]
-        static void PostfixDestroy(ModuleShieldGenerator __instance) {
+        static void PostfixDestroy(ModuleShieldGenerator __instance) => manager.OnObjectDetached(__instance);
-            if (Main.debug.Value) {
-                Console.WriteLine("ModuleShieldGenerator.OnDetaching");
-                originalProperties[__instance].LogValues("Restoring");
-            }
 
-            DoRestoreSingle(__instance);
+        public static void DoPatch() => manager.ApplyAll();
-
-            if (Main.debug.Value) {
-                originalProperties[__instance].LogValues("Restored");
-            }
-
-            originalProperties.Remove(__instance);
-        }
-
-        public static void DoPatch() {
-            if (Main.debug.Value)
-                Console.WriteLine("Modifying {0} ModuleShieldGenerator", originalProperties.Count);
-            foreach (var kvp in originalProperties) {
-                DoRestoreSingle(kvp.Key);
-                DoPatchSingle(kvp.Key);
-            }
-        }
-
-        static void DoPatchSingle(ModuleShieldGenerator moduleShieldGenerator) if (
-            originalProperties.TryGetValue(moduleShieldGenerator, out var props)) props.ApplyTo();
-
-        static void DoRestoreSingle(ModuleShieldGenerator moduleShieldGenerator) if (
-            originalProperties.TryGetValue(moduleShieldGenerator, out var props)) props.RestoreTo();
     }
 }