Revert "[heap_tracker] Use ankerl map instead of rb tree (#249)" (#382)

This reverts commit c9a3baab5d5ba524778492027ef8961da947df2d.

this commit caused issues in ender magnolia or something, need to make
sure I didn't mess up the revert

Reviewed-on: https://git.eden-emu.dev/eden-emu/eden/pulls/382
Reviewed-by: Shinmegumi <shinmegumi@eden-emu.dev>
Reviewed-by: Lizzie <lizzie@eden-emu.dev>
Reviewed-by: MaranBr <maranbr@outlook.com>

7 files changed, 99 insertions, 5 deletions

diff --git a/src/core/arm/dynarmic/arm_dynarmic.cpp b/src/core/arm/dynarmic/arm_dynarmic.cpp
index 9d26db51f7..e6e9fc45be 100644
--- a/src/core/arm/dynarmic/arm_dynarmic.cpp
+++ b/src/core/arm/dynarmic/arm_dynarmic.cpp
@@ -3,9 +3,47 @@
 
 #ifdef __linux__
 
-//#include "common/signal_chain.h"
+#include "common/signal_chain.h"
+
 #include "core/arm/dynarmic/arm_dynarmic.h"
-//#include "core/hle/kernel/k_process.h"
-//#include "core/memory.h"
+#include "core/hle/kernel/k_process.h"
+#include "core/memory.h"
+
+namespace Core {
+
+namespace {
+
+thread_local Core::Memory::Memory* g_current_memory{};
+std::once_flag g_registered{};
+struct sigaction g_old_segv {};
+
+void HandleSigSegv(int sig, siginfo_t* info, void* ctx) {
+    if (g_current_memory && g_current_memory->InvalidateSeparateHeap(info->si_addr)) {
+        return;
+    }
+
+    return g_old_segv.sa_sigaction(sig, info, ctx);
+}
+
+} // namespace
+
+ScopedJitExecution::ScopedJitExecution(Kernel::KProcess* process) {
+    g_current_memory = std::addressof(process->GetMemory());
+}
+
+ScopedJitExecution::~ScopedJitExecution() {
+    g_current_memory = nullptr;
+}
+
+void ScopedJitExecution::RegisterHandler() {
+    std::call_once(g_registered, [] {
+        struct sigaction sa {};
+        sa.sa_sigaction = &HandleSigSegv;
+        sa.sa_flags = SA_SIGINFO | SA_ONSTACK;
+        Common::SigAction(SIGSEGV, std::addressof(sa), std::addressof(g_old_segv));
+    });
+}
+
+} // namespace Core
 
 #endif
diff --git a/src/core/arm/dynarmic/arm_dynarmic.h b/src/core/arm/dynarmic/arm_dynarmic.h
index eef7c31160..53dd188151 100644
--- a/src/core/arm/dynarmic/arm_dynarmic.h
+++ b/src/core/arm/dynarmic/arm_dynarmic.h
@@ -26,4 +26,24 @@ constexpr HaltReason TranslateHaltReason(Dynarmic::HaltReason hr) {
     return static_cast<HaltReason>(hr);
 }
 
+#ifdef __linux__
+
+class ScopedJitExecution {
+public:
+    explicit ScopedJitExecution(Kernel::KProcess* process);
+    ~ScopedJitExecution();
+    static void RegisterHandler();
+};
+
+#else
+
+class ScopedJitExecution {
+public:
+    explicit ScopedJitExecution(Kernel::KProcess* process) {}
+    ~ScopedJitExecution() {}
+    static void RegisterHandler() {}
+};
+
+#endif
+
 } // namespace Core
diff --git a/src/core/arm/dynarmic/arm_dynarmic_32.cpp b/src/core/arm/dynarmic/arm_dynarmic_32.cpp
index 7123497682..2c2c54a1ad 100644
--- a/src/core/arm/dynarmic/arm_dynarmic_32.cpp
+++ b/src/core/arm/dynarmic/arm_dynarmic_32.cpp
@@ -343,11 +343,15 @@ bool ArmDynarmic32::IsInThumbMode() const {
 }
 
 HaltReason ArmDynarmic32::RunThread(Kernel::KThread* thread) {
+    ScopedJitExecution sj(thread->GetOwnerProcess());
+
     m_jit->ClearExclusiveState();
     return TranslateHaltReason(m_jit->Run());
 }
 
 HaltReason ArmDynarmic32::StepThread(Kernel::KThread* thread) {
+    ScopedJitExecution sj(thread->GetOwnerProcess());
+
     m_jit->ClearExclusiveState();
     return TranslateHaltReason(m_jit->Step());
 }
