Dustbin: fix issue for storage systems on multi-planet

Dustbin/Dustbin.cs

@@ -8,56 +8,12 @@ using HarmonyLib;
 
 namespace Dustbin;
 
-public class IsDusbinIndexer
-{
-    private bool[] _store = new bool[256];
-
-    public bool this[int index]
-    {
-        get
-        {
-            if (index < 0 || index >= _store.Length) return false;
-            return _store[index];
-        }
-        set
-        {
-            if (index >= _store.Length)
-            {
-                var oldLen = _store.Length;
-                var newLen = oldLen * 2;
-                var oldArr = _store;
-                _store = new bool[newLen];
-                Array.Copy(oldArr, _store, oldLen);
-            }
-            _store[index] = value;
-        }
-    }
-
-    public void Reset()
-    {
-        _store = new bool[256];
-    }
-
-    public delegate void ForEachFunc(int id);
-    public void ForEachIsDustbin(ForEachFunc func)
-    {
-        var len = _store.Length;
-        for (var i = 0; i < len; i++)
-        {
-            if (_store[i])
-            {
-                func(i);
-            }
-        }
-    }
-}
-
 [BepInPlugin(PluginInfo.PLUGIN_GUID, PluginInfo.PLUGIN_NAME, PluginInfo.PLUGIN_VERSION)]
 [BepInDependency(DSPModSavePlugin.MODGUID)]
 public class Dustbin : BaseUnityPlugin, IModCanSave
 {
     private const ushort ModSaveVersion = 1;
-    private new static readonly BepInEx.Logging.ManualLogSource Logger =
+    public new static readonly BepInEx.Logging.ManualLogSource Logger =
         BepInEx.Logging.Logger.CreateLogSource(PluginInfo.PLUGIN_NAME);
 
     private bool _cfgEnabled = true;

Dustbin/StoragePatch.cs

@@ -1,45 +1,92 @@
-using System.IO;
+using System.Collections.Generic;
+using System.IO;
+using System.Reflection.Emit;
 using HarmonyLib;
 using UnityEngine;
 
 namespace Dustbin;
 
+public class StorageComponentWithDustbin : StorageComponent
+{
+    public bool IsDusbin;
+    public StorageComponentWithDustbin(int size): base(size)
+    {
+        IsDusbin = false;
+    }
+}
+
 [HarmonyPatch]
 public static class StoragePatch
 {
     private static MyCheckBox _storageDustbinCheckBox;
-    private static int lastStorageId;
-    private static IsDusbinIndexer storageIsDustbin = new();
+    private static int _lastStorageId;
 
     public static void Reset()
     {
-        storageIsDustbin.Reset();
-        lastStorageId = 0;
+        _lastStorageId = 0;
     }
 
     public static void Export(BinaryWriter w)
     {
         var tempStream = new MemoryStream();
         var tempWriter = new BinaryWriter(tempStream);
-        int count = 0;
-        storageIsDustbin.ForEachIsDustbin(i =>
+        var factories = GameMain.data.factories;
+        var factoryCount = GameMain.data.factoryCount;
+        for (var i = 0; i < factoryCount; i++)
         {
-            tempWriter.Write(i);
-            count++;
-        });
-        w.Write(count);
-        tempStream.Position = 0;
-        /* FixMe: May BinaryWriter not sync with its BaseStream while subclass overrides Write()? */
-        tempStream.CopyTo(w.BaseStream);
+            var factory = factories[i];
+            var storage = factory?.factoryStorage;
+            var storagePool = storage?.storagePool;
+            if (storagePool == null) continue;
+            var cursor = storage.storageCursor;
+            var count = 0;
+            for (var j = 1; j < cursor; j++)
+            {
+                if (storagePool[j] == null || storagePool[j].id != j) continue;
+                if (storagePool[j] is not StorageComponentWithDustbin { IsDusbin: true }) continue;
+                tempWriter.Write(j);
+                count++;
+            }
+            if (count == 0) continue;
+
+            tempWriter.Flush();
+            tempStream.Position = 0;
+            w.Write((byte)1);
+            w.Write(factory.planetId);
+            w.Write(count);
+            /* FixMe: May BinaryWriter not sync with its BaseStream while subclass overrides Write()? */
+            tempStream.CopyTo(w.BaseStream);
+            tempStream.SetLength(0);
+        }
         tempWriter.Dispose();
         tempStream.Dispose();
     }
 
     public static void Import(BinaryReader r)
     {
-        for (var count = r.ReadInt32(); count > 0; count--)
+        while (r.PeekChar() == 1)
         {
-            storageIsDustbin[r.ReadInt32()] = true;
+            r.ReadByte();
+            var planetId = r.ReadInt32();
+            var planet = GameMain.data.galaxy.PlanetById(planetId);
+            var storagePool = planet?.factory?.factoryStorage?.storagePool;
+            if (storagePool == null)
+            {
+                for (var count = r.ReadInt32(); count > 0; count--)
+                {
+                    r.ReadInt32();
+                }
+                continue;
+            }
+            for (var count = r.ReadInt32(); count > 0; count--)
+            {
+                var id = r.ReadInt32();
+                if (id > 0 && id < storagePool.Length && storagePool[id] != null && storagePool[id].id == id &&
+                    storagePool[id] is StorageComponentWithDustbin comp)
+                {
+                    comp.IsDusbin = true;
+                }
+            }
         }
     }
 
@@ -52,8 +99,11 @@ public static class StoragePatch
         _storageDustbinCheckBox.OnChecked += () =>
         {
             var storageId = window.storageId;
-            if (storageId <= 0 || window.factoryStorage.storagePool[storageId].id != storageId) return;
-            storageIsDustbin[storageId] = _storageDustbinCheckBox.Checked;
+            if (storageId <= 0) return;
+            var storagePool = window.factoryStorage.storagePool;
+            if (storagePool[storageId].id != storageId) return;
+            if (storagePool[storageId] is not StorageComponentWithDustbin comp) return;
+            comp.IsDusbin = _storageDustbinCheckBox.Checked;
         };
     }
     
