From 98c09444198e917e5fee77eaf3366ff03075abae Mon Sep 17 00:00:00 2001
From: kd-11 <karokidii@gmail.com>
Date: Mon, 22 May 2023 20:45:16 +0300
Subject: [PATCH] vk: Fix some sampler load crashes

---
 rpcs3/Emu/RSX/VK/VKResourceManager.h |  6 +++---
 rpcs3/Emu/RSX/VK/vkutils/sampler.cpp | 11 ++++++-----
 rpcs3/Emu/RSX/VK/vkutils/sampler.h   |  4 ++--
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/rpcs3/Emu/RSX/VK/VKResourceManager.h b/rpcs3/Emu/RSX/VK/VKResourceManager.h
index ac2e579e74..bc9b459482 100644
--- a/rpcs3/Emu/RSX/VK/VKResourceManager.h
+++ b/rpcs3/Emu/RSX/VK/VKResourceManager.h
@@ -136,6 +136,7 @@ namespace vk
 
 			if (const auto found = m_sampler_pool.find(key))
 			{
+				found->add_ref();
 				return found;
 			}
 
@@ -145,9 +146,8 @@ namespace vk
 				min_filter, mag_filter, mipmap_mode, border_color,
 				depth_compare, depth_compare_mode);
 
-			result->add_ref();
-			auto ret = result.get();
-			m_sampler_pool.emplace(key, result);
+			auto ret = m_sampler_pool.emplace(key, result);
+			ret->add_ref();
 			return ret;
 		}
 
diff --git a/rpcs3/Emu/RSX/VK/vkutils/sampler.cpp b/rpcs3/Emu/RSX/VK/vkutils/sampler.cpp
index f5e957329f..586656c621 100644
--- a/rpcs3/Emu/RSX/VK/vkutils/sampler.cpp
+++ b/rpcs3/Emu/RSX/VK/vkutils/sampler.cpp
@@ -154,7 +154,7 @@ namespace vk
 		m_custom_color_sampler_pool.clear();
 	}
 
-	vk::sampler* sampler_pool_t::find(const sampler_pool_key_t& key) const
+	cached_sampler_object_t* sampler_pool_t::find(const sampler_pool_key_t& key) const
 	{
 		if (!key.border_color_key) [[ likely ]]
 		{
@@ -174,17 +174,18 @@ namespace vk
 		return nullptr;
 	}
 
-	void sampler_pool_t::emplace(const sampler_pool_key_t& key, std::unique_ptr<cached_sampler_object_t>& object)
+	cached_sampler_object_t* sampler_pool_t::emplace(const sampler_pool_key_t& key, std::unique_ptr<cached_sampler_object_t>& object)
 	{
 		object->key = key;
 
 		if (!key.border_color_key) [[ likely ]]
 		{
-			m_generic_sampler_pool.emplace(key.base_key, std::move(object));
-			return;
+			const auto [iterator, _unused] = m_generic_sampler_pool.emplace(key.base_key, std::move(object));
+			return iterator->second.get();
 		}
 
-		m_custom_color_sampler_pool.emplace (key.base_key, std::move(object));
+		const auto [iterator, _unused] = m_custom_color_sampler_pool.emplace(key.base_key, std::move(object));
+		return iterator->second.get();
 	}
 
 	std::vector<std::unique_ptr<cached_sampler_object_t>> sampler_pool_t::collect(std::function<bool(const cached_sampler_object_t&)> predicate)
diff --git a/rpcs3/Emu/RSX/VK/vkutils/sampler.h b/rpcs3/Emu/RSX/VK/vkutils/sampler.h
index f231350b43..ff86ffe9b4 100644
--- a/rpcs3/Emu/RSX/VK/vkutils/sampler.h
+++ b/rpcs3/Emu/RSX/VK/vkutils/sampler.h
@@ -80,9 +80,9 @@ namespace vk
 
 		void clear();
 
-		vk::sampler* find(const sampler_pool_key_t& key) const;
+		cached_sampler_object_t* find(const sampler_pool_key_t& key) const;
 
-		void emplace(const sampler_pool_key_t& key, std::unique_ptr<cached_sampler_object_t>& object);
+		cached_sampler_object_t* emplace(const sampler_pool_key_t& key, std::unique_ptr<cached_sampler_object_t>& object);
 
 		std::vector<std::unique_ptr<cached_sampler_object_t>> collect(std::function<bool(const cached_sampler_object_t&)> predicate);
 	};