4 files changed, 34 insertions, 15 deletions

diff --git a/src/core/hle/kernel/k_page_table.cpp b/src/core/hle/kernel/k_page_table.cpp
index 599013cf6..47ea3c89c 100644
--- a/src/core/hle/kernel/k_page_table.cpp
+++ b/src/core/hle/kernel/k_page_table.cpp
@@ -346,7 +346,8 @@ ResultCode KPageTable::MapCodeMemory(VAddr dst_address, VAddr src_address, std::
     return ResultSuccess;
 }
 
-ResultCode KPageTable::UnmapCodeMemory(VAddr dst_address, VAddr src_address, std::size_t size) {
+ResultCode KPageTable::UnmapCodeMemory(VAddr dst_address, VAddr src_address, std::size_t size,
+                                       ICacheInvalidationStrategy icache_invalidation_strategy) {
     // Validate the mapping request.
     R_UNLESS(this->CanContain(dst_address, size, KMemoryState::AliasCode),
              ResultInvalidMemoryRegion);
@@ -396,7 +397,11 @@ ResultCode KPageTable::UnmapCodeMemory(VAddr dst_address, VAddr src_address, std
     bool reprotected_pages = false;
     SCOPE_EXIT({
         if (reprotected_pages && any_code_pages) {
-            system.InvalidateCpuInstructionCacheRange(dst_address, size);
+            if (icache_invalidation_strategy == ICacheInvalidationStrategy::InvalidateRange) {
+                system.InvalidateCpuInstructionCacheRange(dst_address, size);
+            } else {
+                system.InvalidateCpuInstructionCaches();
+            }
         }
     });
 
@@ -563,6 +568,8 @@ ResultCode KPageTable::UnmapProcessMemory(VAddr dst_addr, std::size_t size,
     block_manager->Update(dst_addr, num_pages, KMemoryState::Free, KMemoryPermission::None,
                           KMemoryAttribute::None);
 
+    system.InvalidateCpuInstructionCaches();
+
     return ResultSuccess;
 }
 
diff --git a/src/core/hle/kernel/k_page_table.h b/src/core/hle/kernel/k_page_table.h
index bfabdf38c..dd6022975 100644
--- a/src/core/hle/kernel/k_page_table.h
+++ b/src/core/hle/kernel/k_page_table.h
@@ -26,6 +26,8 @@ class KMemoryBlockManager;
 
 class KPageTable final {
 public:
+    enum class ICacheInvalidationStrategy : u32 { InvalidateRange, InvalidateAll };
+
     YUZU_NON_COPYABLE(KPageTable);
     YUZU_NON_MOVEABLE(KPageTable);
 
@@ -38,7 +40,8 @@ public:
     ResultCode MapProcessCode(VAddr addr, std::size_t pages_count, KMemoryState state,
                               KMemoryPermission perm);
     ResultCode MapCodeMemory(VAddr dst_address, VAddr src_address, std::size_t size);
-    ResultCode UnmapCodeMemory(VAddr dst_address, VAddr src_address, std::size_t size);
+    ResultCode UnmapCodeMemory(VAddr dst_address, VAddr src_address, std::size_t size,
+                               ICacheInvalidationStrategy icache_invalidation_strategy);
     ResultCode UnmapProcessMemory(VAddr dst_addr, std::size_t size, KPageTable& src_page_table,
                                   VAddr src_addr);
     ResultCode MapPhysicalMemory(VAddr addr, std::size_t size);
diff --git a/src/core/hle/kernel/svc.cpp b/src/core/hle/kernel/svc.cpp
index 976d63234..0c86435b5 100644
--- a/src/core/hle/kernel/svc.cpp
+++ b/src/core/hle/kernel/svc.cpp
@@ -1713,7 +1713,8 @@ static ResultCode UnmapProcessCodeMemory(Core::System& system, Handle process_ha
         return ResultInvalidMemoryRegion;
     }
 
-    return page_table.UnmapCodeMemory(dst_address, src_address, size);
+    return page_table.UnmapCodeMemory(dst_address, src_address, size,
+                                      KPageTable::ICacheInvalidationStrategy::InvalidateAll);
 }
 
 /// Exits the current process
diff --git a/src/core/hle/service/ldr/ldr.cpp b/src/core/hle/service/ldr/ldr.cpp
index 2477c5612..cf727c167 100644
--- a/src/core/hle/service/ldr/ldr.cpp
+++ b/src/core/hle/service/ldr/ldr.cpp
@@ -389,8 +389,12 @@ public:
 
             if (bss_size) {
                 auto block_guard = detail::ScopeExit([&] {
-                    page_table.UnmapCodeMemory(addr + nro_size, bss_addr, bss_size);
-                    page_table.UnmapCodeMemory(addr, nro_addr, nro_size);
+                    page_table.UnmapCodeMemory(
+                        addr + nro_size, bss_addr, bss_size,
+                        Kernel::KPageTable::ICacheInvalidationStrategy::InvalidateRange);
+                    page_table.UnmapCodeMemory(
+                        addr, nro_addr, nro_size,
+                        Kernel::KPageTable::ICacheInvalidationStrategy::InvalidateRange);
                 });
 
                 const ResultCode result{
@@ -570,17 +574,21 @@ public:
         auto& page_table{system.CurrentProcess()->PageTable()};
 
         if (info.bss_size != 0) {
-            CASCADE_CODE(page_table.UnmapCodeMemory(info.nro_address + info.text_size +
-                                                        info.ro_size + info.data_size,
-                                                    info.bss_address, info.bss_size));
+            CASCADE_CODE(page_table.UnmapCodeMemory(
+                info.nro_address + info.text_size + info.ro_size + info.data_size, info.bss_address,
+                info.bss_size, Kernel::KPageTable::ICacheInvalidationStrategy::InvalidateRange));
         }
 
-        CASCADE_CODE(page_table.UnmapCodeMemory(info.nro_address + info.text_size + info.ro_size,
-                                                info.src_addr + info.text_size + info.ro_size,
-                                                info.data_size));
-        CASCADE_CODE(page_table.UnmapCodeMemory(info.nro_address + info.text_size,
-                                                info.src_addr + info.text_size, info.ro_size));
-        CASCADE_CODE(page_table.UnmapCodeMemory(info.nro_address, info.src_addr, info.text_size));
+        CASCADE_CODE(page_table.UnmapCodeMemory(
+            info.nro_address + info.text_size + info.ro_size,
+            info.src_addr + info.text_size + info.ro_size, info.data_size,
+            Kernel::KPageTable::ICacheInvalidationStrategy::InvalidateRange));
+        CASCADE_CODE(page_table.UnmapCodeMemory(
+            info.nro_address + info.text_size, info.src_addr + info.text_size, info.ro_size,
+            Kernel::KPageTable::ICacheInvalidationStrategy::InvalidateRange));
+        CASCADE_CODE(page_table.UnmapCodeMemory(
+            info.nro_address, info.src_addr, info.text_size,
+            Kernel::KPageTable::ICacheInvalidationStrategy::InvalidateRange));
         return ResultSuccess;
     }