From 7556d911bfc60ab775f948f8a060584eabe2d94e Mon Sep 17 00:00:00 2001
From: Nick <nicholasconnors@hotmail.ca>
Date: Fri, 5 Aug 2022 19:49:55 -0400
Subject: [PATCH] Only unload stuff that isn't in use

---
 NewHorizons/Handlers/StreamingHandler.cs      | 45 +++++++++++++------
 .../Patches/StreamingManagerPatches.cs        | 17 +++++++
 2 files changed, 49 insertions(+), 13 deletions(-)
 create mode 100644 NewHorizons/Patches/StreamingManagerPatches.cs

diff --git a/NewHorizons/Handlers/StreamingHandler.cs b/NewHorizons/Handlers/StreamingHandler.cs
index 4ea8dbb0..56d455a6 100644
--- a/NewHorizons/Handlers/StreamingHandler.cs
+++ b/NewHorizons/Handlers/StreamingHandler.cs
@@ -1,4 +1,4 @@
-using System.Collections.Generic;
+using System.Collections.Generic;
 using System.Linq;
 using UnityEngine;
 using Logger = NewHorizons.Utility.Logger;
@@ -12,11 +12,13 @@ namespace NewHorizons.Handlers
     {
         private static readonly Dictionary<Material, string> _materialCache = new();
         private static readonly Dictionary<GameObject, string[]> _objectCache = new();
+        private static readonly Dictionary<string, List<Sector>> _sectorCache = new();
 
         public static void Init()
         {
             _materialCache.Clear();
             _objectCache.Clear();
+            _sectorCache.Clear();
         }
 
         /// <summary>
@@ -71,22 +73,39 @@ namespace NewHorizons.Handlers
             foreach (var assetBundle in assetBundles)
             {
                 StreamingManager.LoadStreamingAssets(assetBundle);
+
+                // Track the sector even if its null. null means stay loaded forever
+                if (!_sectorCache.ContainsKey(assetBundle)) _sectorCache.Add(assetBundle, new List<Sector>());
+                if (!_sectorCache[assetBundle].Contains(sector)) _sectorCache[assetBundle].Add(sector);
             }
 
-            if (!sector)
+            if (sector)
             {
-                Logger.LogWarning($"StreamingHandler for {obj} has null sector. " +
-                    "This can lead to the thing being unloaded permanently.");
-                return;
-            }
-
-            sector.OnOccupantEnterSector += _ =>
-            {
-                foreach (var assetBundle in assetBundles)
+                sector.OnOccupantEnterSector += _ =>
                 {
-                    StreamingManager.LoadStreamingAssets(assetBundle);
-                }
-            };
+                    foreach (var assetBundle in assetBundles)
+                        StreamingManager.LoadStreamingAssets(assetBundle);
+                };
+                /*
+                sector.OnOccupantExitSector += _ =>
+                {
+                    // UnloadStreamingAssets is patched to check IsBundleInUse first before unloading
+                    if (!sector.ContainsAnyOccupants(DynamicOccupant.Player | DynamicOccupant.Probe))
+                        foreach (var assetBundle in assetBundles)
+                            StreamingManager.UnloadStreamingAssets(assetBundle);
+                };
+                */
+            }
+        }
+
+        public static bool IsBundleInUse(string assetBundle)
+        {
+            // If a sector in the list is null then it is always in use
+            if(_sectorCache.TryGetValue(assetBundle, out var sectors))
+                foreach (var sector in sectors)
+                    if (sector == null || sector.ContainsAnyOccupants(DynamicOccupant.Player | DynamicOccupant.Probe)) 
+                        return true;
+            return false;   
         }
     }
 }
\ No newline at end of file
diff --git a/NewHorizons/Patches/StreamingManagerPatches.cs b/NewHorizons/Patches/StreamingManagerPatches.cs
new file mode 100644
index 00000000..5f422072
--- /dev/null
+++ b/NewHorizons/Patches/StreamingManagerPatches.cs
@@ -0,0 +1,17 @@
+using HarmonyLib;
+using NewHorizons.Handlers;
+
+namespace NewHorizons.Patches
+{
+    [HarmonyPatch]
+    public static class StreamingManagerPatches
+    {
+        [HarmonyPrefix]
+        [HarmonyPatch(typeof(StreamingManager), nameof(StreamingManager.UnloadStreamingAssets))]
+        public static bool StreamingManager_UnloadStreamingAssets(string assetBundleName)
+        {
+            // Only let it unload stuff that isn't being used
+            return !StreamingHandler.IsBundleInUse(assetBundleName);
+        } 
+    }
+}