1 files changed, 66 insertions, 13 deletions
diff --git a/src/Ryujinx.Graphics.Gpu/Image/TextureCache.cs b/src/Ryujinx.Graphics.Gpu/Image/TextureCache.cs
index c3243cf2..97d78a34 100644
--- a/src/Ryujinx.Graphics.Gpu/Image/TextureCache.cs
+++ b/src/Ryujinx.Graphics.Gpu/Image/TextureCache.cs
@@ -64,7 +64,7 @@ namespace Ryujinx.Graphics.Gpu.Image
         }
 
         /// <summary>
-        /// Handles removal of textures written to a memory region being unmapped.
+        /// Handles marking of textures written to a memory region being (partially) remapped.
         /// </summary>
         /// <param name="sender">Sender object</param>
         /// <param name="e">Event arguments</param>
@@ -80,26 +80,41 @@ namespace Ryujinx.Graphics.Gpu.Image
                 overlapCount = _textures.FindOverlaps(unmapped, ref overlaps);
             }
 
-            for (int i = 0; i < overlapCount; i++)
+            if (overlapCount > 0)
             {
-                overlaps[i].Unmapped(unmapped);
+                for (int i = 0; i < overlapCount; i++)
+                {
+                    overlaps[i].Unmapped(unmapped);
+                }
             }
 
-            // If any range was previously unmapped, we also need to purge
-            // all partially mapped texture, as they might be fully mapped now.
-            for (int i = 0; i < unmapped.Count; i++)
+            lock (_partiallyMappedTextures)
             {
-                if (unmapped.GetSubRange(i).Address == MemoryManager.PteUnmapped)
+                if (overlapCount > 0 || _partiallyMappedTextures.Count > 0)
                 {
-                    lock (_partiallyMappedTextures)
+                    e.AddRemapAction(() =>
                     {
-                        foreach (var texture in _partiallyMappedTextures)
+                        lock (_partiallyMappedTextures)
                         {
-                            texture.Unmapped(unmapped);
-                        }
-                    }
+                            if (overlapCount > 0)
+                            {
+                                for (int i = 0; i < overlapCount; i++)
+                                {
+                                    _partiallyMappedTextures.Add(overlaps[i]);
+                                }
+                            }
 
-                    break;
+                            // Any texture that has been unmapped at any point or is partially unmapped
+                            // should update their pool references after the remap completes.
+
+                            MultiRange unmapped = ((MemoryManager)sender).GetPhysicalRegions(e.Address, e.Size);
+
+                            foreach (var texture in _partiallyMappedTextures)
+                            {
+                                texture.UpdatePoolMappings();
+                            }
+                        }
+                    });
                 }
             }
         }
@@ -1136,6 +1151,44 @@ namespace Ryujinx.Graphics.Gpu.Image
         }
 
         /// <summary>
+        /// Queries a texture's memory range and marks it as partially mapped or not.
+        /// Partially mapped textures re-evaluate their memory range after each time GPU memory is mapped.
+        /// </summary>
+        /// <param name="memoryManager">GPU memory manager where the texture is mapped</param>
+        /// <param name="address">The virtual address of the texture</param>
+        /// <param name="texture">The texture to be marked</param>
+        /// <returns>The physical regions for the texture, found when evaluating whether the texture was partially mapped</returns>
+        public MultiRange UpdatePartiallyMapped(MemoryManager memoryManager, ulong address, Texture texture)
+        {
+            MultiRange range;
+            lock (_partiallyMappedTextures)
+            {
+                range = memoryManager.GetPhysicalRegions(address, texture.Size);
+                bool partiallyMapped = false;
+
+                for (int i = 0; i < range.Count; i++)
+                {
+                    if (range.GetSubRange(i).Address == MemoryManager.PteUnmapped)
+                    {
+                        partiallyMapped = true;
+                        break;
+                    }
+                }
+
+                if (partiallyMapped)
+                {
+                    _partiallyMappedTextures.Add(texture);
+                }
+                else
+                {
+                    _partiallyMappedTextures.Remove(texture);
+                }
+            }
+
+            return range;
+        }
+
+        /// <summary>
         /// Adds a texture to the short duration cache. This typically keeps it alive for two ticks.
         /// </summary>
         /// <param name="texture">Texture to add to the short cache</param>