From 4b911263eb7f2cd031a98656c2b27c5ae19d8a79 Mon Sep 17 00:00:00 2001
From: Soar Qin <soarchin@gmail.com>
Date: Sat, 30 Nov 2024 23:09:16 +0800
Subject: [PATCH] fix a possible crash

---
 UXAssist/Functions/DysonSphereFunctions.cs | 36 ++++++++++++++++++++--
 1 file changed, 33 insertions(+), 3 deletions(-)

diff --git a/UXAssist/Functions/DysonSphereFunctions.cs b/UXAssist/Functions/DysonSphereFunctions.cs
index 45ef590..04984d5 100644
--- a/UXAssist/Functions/DysonSphereFunctions.cs
+++ b/UXAssist/Functions/DysonSphereFunctions.cs
@@ -1,4 +1,6 @@
-namespace UXAssist.Functions;
+using System.Linq;
+
+namespace UXAssist.Functions;
 
 public static class DysonSphereFunctions
 {
@@ -19,26 +21,54 @@ public static class DysonSphereFunctions
         if (star == null) return;
         var dysonSpheres = GameMain.data?.dysonSpheres;
         if (dysonSpheres == null) return;
+        var dysonEditor = UIRoot.instance?.uiGame.dysonEditor;
         if (index < 0)
         {
             if (dysonSpheres[star.index] == null) return;
+            var idsToRemove = dysonSpheres[star.index].layersSorted.Select(layer => layer.id).ToArray();
             var dysonSphere = new DysonSphere();
             dysonSpheres[star.index] = dysonSphere;
             dysonSphere.Init(GameMain.data, star);
             dysonSphere.ResetNew();
+            if (dysonEditor == null) return;
+            foreach (var id in idsToRemove)
+            {
+                if (!dysonEditor.IsRender(id, false, true))
+                {
+                    dysonEditor.SwitchRenderState(id, false, true);
+                }
+                if (!dysonEditor.IsRender(id, false, false))
+                {
+                    dysonEditor.SwitchRenderState(id, false, false);
+                }
+            }
+            dysonEditor.selection.ClearAllSelection();
             return;
         }
 
         var ds = dysonSpheres[star.index];
         if (ds?.layersIdBased[index] == null) return;
         var pool = ds.rocketPool;
+        var idToRemove = -1;
         for (var id = ds.rocketCursor - 1; id > 0; id--)
         {
             if (pool[id].id != id) continue;
             if (pool[id].nodeLayerId != index) continue;
             ds.RemoveDysonRocket(id);
+            idToRemove = id;
+            break;
         }
-
-        ds.RemoveLayer(index);
+        if (idToRemove < 0) return;
+        ds.RemoveLayer(idToRemove);
+        if (dysonEditor == null) return;
+        if (!dysonEditor.IsRender(idToRemove, false, true))
+        {
+            dysonEditor.SwitchRenderState(idToRemove, false, true);
+        }
+        if (!dysonEditor.IsRender(idToRemove, false, false))
+        {
+            dysonEditor.SwitchRenderState(idToRemove, false, false);
+        }
+        dysonEditor.selection.ClearAllSelection();
     }
 }