Moved CustomWidgetChassis logic to avoid NRE

BattletechModders · Dec 1, 2023 · d917050 · d917050
1 parent 6d0398c
commit d917050
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 13 deletions.
diff --git a/source/Features/MechLabSlots/CustomWidgetsFixMechLab.cs b/source/Features/MechLabSlots/CustomWidgetsFixMechLab.cs
@@ -22,33 +22,28 @@ internal static bool IsCustomWidget(MechLabLocationWidget widget)
 
     internal static void Setup(MechLabPanel mechLabPanel)
     {
-        var cwcc = mechLabPanel.activeMechDef.Chassis.GetComponent<CustomWidgetChassisCustom>();
-
         SetupWidget(
             "TopLeftWidget",
             ref TopLeftWidget,
             mechLabPanel,
             mechLabPanel.rightArmWidget,
-            MechLabSlotsFeature.settings.TopLeftWidget,
-            cwcc?.TopLeftWidgetEnabled
+            MechLabSlotsFeature.settings.TopLeftWidget
         );
 
         SetupWidget(
             "TopRightWidget",
             ref TopRightWidget,
             mechLabPanel,
             mechLabPanel.leftArmWidget,
-            MechLabSlotsFeature.settings.TopRightWidget,
-            cwcc?.TopRightWidgetEnabled
+            MechLabSlotsFeature.settings.TopRightWidget
         );
     }
 
     internal static void SetupWidget(string id,
         ref MechLabLocationWidget? topWidget,
         MechLabPanel mechLabPanel,
         MechLabLocationWidget armWidget,
-        MechLabSlotsSettings.WidgetSettings settings,
-        bool? chassisEnabled)
+        MechLabSlotsSettings.WidgetSettings settings)
     {
         GameObject go;
         if (topWidget == null)
@@ -100,11 +95,6 @@ internal static void Setup(MechLabPanel mechLabPanel)
             var mechRectTransform = parent.parent.GetComponent<RectTransform>();
             LayoutRebuilder.ForceRebuildLayoutImmediate(mechRectTransform);
         }
-
-        if (chassisEnabled.HasValue)
-        {
-            go.SetActive(chassisEnabled.Value);
-        }
     }
 
     internal static void OnAdditem_SetParent(Transform @this, Transform parent, bool worldPositionStays)
@@ -202,4 +192,29 @@ internal static void RefreshDropHighlights(MechLabLocationWidget widget, IMechLa
         cRef = null;
         return false;
     }
+
+    internal static void ShowOrHideCustomWidgets(MechLabLocationWidget widget)
+    {
+        if (widget.loadout.Location != ChassisLocations.CenterTorso)
+        {
+            return;
+        }
+
+        var custom = widget.mechLab.activeMechDef.Chassis.GetComponent<CustomWidgetChassisCustom>();
+        if (custom == null)
+        {
+            return;
+        }
+
+        static void ShowOrHideCustomWidget(MechLabLocationWidget? customWidget, bool? enabled)
+        {
+            if (customWidget == null || !enabled.HasValue)
+            {
+                return;
+            }
+            customWidget.gameObject.SetActive(enabled.Value);
+        }
+        ShowOrHideCustomWidget(TopLeftWidget, custom.TopLeftWidgetEnabled);
+        ShowOrHideCustomWidget(TopRightWidget, custom.TopRightWidgetEnabled);
+    }
 }
diff --git a/source/Features/MechLabSlots/Patches/MechLabLocationWidget_SetData_Patch.cs b/source/Features/MechLabSlots/Patches/MechLabLocationWidget_SetData_Patch.cs
@@ -31,6 +31,7 @@ public static void Postfix(MechLabLocationWidget __instance, int ___maxSlots, re
         MechLabSlotsFixer.FixSlots(widgetLayout, ___maxSlots);
         DynamicSlotsFeature.PrepareWidget(widgetLayout);
         AdjustMechLabLocationNaming(widget, loadout.Location);
+        CustomWidgetsFixMechLab.ShowOrHideCustomWidgets(widget);
     }
 
     private static void AdjustMechLabLocationNaming(MechLabLocationWidget widget, ChassisLocations location)