From 90cf47cdce4e8b4641ef25071915aca1e0668902 Mon Sep 17 00:00:00 2001
From: kd-11 <karokidii@gmail.com>
Date: Fri, 16 Dec 2022 00:08:30 +0300
Subject: [PATCH] rsx: Handle some corner cases in surface locking

---
 rpcs3/Emu/RSX/Common/surface_utils.h       | 24 ++++++++++++++++++++++
 rpcs3/Emu/RSX/Common/texture_cache_utils.h |  4 ++--
 rpcs3/Emu/RSX/GL/GLTextureCache.h          |  4 ++--
 rpcs3/Emu/RSX/VK/VKTextureCache.h          |  4 ++--
 4 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/rpcs3/Emu/RSX/Common/surface_utils.h b/rpcs3/Emu/RSX/Common/surface_utils.h
index 5ff9490867..0ffa1e9009 100644
--- a/rpcs3/Emu/RSX/Common/surface_utils.h
+++ b/rpcs3/Emu/RSX/Common/surface_utils.h
@@ -722,6 +722,30 @@ namespace rsx
 			release();
 		}
 
+		void on_swap_out()
+		{
+			if (is_locked())
+			{
+				on_unlock();
+			}
+			else
+			{
+				release();
+			}
+		}
+
+		void on_swap_in(bool lock)
+		{
+			if (!is_locked() && lock)
+			{
+				on_lock();
+			}
+			else
+			{
+				add_ref();
+			}
+		}
+
 		bool is_locked() const
 		{
 			return texture_cache_metadata.locked;
diff --git a/rpcs3/Emu/RSX/Common/texture_cache_utils.h b/rpcs3/Emu/RSX/Common/texture_cache_utils.h
index 002df8f8ac..7c3704312c 100644
--- a/rpcs3/Emu/RSX/Common/texture_cache_utils.h
+++ b/rpcs3/Emu/RSX/Common/texture_cache_utils.h
@@ -1375,17 +1375,17 @@ namespace rsx
 			if (context == rsx::texture_upload_context::framebuffer_storage && !Emu.IsStopped())
 			{
 				// Lock, unlock
+				auto surface = derived()->get_render_target();
+
 				if (prot == utils::protection::no && old_prot != utils::protection::no)
 				{
 					// Locked memory. We have to take ownership of the object in the surface cache as well
-					auto surface = derived()->get_render_target();
 					surface->on_lock();
 				}
 				else if (old_prot == utils::protection::no && prot != utils::protection::no)
 				{
 					// Release the surface, the cache can remove it if needed
 					ensure(prot == utils::protection::rw);
-					auto surface = derived()->get_render_target();
 					surface->on_unlock();
 				}
 			}
diff --git a/rpcs3/Emu/RSX/GL/GLTextureCache.h b/rpcs3/Emu/RSX/GL/GLTextureCache.h
index fb0ba38893..62ecfc2e1e 100644
--- a/rpcs3/Emu/RSX/GL/GLTextureCache.h
+++ b/rpcs3/Emu/RSX/GL/GLTextureCache.h
@@ -70,12 +70,12 @@ namespace gl
 			if (vram_texture && !managed_texture && get_protection() == utils::protection::no)
 			{
 				// In-place image swap, still locked. Likely a color buffer that got rebound as depth buffer or vice-versa.
-				gl::as_rtt(vram_texture)->release();
+				gl::as_rtt(vram_texture)->on_swap_out();
 
 				if (!managed)
 				{
 					// Incoming is also an external resource, reference it immediately
-					gl::as_rtt(image)->add_ref();
+					gl::as_rtt(image)->on_swap_in(is_locked());
 				}
 			}
 
diff --git a/rpcs3/Emu/RSX/VK/VKTextureCache.h b/rpcs3/Emu/RSX/VK/VKTextureCache.h
index e57b2da0c1..5d16b47c45 100644
--- a/rpcs3/Emu/RSX/VK/VKTextureCache.h
+++ b/rpcs3/Emu/RSX/VK/VKTextureCache.h
@@ -50,12 +50,12 @@ namespace vk
 			if (vram_texture && !managed_texture && get_protection() == utils::protection::no)
 			{
 				// In-place image swap, still locked. Likely a color buffer that got rebound as depth buffer or vice-versa.
-				vk::as_rtt(vram_texture)->release();
+				vk::as_rtt(vram_texture)->on_swap_out();
 
 				if (!managed)
 				{
 					// Incoming is also an external resource, reference it immediately
-					vk::as_rtt(image)->add_ref();
+					vk::as_rtt(image)->on_swap_in(is_locked());
 				}
 			}