From b6ad99ee21695955a78082fda748bd7aaaf1cf97 Mon Sep 17 00:00:00 2001
From: "Herman S." <429230+has207@users.noreply.github.com>
Date: Wed, 25 Mar 2026 14:15:59 +0900
Subject: [PATCH] [A64] Add thread-safe atomics to A64Function

Also adds some bounds check to LookupFunction.

Co-Authored-By: Reality <reality@xenios.jp>
---
 src/xenia/cpu/backend/a64/a64_function.cc |  9 +++++----
 src/xenia/cpu/backend/a64/a64_function.h  | 14 ++++++++++----
 src/xenia/cpu/backend/code_cache_base.h   | 10 +++++++++-
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/src/xenia/cpu/backend/a64/a64_function.cc b/src/xenia/cpu/backend/a64/a64_function.cc
index b27f71fae..9fbdd7294 100644
--- a/src/xenia/cpu/backend/a64/a64_function.cc
+++ b/src/xenia/cpu/backend/a64/a64_function.cc
@@ -26,18 +26,19 @@ A64Function::~A64Function() {
 }
 
 void A64Function::Setup(uint8_t* machine_code, size_t machine_code_length) {
-  machine_code_ = machine_code;
-  machine_code_length_ = machine_code_length;
+  machine_code_length_.store(machine_code_length, std::memory_order_relaxed);
+  machine_code_.store(machine_code, std::memory_order_release);
 }
 
 bool A64Function::CallImpl(ThreadState* thread_state, uint32_t return_address) {
   auto backend =
       reinterpret_cast<A64Backend*>(thread_state->processor()->backend());
   auto thunk = backend->host_to_guest_thunk();
-  if (!thunk || !machine_code_) {
+  auto* code = machine_code_.load(std::memory_order_acquire);
+  if (!thunk || !code) {
     return false;
   }
-  thunk(machine_code_, thread_state->context(),
+  thunk(code, thread_state->context(),
         reinterpret_cast<void*>(uintptr_t(return_address)));
   return true;
 }
diff --git a/src/xenia/cpu/backend/a64/a64_function.h b/src/xenia/cpu/backend/a64/a64_function.h
index 42f72dbd3..9be9ce5ba 100644
--- a/src/xenia/cpu/backend/a64/a64_function.h
+++ b/src/xenia/cpu/backend/a64/a64_function.h
@@ -10,6 +10,8 @@
 #ifndef XENIA_CPU_BACKEND_A64_A64_FUNCTION_H_
 #define XENIA_CPU_BACKEND_A64_A64_FUNCTION_H_
 
+#include <atomic>
+
 #include "xenia/cpu/function.h"
 #include "xenia/cpu/thread_state.h"
 
@@ -23,8 +25,12 @@ class A64Function : public GuestFunction {
   A64Function(Module* module, uint32_t address);
   ~A64Function() override;
 
-  uint8_t* machine_code() const override { return machine_code_; }
-  size_t machine_code_length() const override { return machine_code_length_; }
+  uint8_t* machine_code() const override {
+    return machine_code_.load(std::memory_order_acquire);
+  }
+  size_t machine_code_length() const override {
+    return machine_code_length_.load(std::memory_order_acquire);
+  }
 
   void Setup(uint8_t* machine_code, size_t machine_code_length);
 
@@ -32,8 +38,8 @@ class A64Function : public GuestFunction {
   bool CallImpl(ThreadState* thread_state, uint32_t return_address) override;
 
  private:
-  uint8_t* machine_code_ = nullptr;
-  size_t machine_code_length_ = 0;
+  std::atomic<uint8_t*> machine_code_{nullptr};
+  std::atomic<size_t> machine_code_length_{0};
 };
 
 }  // namespace a64
diff --git a/src/xenia/cpu/backend/code_cache_base.h b/src/xenia/cpu/backend/code_cache_base.h
index 0e4e9d0f2..976030c59 100644
--- a/src/xenia/cpu/backend/code_cache_base.h
+++ b/src/xenia/cpu/backend/code_cache_base.h
@@ -234,7 +234,15 @@ class CodeCacheBase : public CodeCache {
   }
 
   GuestFunction* LookupFunction(uint64_t host_pc) override {
-    uint32_t key = uint32_t(host_pc - kGeneratedCodeExecuteBase);
+    if (generated_code_map_.empty()) {
+      return nullptr;
+    }
+    const uint64_t code_base = execute_base_address();
+    const uint64_t code_end = code_base + total_size();
+    if (host_pc < code_base || host_pc >= code_end) {
+      return nullptr;
+    }
+    uint32_t key = uint32_t(host_pc - code_base);
     void* fn_entry = std::bsearch(
         &key, generated_code_map_.data(), generated_code_map_.size(),
         sizeof(std::pair<uint32_t, Function*>),