From 9464432776d4d4d3ded651e4a3c97b902af20d9b Mon Sep 17 00:00:00 2001
From: Metious <71298690+Metious@users.noreply.github.com>
Date: Thu, 25 Apr 2024 06:01:18 +0330
Subject: [PATCH] fix: Prefab cache now resets (#544)

* Prefab cache now resets on world quit

* Fixed `GameObjectExtensions.IsPrefab` sometimes failing

* Fixed null check
---
 Nautilus/Assets/ModPrefabCache.cs           | 43 +++++++++++++++++++--
 Nautilus/Extensions/GameObjectExtensions.cs |  4 +-
 2 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/Nautilus/Assets/ModPrefabCache.cs b/Nautilus/Assets/ModPrefabCache.cs
index e3cecb08..f95564d4 100644
--- a/Nautilus/Assets/ModPrefabCache.cs
+++ b/Nautilus/Assets/ModPrefabCache.cs
@@ -38,6 +38,7 @@ public static bool IsPrefabCached(string classId)
     /// Any prefab with the matching <paramref name="classId"/> will be removed from the cache.
     /// </summary>
     /// <param name="classId">The class id of the prefab that will be removed.</param>
+    /// <remarks>This operation is extremely dangerous on custom prefabs that are directly registering an asset bundle prefab as it may make the prefab unusable in the current session.<br/>Avoid using this method unless you know what you're doing.</remarks>
     public static void RemovePrefabFromCache(string classId)
     {
         if(_cacheInstance == null)
@@ -85,7 +86,9 @@ private void Awake()
 
         gameObject.AddComponent<SceneCleanerPreserve>();
         DontDestroyOnLoad(gameObject);
+        
         SaveUtils.RegisterOnQuitEvent(ModPrefabCache.RunningPrefabs.Clear);
+        SaveUtils.RegisterOnQuitEvent(RemoveFakePrefabs);
     }
 
     public void EnterPrefabIntoCache(GameObject prefab)
@@ -122,12 +125,44 @@ public void EnterPrefabIntoCache(GameObject prefab)
 
     public void RemoveCachedPrefab(string classId)
     {
-        if (Entries.TryGetValue(classId, out var prefab))
+        if (!Entries.TryGetValue(classId, out var prefab))
         {
-            if(!prefab.IsPrefab())
-                Destroy(prefab);
-            InternalLogger.Debug($"ModPrefabCache: removed prefab {classId}");
+            return;
+        }
+        
+        if (!prefab)
+        {
+            InternalLogger.Debug($"ModPrefabCache: Prefab for '{classId}' is null; removing entry.");
             Entries.Remove(classId);
+            return;
+        }
+            
+        if (!prefab.IsPrefab())
+        {
+            Destroy(prefab);
+        }
+            
+        InternalLogger.Debug($"ModPrefabCache: removing prefab '{classId}'");
+        Entries.Remove(classId);
+    }
+
+    private void RemoveFakePrefabs()
+    {
+        foreach (var prefab in new Dictionary<string, GameObject>(Entries))
+        {
+            if (prefab.Value.Exists() is null)
+            {
+                Entries.Remove(prefab.Key);
+                continue;
+            }
+
+            if (prefab.Value.IsPrefab())
+            {
+                continue;
+            }
+            
+            Destroy(prefab.Value);
+            Entries.Remove(prefab.Key);
         }
     }
 
diff --git a/Nautilus/Extensions/GameObjectExtensions.cs b/Nautilus/Extensions/GameObjectExtensions.cs
index 83e63c3a..326d9f76 100644
--- a/Nautilus/Extensions/GameObjectExtensions.cs
+++ b/Nautilus/Extensions/GameObjectExtensions.cs
@@ -5,6 +5,7 @@
 using Nautilus.Utility;
 using UnityEngine;
 using UnityEngine.AddressableAssets;
+using UnityEngine.SceneManagement;
 using Object = UnityEngine.Object;
 
 namespace Nautilus.Extensions;
@@ -136,9 +137,10 @@ public static Transform SearchChild(this Transform transform, string name)
     /// </summary>
     /// <param name="gameObject">The game object to check.</param>
     /// <returns>True if this game object is a proper prefab, otherwise false.</returns>
+    /// <exception cref="System.NullReferenceException"><paramref name="gameObject"/> is null.</exception>
     public static bool IsPrefab(this GameObject gameObject)
     {
-        return gameObject.transform.parent == null && !gameObject.activeInHierarchy && gameObject.activeSelf;
+        return gameObject.scene.name is null && gameObject.scene.loadingState is Scene.LoadingState.NotLoaded;
     }
     
     /// <summary>