From 1711ac2fb4e809567b1a9b334e41cecd735ac7a8 Mon Sep 17 00:00:00 2001
From: Metious <71298690+Metious@users.noreply.github.com>
Date: Mon, 6 May 2024 04:13:52 +0330
Subject: [PATCH] feat: Ban custom prefabs from getting cached (#545)

* Custom prefabs now can opt-out from cache

* Extension method for custom prefab to remove from cache

* Use taskResult instead of looking in cache

* Added more logs
---
 Nautilus/Assets/CustomPrefabExtensions.cs | 32 +++++++++++++++++++++++
 Nautilus/Assets/ModPrefabCache.cs         | 23 +++++++++++++---
 Nautilus/Assets/ModPrefabRequest.cs       |  3 ++-
 3 files changed, 53 insertions(+), 5 deletions(-)
 create mode 100644 Nautilus/Assets/CustomPrefabExtensions.cs

diff --git a/Nautilus/Assets/CustomPrefabExtensions.cs b/Nautilus/Assets/CustomPrefabExtensions.cs
new file mode 100644
index 00000000..23081449
--- /dev/null
+++ b/Nautilus/Assets/CustomPrefabExtensions.cs
@@ -0,0 +1,32 @@
+using Nautilus.Utility;
+
+namespace Nautilus.Assets;
+
+/// <summary>
+/// Represents extension methods for the <see cref="CustomPrefab"/> class.
+/// </summary>
+public static class CustomPrefabExtensions
+{
+    /// <summary>
+    /// Removes the current prefab from the prefab cache and doesn't allow it to get cached later.
+    /// </summary>
+    /// <param name="customPrefab">The custom prefab to remove from the prefab cache.</param>
+    /// <returns>A reference to this instance after the operation has completed.</returns>
+    public static ICustomPrefab RemoveFromCache(this CustomPrefab customPrefab)
+    {
+        if (customPrefab.Info == default)
+        {
+            return customPrefab;
+        }
+
+        if (string.IsNullOrWhiteSpace(customPrefab.Info.ClassID))
+        {
+            InternalLogger.Error($"Couldn't remove prefab '{customPrefab.Info}' from cache because the class ID is null.");
+            return customPrefab;
+        }
+        
+        ModPrefabCache.RemovePrefabFromCache(customPrefab.Info.ClassID);
+
+        return customPrefab;
+    }
+}
\ No newline at end of file
diff --git a/Nautilus/Assets/ModPrefabCache.cs b/Nautilus/Assets/ModPrefabCache.cs
index f95564d4..df6a7a2a 100644
--- a/Nautilus/Assets/ModPrefabCache.cs
+++ b/Nautilus/Assets/ModPrefabCache.cs
@@ -1,6 +1,7 @@
 using Nautilus.Utility;
 using System.Collections.Generic;
 using Nautilus.Extensions;
+using Nautilus.Handlers;
 using UnityEngine;
 
 namespace Nautilus.Assets;
@@ -16,7 +17,7 @@ public static class ModPrefabCache
     private static ModPrefabCacheInstance _cacheInstance;
 
     /// <summary> Adds the given prefab to the cache. </summary>
-    /// <param name="prefab"> The prefab object that is disabled and cached. </param>
+    /// <param name="prefab"> The prefab object that is disabled and cached.</param>
     public static void AddPrefab(GameObject prefab)
     {
         EnsureCacheExists();
@@ -41,8 +42,12 @@ public static bool IsPrefabCached(string classId)
     /// <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)
+        if (_cacheInstance == null)
+        {
+            InternalLogger.Debug($"Removed '{classId}' from prefab cache.");
+            ModPrefabCacheInstance.BannedPrefabs.Add(classId);
             return;
+        }
 
         _cacheInstance.RemoveCachedPrefab(classId);
     }
@@ -74,7 +79,10 @@ private static void EnsureCacheExists()
 
 internal class ModPrefabCacheInstance: MonoBehaviour
 {
-    public Dictionary<string, GameObject> Entries { get; } = new Dictionary<string, GameObject>();
+    public Dictionary<string, GameObject> Entries { get; } = new();
+    
+    // Prefabs that are banned from getting cached
+    internal static readonly HashSet<string> BannedPrefabs = new();
 
     private Transform _prefabRoot;
 
@@ -101,6 +109,11 @@ public void EnterPrefabIntoCache(GameObject prefab)
             return;
         }
 
+        if (BannedPrefabs.Contains(prefabIdentifier.classId))
+        {
+            return;
+        }
+
         if (!Entries.ContainsKey(prefabIdentifier.classId))
         {
             Entries.Add(prefabIdentifier.classId, prefab);
@@ -125,6 +138,8 @@ public void EnterPrefabIntoCache(GameObject prefab)
 
     public void RemoveCachedPrefab(string classId)
     {
+        BannedPrefabs.Add(classId);
+
         if (!Entries.TryGetValue(classId, out var prefab))
         {
             return;
@@ -142,8 +157,8 @@ public void RemoveCachedPrefab(string classId)
             Destroy(prefab);
         }
             
-        InternalLogger.Debug($"ModPrefabCache: removing prefab '{classId}'");
         Entries.Remove(classId);
+        InternalLogger.Debug($"ModPrefabCache: removing prefab '{classId}'");
     }
 
     private void RemoveFakePrefabs()
diff --git a/Nautilus/Assets/ModPrefabRequest.cs b/Nautilus/Assets/ModPrefabRequest.cs
index 35e7d308..a20972e0 100644
--- a/Nautilus/Assets/ModPrefabRequest.cs
+++ b/Nautilus/Assets/ModPrefabRequest.cs
@@ -50,7 +50,7 @@ public object Current
 
     public bool TryGetPrefab(out GameObject result)
     {
-        return ModPrefabCache.TryGetPrefabFromCache(prefabInfo.ClassID, out result) && result != null;
+        return result = taskResult.Get();
     }
 
     public bool MoveNext()
@@ -66,6 +66,7 @@ public bool MoveNext()
     public void Reset()
     {
         task.Reset();
+        taskResult.Set(null);
         Done = false;
     }