7 files changed, 135 insertions, 19 deletions
diff --git a/src/Ryujinx.Graphics.Gpu/Image/Texture.cs b/src/Ryujinx.Graphics.Gpu/Image/Texture.cs
index 84808a84..f0df55e6 100644
--- a/src/Ryujinx.Graphics.Gpu/Image/Texture.cs
+++ b/src/Ryujinx.Graphics.Gpu/Image/Texture.cs
@@ -1610,6 +1610,8 @@ namespace Ryujinx.Graphics.Gpu.Image
         /// </summary>
         public void UpdatePoolMappings()
         {
+            ChangedMapping = true;
+
             lock (_poolOwners)
             {
                 ulong address = 0;
@@ -1685,8 +1687,6 @@ namespace Ryujinx.Graphics.Gpu.Image
             {
                 Group.ClearModified(unmapRange);
             }
-
-            UpdatePoolMappings();
         }
 
         /// <summary>
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>
diff --git a/src/Ryujinx.Graphics.Gpu/Image/TexturePool.cs b/src/Ryujinx.Graphics.Gpu/Image/TexturePool.cs
index 5277e789..dbcb2e75 100644
--- a/src/Ryujinx.Graphics.Gpu/Image/TexturePool.cs
+++ b/src/Ryujinx.Graphics.Gpu/Image/TexturePool.cs
@@ -272,7 +272,15 @@ namespace Ryujinx.Graphics.Gpu.Image
 
                     ulong address = descriptor.UnpackAddress();
 
-                    MultiRange range = _channel.MemoryManager.GetPhysicalRegions(address, texture.Size);
+                    if (!descriptor.Equals(ref DescriptorCache[request.ID]))
+                    {
+                        // If the pool entry has already been replaced, just remove the texture.
+
+                        texture.DecrementReferenceCount();
+                        continue;
+                    }
+
+                    MultiRange range = _channel.MemoryManager.Physical.TextureCache.UpdatePartiallyMapped(_channel.MemoryManager, address, texture);
 
                     // If the texture is not mapped at all, delete its reference.
 
diff --git a/src/Ryujinx.Graphics.Gpu/Memory/MemoryManager.cs b/src/Ryujinx.Graphics.Gpu/Memory/MemoryManager.cs
index b0f7e799..0d4a41f0 100644
--- a/src/Ryujinx.Graphics.Gpu/Memory/MemoryManager.cs
+++ b/src/Ryujinx.Graphics.Gpu/Memory/MemoryManager.cs
@@ -366,6 +366,22 @@ namespace Ryujinx.Graphics.Gpu.Memory
         }
 
         /// <summary>
+        /// Runs remap actions that are added to an unmap event.
+        /// These must run after the mapping completes.
+        /// </summary>
+        /// <param name="e">Event with remap actions</param>
+        private void RunRemapActions(UnmapEventArgs e)
+        {
+            if (e.RemapActions != null)
+            {
+                foreach (Action action in e.RemapActions)
+                {
+                    action();
+                }
+            }
+        }
+
+        /// <summary>
         /// Maps a given range of pages to the specified CPU virtual address.
         /// </summary>
         /// <remarks>
@@ -379,12 +395,15 @@ namespace Ryujinx.Graphics.Gpu.Memory
         {
             lock (_pageTable)
             {
-                MemoryUnmapped?.Invoke(this, new UnmapEventArgs(va, size));
+                UnmapEventArgs e = new(va, size);
+                MemoryUnmapped?.Invoke(this, e);
 
                 for (ulong offset = 0; offset < size; offset += PageSize)
                 {
                     SetPte(va + offset, PackPte(pa + offset, kind));
                 }
+
+                RunRemapActions(e);
             }
         }
 
@@ -398,12 +417,15 @@ namespace Ryujinx.Graphics.Gpu.Memory
             lock (_pageTable)
             {
                 // Event handlers are not expected to be thread safe.
-                MemoryUnmapped?.Invoke(this, new UnmapEventArgs(va, size));
+                UnmapEventArgs e = new(va, size);
+                MemoryUnmapped?.Invoke(this, e);
 
                 for (ulong offset = 0; offset < size; offset += PageSize)
                 {
                     SetPte(va + offset, PteUnmapped);
                 }
+
+                RunRemapActions(e);
             }
         }
 
diff --git a/src/Ryujinx.Graphics.Gpu/Memory/UnmapEventArgs.cs b/src/Ryujinx.Graphics.Gpu/Memory/UnmapEventArgs.cs
index 305747f8..837b5aab 100644
--- a/src/Ryujinx.Graphics.Gpu/Memory/UnmapEventArgs.cs
+++ b/src/Ryujinx.Graphics.Gpu/Memory/UnmapEventArgs.cs
@@ -1,14 +1,24 @@
-namespace Ryujinx.Graphics.Gpu.Memory
+using System;
+using System.Collections.Generic;
+
+namespace Ryujinx.Graphics.Gpu.Memory
 {
     public class UnmapEventArgs
     {
         public ulong Address { get; }
         public ulong Size { get; }
+        public List<Action> RemapActions { get; private set; }
 
         public UnmapEventArgs(ulong address, ulong size)
         {
             Address = address;
             Size = size;
         }
+
+        public void AddRemapAction(Action action)
+        {
+            RemapActions ??= new List<Action>();
+            RemapActions.Add(action);
+        }
     }
 }
diff --git a/src/Ryujinx.Memory/Range/MemoryRange.cs b/src/Ryujinx.Memory/Range/MemoryRange.cs
index 7465fbcb..c7ee6db2 100644
--- a/src/Ryujinx.Memory/Range/MemoryRange.cs
+++ b/src/Ryujinx.Memory/Range/MemoryRange.cs
@@ -57,5 +57,19 @@
 
             return thisAddress < otherEndAddress && otherAddress < thisEndAddress;
         }
+
+        /// <summary>
+        /// Returns a string summary of the memory range.
+        /// </summary>
+        /// <returns>A string summary of the memory range</returns>
+        public override string ToString()
+        {
+            if (Address == ulong.MaxValue)
+            {
+                return $"[Unmapped 0x{Size:X}]";
+            }
+
+            return $"[0x{Address:X}, 0x{EndAddress:X})";
+        }
     }
 }
diff --git a/src/Ryujinx.Memory/Range/MultiRange.cs b/src/Ryujinx.Memory/Range/MultiRange.cs
index 9dbd76ec..42ef24be 100644
--- a/src/Ryujinx.Memory/Range/MultiRange.cs
+++ b/src/Ryujinx.Memory/Range/MultiRange.cs
@@ -319,5 +319,14 @@ namespace Ryujinx.Memory.Range
 
             return hash.ToHashCode();
         }
+
+        /// <summary>
+        /// Returns a string summary of the ranges contained in the MultiRange.
+        /// </summary>
+        /// <returns>A string summary of the ranges contained within</returns>
+        public override string ToString()
+        {
+            return HasSingleRange ? _singleRange.ToString() : string.Join(", ", _ranges);
+        }
     }
 }