@@ -62,12 +112,13 @@ public static class StoragePatch
     private static void UIStorageWindow__OnUpdate_Postfix(UIStorageWindow __instance)
     {
         var storageId = __instance.storageId;
-        if (lastStorageId == storageId) return;
-        lastStorageId = storageId;
-        if (storageId > 0 && __instance.factoryStorage.storagePool[storageId].id == storageId)
-        {
-            _storageDustbinCheckBox.Checked = storageIsDustbin[storageId];
-        }
+        if (_lastStorageId == storageId) return;
+        _lastStorageId = storageId;
+        if (storageId <= 0) return;
+        var storagePool = __instance.factoryStorage.storagePool;
+        if (storagePool[storageId].id != storageId) return;
+        if (storagePool[storageId] is not StorageComponentWithDustbin comp) return;
+        _storageDustbinCheckBox.Checked = comp.IsDusbin;
         if (__instance.transform is RectTransform rectTrans) {
             _storageDustbinCheckBox.rectTrans.anchoredPosition3D = new Vector3(50, 58 - rectTrans.sizeDelta.y, 0);
         }
@@ -83,7 +134,7 @@ public static class StoragePatch
     private static bool StorageComponent_AddItem_Prefix(ref int __result, StorageComponent __instance, int itemId, int count, int inc,
         ref int remainInc, bool useBan = false)
     {
-        if (!storageIsDustbin[__instance.id]) return true;
+        if (__instance is not StorageComponentWithDustbin { IsDusbin: true }) return true;
         remainInc = inc;
         __result = count;
         var fluidArr = Dustbin.IsFluid;
@@ -107,25 +158,34 @@ public static class StoragePatch
         return false;
     }
 
-    [HarmonyPrefix]
-    [HarmonyPatch(typeof(FactoryStorage), "RemoveStorageComponent")]
-    private static void FactoryStorage_RemoveStorageComponent_Prefix(FactoryStorage __instance, int id)
-    {
-        var storage = __instance.storagePool[id];
-        if (storage != null && storage.id != 0)
-        {
-            storageIsDustbin[id] = false;
-        }
-    }
-
     /* We keep this to make MOD compatible with older version */
     [HarmonyPostfix]
     [HarmonyPatch(typeof(StorageComponent), "Import")]
     private static void StorageComponent_Import_Postfix(StorageComponent __instance)
     {
-        if (__instance.bans >= 0)
+        if (__instance.bans >= 0 || __instance is not StorageComponentWithDustbin comp)
             return;
-        storageIsDustbin[__instance.id] = true;
         __instance.bans = -__instance.bans - 1;
+        comp.IsDusbin = true;
+    }
+
+    /* Replace: new StorageComponent(int) => new StorageComponentWithDustbin(int) */
+    [HarmonyTranspiler]
+    [HarmonyPatch(typeof(FactoryStorage), "Import")]
+    [HarmonyPatch(typeof(FactoryStorage), "NewStorageComponent")]
+    private static IEnumerable<CodeInstruction> FactoryStorage_NewStorageComponent_Transpiler(
+        IEnumerable<CodeInstruction> instructions)
+    {
+        foreach (var instr in instructions)
+        {
+            if (instr.opcode == OpCodes.Newobj && instr.OperandIs(AccessTools.Constructor(typeof(StorageComponent), new [] { typeof(int) })))
+            {
+                yield return new CodeInstruction(OpCodes.Newobj, AccessTools.Constructor(typeof(StorageComponentWithDustbin), new [] { typeof(int) }));
+            }
+            else
+            {
+                yield return instr;
+            }
+        }
     }
 }