@@ -389,6 +393,7 @@ ArmDynarmic32::ArmDynarmic32(System& system, bool uses_wall_clock, Kernel::KProc
       m_cp15(std::make_shared<DynarmicCP15>(*this)), m_core_index{core_index} {
     auto& page_table_impl = process->GetPageTable().GetBasePageTable().GetImpl();
     m_jit = MakeJit(&page_table_impl);
+    ScopedJitExecution::RegisterHandler();
 }
 
 ArmDynarmic32::~ArmDynarmic32() = default;
diff --git a/src/core/arm/dynarmic/arm_dynarmic_64.cpp b/src/core/arm/dynarmic/arm_dynarmic_64.cpp
index 2745aeb862..438b7b691c 100644
--- a/src/core/arm/dynarmic/arm_dynarmic_64.cpp
+++ b/src/core/arm/dynarmic/arm_dynarmic_64.cpp
@@ -374,11 +374,15 @@ std::shared_ptr<Dynarmic::A64::Jit> ArmDynarmic64::MakeJit(Common::PageTable* pa
 }
 
 HaltReason ArmDynarmic64::RunThread(Kernel::KThread* thread) {
+    ScopedJitExecution sj(thread->GetOwnerProcess());
+
     m_jit->ClearExclusiveState();
     return TranslateHaltReason(m_jit->Run());
 }
 
 HaltReason ArmDynarmic64::StepThread(Kernel::KThread* thread) {
+    ScopedJitExecution sj(thread->GetOwnerProcess());
+
     m_jit->ClearExclusiveState();
     return TranslateHaltReason(m_jit->Step());
 }
@@ -418,6 +422,7 @@ ArmDynarmic64::ArmDynarmic64(System& system, bool uses_wall_clock, Kernel::KProc
     auto& page_table = process->GetPageTable().GetBasePageTable();
     auto& page_table_impl = page_table.GetImpl();
     m_jit = MakeJit(&page_table_impl, page_table.GetAddressSpaceWidth());
+    ScopedJitExecution::RegisterHandler();
 }
 
 ArmDynarmic64::~ArmDynarmic64() = default;
diff --git a/src/core/hle/kernel/k_process.cpp b/src/core/hle/kernel/k_process.cpp
index cf03353f84..80566b7e77 100644
--- a/src/core/hle/kernel/k_process.cpp
+++ b/src/core/hle/kernel/k_process.cpp
@@ -1266,6 +1266,10 @@ void KProcess::InitializeInterfaces() {
 
 #ifdef HAS_NCE
     if (this->IsApplication() && Settings::IsNceEnabled()) {
+        // Register the scoped JIT handler before creating any NCE instances
+        // so that its signal handler will appear first in the signal chain.
+        Core::ScopedJitExecution::RegisterHandler();
+
         for (size_t i = 0; i < Core::Hardware::NUM_CPU_CORES; i++) {
             m_arm_interfaces[i] = std::make_unique<Core::ArmNce>(m_kernel.System(), true, i);
         }
diff --git a/src/core/memory.cpp b/src/core/memory.cpp
index 08391cd815..0035c626e2 100644
--- a/src/core/memory.cpp
+++ b/src/core/memory.cpp
@@ -61,7 +61,8 @@ struct Memory::Impl {
         }
 
 #ifdef __linux__
-        buffer.emplace(system.DeviceMemory().buffer);
+        heap_tracker.emplace(system.DeviceMemory().buffer);
+        buffer = std::addressof(*heap_tracker);
 #else
         buffer = std::addressof(system.DeviceMemory().buffer);
 #endif
@@ -1023,8 +1024,9 @@ struct Memory::Impl {
     std::span<Core::GPUDirtyMemoryManager> gpu_dirty_managers;
     std::mutex sys_core_guard;
 
+    std::optional<Common::HeapTracker> heap_tracker;
 #ifdef __linux__
-    std::optional<Common::HeapTracker> buffer;
+    Common::HeapTracker* buffer{};
 #else
     Common::HostMemory* buffer{};
 #endif
@@ -1228,7 +1230,22 @@ bool Memory::InvalidateNCE(Common::ProcessAddress vaddr, size_t size) {
     if (rasterizer) {
         impl->InvalidateGPUMemory(ptr, size);
     }
+
+#ifdef __linux__
+    if (!rasterizer && mapped) {
+        impl->buffer->DeferredMapSeparateHeap(GetInteger(vaddr));
+    }
+#endif
+
     return mapped && ptr != nullptr;
 }
 
+bool Memory::InvalidateSeparateHeap(void* fault_address) {
+#ifdef __linux__
+    return impl->buffer->DeferredMapSeparateHeap(static_cast<u8*>(fault_address));
+#else
+    return false;
+#endif
+}
+
 } // namespace Core::Memory
diff --git a/src/core/memory.h b/src/core/memory.h
index 99108ecf0d..dcca26892b 100644
--- a/src/core/memory.h
+++ b/src/core/memory.h
@@ -487,8 +487,13 @@ public:
      *              marked as debug or non-debug.
      */
     void MarkRegionDebug(Common::ProcessAddress vaddr, u64 size, bool debug);
+
     void SetGPUDirtyManagers(std::span<Core::GPUDirtyMemoryManager> managers);
+
     bool InvalidateNCE(Common::ProcessAddress vaddr, size_t size);
+
+    bool InvalidateSeparateHeap(void* fault_address);
+
 private:
     Core::System& system;