12 files changed, 56 insertions, 27 deletions

diff --git a/src/core/arm/arm_interface.h b/src/core/arm/arm_interface.h
index ccd43f431..2aa017a54 100644
--- a/src/core/arm/arm_interface.h
+++ b/src/core/arm/arm_interface.h
@@ -41,6 +41,9 @@ public:
     /// Clear all instruction cache
     virtual void ClearInstructionCache() = 0;
 
+    /// Notify CPU emulation that page tables have changed
+    virtual void PageTableChanged() = 0;
+
     /**
      * Set the Program Counter to an address
      * @param addr Address to set PC to
diff --git a/src/core/arm/dynarmic/arm_dynarmic.cpp b/src/core/arm/dynarmic/arm_dynarmic.cpp
index 34c5aa381..42ae93ae8 100644
--- a/src/core/arm/dynarmic/arm_dynarmic.cpp
+++ b/src/core/arm/dynarmic/arm_dynarmic.cpp
@@ -41,7 +41,7 @@ static bool IsReadOnlyMemory(u32 vaddr) {
 }
 
 static Dynarmic::UserCallbacks GetUserCallbacks(
-    const std::shared_ptr<ARMul_State>& interpeter_state) {
+    const std::shared_ptr<ARMul_State>& interpeter_state, Memory::PageTable* current_page_table) {
     Dynarmic::UserCallbacks user_callbacks{};
     user_callbacks.InterpreterFallback = &InterpreterFallback;
     user_callbacks.user_arg = static_cast<void*>(interpeter_state.get());
@@ -56,16 +56,14 @@ static Dynarmic::UserCallbacks GetUserCallbacks(
     user_callbacks.memory.Write16 = &Memory::Write16;
     user_callbacks.memory.Write32 = &Memory::Write32;
     user_callbacks.memory.Write64 = &Memory::Write64;
-    // TODO(Subv): Re-add the page table pointers once dynarmic supports switching page tables at
-    // runtime.
-    user_callbacks.page_table = nullptr;
+    user_callbacks.page_table = &current_page_table->pointers;
     user_callbacks.coprocessors[15] = std::make_shared<DynarmicCP15>(interpeter_state);
     return user_callbacks;
 }
 
 ARM_Dynarmic::ARM_Dynarmic(PrivilegeMode initial_mode) {
     interpreter_state = std::make_shared<ARMul_State>(initial_mode);
-    jit = std::make_unique<Dynarmic::Jit>(GetUserCallbacks(interpreter_state));
+    PageTableChanged();
 }
 
 void ARM_Dynarmic::SetPC(u32 pc) {
@@ -136,6 +134,7 @@ void ARM_Dynarmic::AddTicks(u64 ticks) {
 MICROPROFILE_DEFINE(ARM_Jit, "ARM JIT", "ARM JIT", MP_RGB(255, 64, 64));
 
 void ARM_Dynarmic::ExecuteInstructions(int num_instructions) {
+    ASSERT(Memory::GetCurrentPageTable() == current_page_table);
     MICROPROFILE_SCOPE(ARM_Jit);
 
     std::size_t ticks_executed = jit->Run(static_cast<unsigned>(num_instructions));
@@ -178,3 +177,16 @@ void ARM_Dynarmic::PrepareReschedule() {
 void ARM_Dynarmic::ClearInstructionCache() {
     jit->ClearCache();
 }
+
+void ARM_Dynarmic::PageTableChanged() {
+    current_page_table = Memory::GetCurrentPageTable();
+
+    auto iter = jits.find(current_page_table);
+    if (iter != jits.end()) {
+        jit = iter->second.get();
+        return;
+    }
+
+    jit = new Dynarmic::Jit(GetUserCallbacks(interpreter_state, current_page_table));
+    jits.emplace(current_page_table, std::unique_ptr<Dynarmic::Jit>(jit));
+}
diff --git a/src/core/arm/dynarmic/arm_dynarmic.h b/src/core/arm/dynarmic/arm_dynarmic.h
index 834dc989e..96148a1a5 100644
--- a/src/core/arm/dynarmic/arm_dynarmic.h
+++ b/src/core/arm/dynarmic/arm_dynarmic.h
@@ -4,12 +4,17 @@
 
 #pragma once
 
+#include <map>
 #include <memory>
 #include <dynarmic/dynarmic.h>
 #include "common/common_types.h"
 #include "core/arm/arm_interface.h"
 #include "core/arm/skyeye_common/armstate.h"
 
+namespace Memory {
+struct PageTable;
+} // namespace Memory
+
 class ARM_Dynarmic final : public ARM_Interface {
 public:
     ARM_Dynarmic(PrivilegeMode initial_mode);
@@ -36,8 +41,11 @@ public:
     void ExecuteInstructions(int num_instructions) override;
 
     void ClearInstructionCache() override;
+    void PageTableChanged() override;
 
 private:
-    std::unique_ptr<Dynarmic::Jit> jit;
+    Dynarmic::Jit* jit = nullptr;
+    Memory::PageTable* current_page_table = nullptr;
+    std::map<Memory::PageTable*, std::unique_ptr<Dynarmic::Jit>> jits;
     std::shared_ptr<ARMul_State> interpreter_state;
 };
diff --git a/src/core/arm/dyncom/arm_dyncom.cpp b/src/core/arm/dyncom/arm_dyncom.cpp
index 81f9bf99e..da955c9b9 100644
--- a/src/core/arm/dyncom/arm_dyncom.cpp
+++ b/src/core/arm/dyncom/arm_dyncom.cpp
@@ -25,6 +25,10 @@ void ARM_DynCom::ClearInstructionCache() {
     trans_cache_buf_top = 0;
 }
 
+void ARM_DynCom::PageTableChanged() {
+    ClearInstructionCache();
+}
+
 void ARM_DynCom::SetPC(u32 pc) {
     state->Reg[15] = pc;
 }
diff --git a/src/core/arm/dyncom/arm_dyncom.h b/src/core/arm/dyncom/arm_dyncom.h
index 62c174f3c..0ae535671 100644
--- a/src/core/arm/dyncom/arm_dyncom.h
+++ b/src/core/arm/dyncom/arm_dyncom.h
@@ -16,6 +16,7 @@ public:
     ~ARM_DynCom();
 
     void ClearInstructionCache() override;
+    void PageTableChanged() override;
 
     void SetPC(u32 pc) override;
     u32 GetPC() const override;
diff --git a/src/core/hle/kernel/thread.cpp b/src/core/hle/kernel/thread.cpp
index 324415a36..61378211f 100644
--- a/src/core/hle/kernel/thread.cpp
+++ b/src/core/hle/kernel/thread.cpp
@@ -178,16 +178,13 @@ static void SwitchContext(Thread* new_thread) {
         ready_queue.remove(new_thread->current_priority, new_thread);
         new_thread->status = THREADSTATUS_RUNNING;
 
-        Core::CPU().LoadContext(new_thread->context);
-        Core::CPU().SetCP15Register(CP15_THREAD_URO, new_thread->GetTLSAddress());
-
         if (previous_process != current_thread->owner_process) {
             Kernel::g_current_process = current_thread->owner_process;
-            Memory::current_page_table = &Kernel::g_current_process->vm_manager.page_table;
-            // We have switched processes and thus, page tables, clear the instruction cache so we
-            // don't keep stale data from the previous process.
-            Core::CPU().ClearInstructionCache();
+            SetCurrentPageTable(&Kernel::g_current_process->vm_manager.page_table);
         }
+
+        Core::CPU().LoadContext(new_thread->context);
+        Core::CPU().SetCP15Register(CP15_THREAD_URO, new_thread->GetTLSAddress());
     } else {
         current_thread = nullptr;
         // Note: We do not reset the current process and current page table when idling because
diff --git a/src/core/loader/3dsx.cpp b/src/core/loader/3dsx.cpp
index 69cdc0867..a03515e6e 100644
--- a/src/core/loader/3dsx.cpp
+++ b/src/core/loader/3dsx.cpp
@@ -270,7 +270,7 @@ ResultStatus AppLoader_THREEDSX::Load() {
     Kernel::g_current_process = Kernel::Process::Create(std::move(codeset));
     Kernel::g_current_process->svc_access_mask.set();
     Kernel::g_current_process->address_mappings = default_address_mappings;
-    Memory::current_page_table = &Kernel::g_current_process->vm_manager.page_table;
+    Memory::SetCurrentPageTable(&Kernel::g_current_process->vm_manager.page_table);
 
     // Attach the default resource limit (APPLICATION) to the process
     Kernel::g_current_process->resource_limit =
diff --git a/src/core/loader/elf.cpp b/src/core/loader/elf.cpp
index 2f27606a1..2de1f4e81 100644
--- a/src/core/loader/elf.cpp
+++ b/src/core/loader/elf.cpp
@@ -397,7 +397,7 @@ ResultStatus AppLoader_ELF::Load() {
     Kernel::g_current_process = Kernel::Process::Create(std::move(codeset));
     Kernel::g_current_process->svc_access_mask.set();
     Kernel::g_current_process->address_mappings = default_address_mappings;
-    Memory::current_page_table = &Kernel::g_current_process->vm_manager.page_table;
+    Memory::SetCurrentPageTable(&Kernel::g_current_process->vm_manager.page_table);
 
     // Attach the default resource limit (APPLICATION) to the process
     Kernel::g_current_process->resource_limit =
diff --git a/src/core/loader/ncch.cpp b/src/core/loader/ncch.cpp
index bef7fa567..c46d7cfc6 100644
--- a/src/core/loader/ncch.cpp
+++ b/src/core/loader/ncch.cpp
@@ -108,7 +108,7 @@ ResultStatus AppLoader_NCCH::LoadExec() {
         codeset->memory = std::make_shared<std::vector<u8>>(std::move(code));
 
         Kernel::g_current_process = Kernel::Process::Create(std::move(codeset));
-        Memory::current_page_table = &Kernel::g_current_process->vm_manager.page_table;
+        Memory::SetCurrentPageTable(&Kernel::g_current_process->vm_manager.page_table);
 
         // Attach a resource limit to the process based on the resource limit category
         Kernel::g_current_process->resource_limit =
diff --git a/src/core/memory.cpp b/src/core/memory.cpp
index 68a6b1ac2..a6b5f6c99 100644
--- a/src/core/memory.cpp
+++ b/src/core/memory.cpp
@@ -9,6 +9,8 @@
 #include "common/common_types.h"
 #include "common/logging/log.h"
 #include "common/swap.h"
+#include "core/arm/arm_interface.h"
+#include "core/core.h"
 #include "core/hle/kernel/memory.h"
 #include "core/hle/kernel/process.h"
 #include "core/hle/lock.h"
@@ -22,10 +24,17 @@ namespace Memory {
 static std::array<u8, Memory::VRAM_SIZE> vram;
 static std::array<u8, Memory::N3DS_EXTRA_RAM_SIZE> n3ds_extra_ram;
 
-PageTable* current_page_table = nullptr;
+static PageTable* current_page_table = nullptr;
 
-std::array<u8*, PAGE_TABLE_NUM_ENTRIES>* GetCurrentPageTablePointers() {
-    return &current_page_table->pointers;
+void SetCurrentPageTable(PageTable* page_table) {
+    current_page_table = page_table;
+    if (Core::System::GetInstance().IsPoweredOn()) {
+        Core::CPU().PageTableChanged();
+    }
+}
+
+PageTable* GetCurrentPageTable() {
+    return current_page_table;
 }
 
 static void MapPages(PageTable& page_table, u32 base, u32 size, u8* memory, PageType type) {
diff --git a/src/core/memory.h b/src/core/memory.h
index b228a48c2..1865bfea0 100644
--- a/src/core/memory.h
+++ b/src/core/memory.h
@@ -182,7 +182,8 @@ enum : VAddr {
 };
 
 /// Currently active page table
-extern PageTable* current_page_table;
+void SetCurrentPageTable(PageTable* page_table);
+PageTable* GetCurrentPageTable();
 
 bool IsValidVirtualAddress(const VAddr addr);
 bool IsValidPhysicalAddress(const PAddr addr);
@@ -259,10 +260,4 @@ enum class FlushMode {
  */
 void RasterizerFlushVirtualRegion(VAddr start, u32 size, FlushMode mode);
 
-/**
- * Dynarmic has an optimization to memory accesses when the pointer to the page exists that
- * can be used by setting up the current page table as a callback. This function is used to
- * retrieve the current page table for that purpose.
- */
-std::array<u8*, PAGE_TABLE_NUM_ENTRIES>* GetCurrentPageTablePointers();
 } // namespace Memory
diff --git a/src/tests/core/arm/arm_test_common.cpp b/src/tests/core/arm/arm_test_common.cpp
index 8384ce744..cfe0d503a 100644
--- a/src/tests/core/arm/arm_test_common.cpp
+++ b/src/tests/core/arm/arm_test_common.cpp
@@ -21,7 +21,7 @@ TestEnvironment::TestEnvironment(bool mutable_memory_)
     Memory::MapIoRegion(page_table, 0x00000000, 0x80000000, test_memory);
     Memory::MapIoRegion(page_table, 0x80000000, 0x80000000, test_memory);
 
-    Memory::current_page_table = &page_table;
+    Memory::SetCurrentPageTable(&page_table);
 }
 
 TestEnvironment::~TestEnvironment() {