Dustbin/TankPatch.cs

@@ -1,44 +1,139 @@
-using System.IO;
+using System;
+using System.IO;
 using HarmonyLib;
 
 namespace Dustbin;
 
+using IsDustbinIndexer = DynamicObjectArray<DynamicObjectArray<DynamicValueArray<bool>>>;
+
+public class DynamicValueArray<T> where T: struct
+{
+    private T[] _store = new T[16];
+
+    public T this[int index]
+    {
+        get
+        {
+            if (index < 0 || index >= _store.Length) return default;
+            return _store[index];
+        }
+        set
+        {
+            if (index >= _store.Length)
+            {
+                Array.Resize(ref _store, _store.Length * 2);
+            }
+            _store[index] = value;
+        }
+    }
+
+    public void Reset()
+    {
+        _store = new T[16];
+    }
+
+    public delegate void ForEachFunc(int id, T value);
+    public void ForEach(ForEachFunc func)
+    {
+        var len = _store.Length;
+        for (var i = 0; i < len; i++)
+        {
+            func(i, _store[i]);
+        }
+    }
+}
+
+public class DynamicObjectArray<T> where T: class, new()
+{
+    private T[] _store = new T[16];
+
+    public T this[int index]
+    {
+        get
+        {
+            if (index < 0 || index >= _store.Length) return null;
+            var result = _store[index];
+            if (result == null)
+            {
+                result = new T();
+                _store[index] = result;
+            }
+            return result;
+        }
+    }
+
+    public void Reset()
+    {
+        _store = new T[16];
+    }
+
+    public delegate void ForEachFunc(int id, T value);
+    public void ForEach(ForEachFunc func)
+    {
+        var len = _store.Length;
+        for (var i = 0; i < len; i++)
+        {
+            func(i, _store[i]);
+        }
+    }
+}
+
 [HarmonyPatch]
 public static class TankPatch
 {
     private static MyCheckBox _tankDustbinCheckBox;
     private static int lastTankId;
-    private static IsDusbinIndexer tankIsDustbin = new();
+    private static IsDustbinIndexer tankIsDustbin = new();
 
     public static void Reset()
     {
-        tankIsDustbin.Reset();
         lastTankId = 0;
+        tankIsDustbin.Reset();
     }
 
     public static void Export(BinaryWriter w)
     {
         var tempStream = new MemoryStream();
         var tempWriter = new BinaryWriter(tempStream);
-        int count = 0;
-        tankIsDustbin.ForEachIsDustbin(i =>
+        tankIsDustbin.ForEach((i, star) =>
         {
-            tempWriter.Write(i);
-            count++;
+            star?.ForEach((j, planet) =>
+            {
+                var count = 0;
+                planet?.ForEach((id, v) =>
+                {
+                    if (!v) return;
+                    tempWriter.Write(id);
+                    count++;
+                });
+                if (count == 0) return;
+
+                tempWriter.Flush();
+                tempStream.Position = 0;
+                w.Write((byte)2);
+                var planetId = i * 100 + j;
+                w.Write(planetId);
+                w.Write(count);
+                /* FixMe: May BinaryWriter not sync with its BaseStream while subclass overrides Write()? */
+                tempStream.CopyTo(w.BaseStream);
+                tempStream.SetLength(0);
+            });
         });
-        w.Write(count);
-        tempStream.Position = 0;
-        /* FixMe: May BinaryWriter not sync with its BaseStream while subclass overrides Write()? */
-        tempStream.CopyTo(w.BaseStream);
         tempWriter.Dispose();
         tempStream.Dispose();
     }
 
     public static void Import(BinaryReader r)
     {
-        for (var count = r.ReadInt32(); count > 0; count--)
+        while (r.PeekChar() == 2)
         {
-            tankIsDustbin[r.ReadInt32()] = true;
+            r.ReadByte();
+            var planetId = r.ReadInt32();
+            var data = tankIsDustbin[planetId / 100][planetId % 100];
+            for (var count = r.ReadInt32(); count > 0; count--)
+            {
+                data[r.ReadInt32()] = true;
+            }
         }
     }
 
@@ -52,7 +147,8 @@ public static class TankPatch
         {
             var tankId = window.tankId;
             if (tankId <= 0 || window.storage.tankPool[tankId].id != tankId) return;
-            tankIsDustbin[tankId] = _tankDustbinCheckBox.Checked;
+            var planetId = window.storage.planet.id;
+            tankIsDustbin[planetId / 100][planetId % 100][tankId] = _tankDustbinCheckBox.Checked;
         };
     }
     
@@ -65,15 +161,17 @@ public static class TankPatch
         lastTankId = tankId;
         if (tankId > 0 && __instance.storage.tankPool[tankId].id == tankId)
         {
-            _tankDustbinCheckBox.Checked = tankIsDustbin[tankId];
+            var planetId = __instance.storage.planet.id;
+            _tankDustbinCheckBox.Checked = tankIsDustbin[planetId / 100][planetId % 100][tankId];
         }
     }
 
     [HarmonyPrefix]
     [HarmonyPatch(typeof(TankComponent), "GameTick")]
