From 8f266ca7448597fe9dcb902c69cc28f21243b301 Mon Sep 17 00:00:00 2001
From: Metious <71298690+Metious@users.noreply.github.com>
Date: Sat, 2 Dec 2023 23:03:22 +0330
Subject: [PATCH] fix: Fixed issue when yield returning same request (#498)

* Fixed issue when yield returning same request

The issue occurs because when you yield return the same request multiple times, the state machine advances but it will not wait until the code is completely executed.

* Fixed missing reference when a prefab unloads

* Better access modifiers
---
 Nautilus/Assets/ModPrefabCache.cs          |  2 +-
 Nautilus/Assets/ModPrefabRequest.cs        | 19 +++++++++++++++----
 Nautilus/Patchers/PrefabDatabasePatcher.cs |  7 +++++++
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/Nautilus/Assets/ModPrefabCache.cs b/Nautilus/Assets/ModPrefabCache.cs
index ef74e5a8e..933c546d5 100644
--- a/Nautilus/Assets/ModPrefabCache.cs
+++ b/Nautilus/Assets/ModPrefabCache.cs
@@ -111,7 +111,7 @@ public void EnterPrefabIntoCache(GameObject prefab)
         if(!Entries.ContainsKey(prefabIdentifier.classId))
         {
             Entries.Add(prefabIdentifier.classId, prefab);
-            InternalLogger.Debug($"ModPrefabCache: adding prefab {prefab}");
+            InternalLogger.Debug($"ModPrefabCache: added prefab {prefab}");
         }
         else // this should never happen
         {
diff --git a/Nautilus/Assets/ModPrefabRequest.cs b/Nautilus/Assets/ModPrefabRequest.cs
index dc5e67ed9..c3154ca5f 100644
--- a/Nautilus/Assets/ModPrefabRequest.cs
+++ b/Nautilus/Assets/ModPrefabRequest.cs
@@ -7,12 +7,14 @@
 namespace Nautilus.Assets;
 
 // request for getting ModPrefab asynchronously
-internal class ModPrefabRequest: IPrefabRequest, IEnumerator
+internal class ModPrefabRequest: IPrefabRequest
 {
+    internal bool Done { get; private set; }
+    
     private readonly PrefabInfo prefabInfo;
-
-    private int state = 0;
+    
     private CoroutineTask<GameObject> task;
+    
     private TaskResult<GameObject> taskResult;
 
     public ModPrefabRequest(PrefabInfo prefabInfo)
@@ -50,13 +52,22 @@ public object Current
     public bool TryGetPrefab(out GameObject result)
     {
         result = taskResult.Get();
+        if (!Done)
+        {
+            Done = result;
+        }
         return result != null;
     }
 
     public bool MoveNext()
     {
         Init();
-        return state++ == 0;
+        if (task == null)
+        {
+            return false;
+        }
+        
+        return !TryGetPrefab(out _);
     }
 
     public void Reset() {}
diff --git a/Nautilus/Patchers/PrefabDatabasePatcher.cs b/Nautilus/Patchers/PrefabDatabasePatcher.cs
index 8ea415bc6..172f22511 100644
--- a/Nautilus/Patchers/PrefabDatabasePatcher.cs
+++ b/Nautilus/Patchers/PrefabDatabasePatcher.cs
@@ -80,7 +80,14 @@ private static IPrefabRequest GetModPrefabAsync(string classId)
         }
 
         if(ModPrefabCache.Requests.TryGetValue(prefabInfo.ClassID, out var request))
+        {
+            if (request.Done && !request.TryGetPrefab(out _))
+            {
+                return new ModPrefabRequest(prefabInfo);
+            }
+            
             return request;
+        }
 
         return new ModPrefabRequest(prefabInfo);
     }