From 8657d94e33825900123024101fab0b938fc13131 Mon Sep 17 00:00:00 2001
From: iwubcode <iwubcode@users.noreply.github.com>
Date: Fri, 3 Nov 2023 23:31:39 -0500
Subject: [PATCH] VideoCommon: support resetting custom asset loader

---
 Source/Core/VideoCommon/Assets/CustomAssetLoader.cpp | 10 ++++++++++
 Source/Core/VideoCommon/Assets/CustomAssetLoader.h   | 12 +++++++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/Source/Core/VideoCommon/Assets/CustomAssetLoader.cpp b/Source/Core/VideoCommon/Assets/CustomAssetLoader.cpp
index 119d7444ea..a63aa3f387 100644
--- a/Source/Core/VideoCommon/Assets/CustomAssetLoader.cpp
+++ b/Source/Core/VideoCommon/Assets/CustomAssetLoader.cpp
@@ -30,6 +30,7 @@ void CustomAssetLoader::Init()
       std::this_thread::sleep_for(TIME_BETWEEN_ASSET_MONITOR_CHECKS);
 
       std::lock_guard lk(m_asset_load_lock);
+      std::vector<CustomAssetLibrary::AssetID> dead_asset_ids;
       for (auto& [asset_id, asset_to_monitor] : m_assets_to_monitor)
       {
         if (auto ptr = asset_to_monitor.lock())
@@ -40,6 +41,15 @@ void CustomAssetLoader::Init()
             (void)ptr->Load();
           }
         }
+        else
+        {
+          dead_asset_ids.push_back(asset_id);
+        }
+      }
+
+      for (const auto& dead_asset_id : dead_asset_ids)
+      {
+        m_assets_to_monitor.erase(dead_asset_id);
       }
     }
   });
diff --git a/Source/Core/VideoCommon/Assets/CustomAssetLoader.h b/Source/Core/VideoCommon/Assets/CustomAssetLoader.h
index 90d4f81a0e..55683e169e 100644
--- a/Source/Core/VideoCommon/Assets/CustomAssetLoader.h
+++ b/Source/Core/VideoCommon/Assets/CustomAssetLoader.h
@@ -52,6 +52,17 @@ public:
   std::shared_ptr<MeshAsset> LoadMesh(const CustomAssetLibrary::AssetID& asset_id,
                                       std::shared_ptr<CustomAssetLibrary> library);
 
+  void Reset()
+  {
+    std::lock_guard lk(m_asset_load_lock);
+    m_game_textures.clear();
+    m_pixel_shaders.clear();
+    m_materials.clear();
+    m_meshes.clear();
+    m_total_bytes_loaded = 0;
+    m_assets_to_monitor.clear();
+  }
+
 private:
   // TODO C++20: use a 'derived_from' concept against 'CustomAsset' when available
   template <typename AssetType>
@@ -71,7 +82,6 @@ private:
       {
         std::lock_guard lk(m_asset_load_lock);
         m_total_bytes_loaded -= a->GetByteSizeInMemory();
-        m_assets_to_monitor.erase(a->GetAssetId());
         if (m_max_memory_available >= m_total_bytes_loaded && m_memory_exceeded)
         {
           INFO_LOG_FMT(VIDEO, "Asset memory went below limit, new assets can begin loading.");