-    private static void TankComponent_GameTick_Prefix(ref TankComponent __instance, out long __state)
+    private static void TankComponent_GameTick_Prefix(ref TankComponent __instance, out long __state, PlanetFactory factory)
     {
-        if (tankIsDustbin[__instance.id])
+        var planetId = factory.planetId;
+        if (tankIsDustbin[planetId / 100][planetId % 100][__instance.id])
         {
             __state = ((long)__instance.fluidInc << 36) | ((long)__instance.fluidCount << 16) | (uint)__instance.fluidId;
             __instance.fluidId = __instance.fluidCount = __instance.fluidInc = 0;
@@ -100,17 +198,8 @@ public static class TankPatch
     {
         if (__instance.tankPool[id].id != 0)
         {
-            tankIsDustbin[id] = false;
+            var planetId = __instance.planet.id;
+            tankIsDustbin[planetId / 100][planetId % 100][id] = false;
         }
     }
-
-    /* We keep this to make MOD compatible with older version */
-    [HarmonyPostfix]
-    [HarmonyPatch(typeof(TankComponent), "Import")]
-    private static void TankComponent_Import_Postfix(TankComponent __instance)
-    {
-        if (__instance.fluidInc >= 0) return;
-        tankIsDustbin[__instance.id] = true;
-        __instance.fluidInc = -__instance.fluidInc - 1;
-    }
 }