From 642903a7bf5cd5a091ed492cfd4926b42e8d70b4 Mon Sep 17 00:00:00 2001
From: Billy Laws <blaws05@gmail.com>
Date: Sun, 23 Feb 2025 14:41:03 +0000
Subject: [PATCH] FEXCore: Support tracking TSO range information

---
 FEXCore/Source/Interface/Context/Context.h |  6 ++++++
 FEXCore/Source/Interface/Core/Core.cpp     | 13 +++++++++++++
 FEXCore/include/FEXCore/Core/Context.h     | 10 ++++++++++
 3 files changed, 29 insertions(+)

diff --git a/FEXCore/Source/Interface/Context/Context.h b/FEXCore/Source/Interface/Context/Context.h
index c3074ba25..0361961a2 100644
--- a/FEXCore/Source/Interface/Context/Context.h
+++ b/FEXCore/Source/Interface/Context/Context.h
@@ -186,6 +186,10 @@ public:
 
   void AddThunkTrampolineIRHandler(uintptr_t Entrypoint, uintptr_t GuestThunkEntrypoint) override;
 
+  void AddForceTSOInformation(const IntervalList<uint64_t>& ValidRanges, fextl::set<uint64_t>&& Instructions) override;
+
+  void RemoveForceTSOInformation(uint64_t Address, uint64_t Size) override;
+
 public:
   friend class FEXCore::HLE::SyscallHandler;
 #ifdef JIT_ARM64
@@ -373,5 +377,7 @@ private:
   std::shared_mutex CustomIRMutex;
   std::atomic<bool> HasCustomIRHandlers {};
   fextl::unordered_map<uint64_t, std::tuple<CustomIREntrypointHandler, void*, void*>> CustomIRHandlers;
+  IntervalList<uint64_t> ForceTSOValidRanges; // The ranges for which ForceTSOInstructions has populated data
+  fextl::set<uint64_t> ForceTSOInstructions;
 };
 } // namespace FEXCore::Context
diff --git a/FEXCore/Source/Interface/Core/Core.cpp b/FEXCore/Source/Interface/Core/Core.cpp
index ee8d60a7e..4e505325f 100644
--- a/FEXCore/Source/Interface/Core/Core.cpp
+++ b/FEXCore/Source/Interface/Core/Core.cpp
@@ -971,6 +971,19 @@ void ContextImpl::AddThunkTrampolineIRHandler(uintptr_t Entrypoint, uintptr_t Gu
   }
 }
 
+void ContextImpl::AddForceTSOInformation(const IntervalList<uint64_t>& ValidRanges, fextl::set<uint64_t>&& Instructions) {
+  LogMan::Throw::AFmt(CodeInvalidationMutex.try_lock() == false, "CodeInvalidationMutex needs to be unique_locked here");
+  ForceTSOValidRanges.Insert(ValidRanges);
+  ForceTSOInstructions.merge(Instructions);
+}
+
+void ContextImpl::RemoveForceTSOInformation(uint64_t Address, uint64_t Size) {
+  LogMan::Throw::AFmt(CodeInvalidationMutex.try_lock() == false, "CodeInvalidationMutex needs to be unique_locked here");
+
+  ForceTSOValidRanges.Remove({Address, Address + Size});
+  ForceTSOInstructions.erase(ForceTSOInstructions.lower_bound(Address), ForceTSOInstructions.upper_bound(Address + Size));
+}
+
 void ContextImpl::RemoveCustomIREntrypoint(uintptr_t Entrypoint) {
   LOGMAN_THROW_A_FMT(Config.Is64BitMode || !(Entrypoint >> 32), "64-bit Entrypoint in 32-bit mode {:x}", Entrypoint);
 
diff --git a/FEXCore/include/FEXCore/Core/Context.h b/FEXCore/include/FEXCore/Core/Context.h
index e80710701..27dbf2de1 100644
--- a/FEXCore/include/FEXCore/Core/Context.h
+++ b/FEXCore/include/FEXCore/Core/Context.h
@@ -7,6 +7,7 @@
 #include <FEXCore/Core/CPUID.h>
 #include <FEXCore/IR/IR.h>
 #include <FEXCore/Utils/CompilerDefs.h>
+#include <FEXCore/Utils/IntervalList.h>
 #include <FEXCore/fextl/memory.h>
 #include <FEXCore/fextl/set.h>
 #include <FEXCore/fextl/string.h>
@@ -216,6 +217,15 @@ public:
    */
   FEX_DEFAULT_VISIBILITY virtual void AddThunkTrampolineIRHandler(uintptr_t Entrypoint, uintptr_t GuestThunkEntrypoint) = 0;
 
+  /**
+   * @brief Adds additional per-instruction granularity TSO enable/disable information for the given range.
+   *
+   * @param ValidRanges The set of address ranges covered by this information
+   * @param Instructions The set of instruction addresses within the given ranges for which TSO should be enabled
+   */
+  FEX_DEFAULT_VISIBILITY virtual void AddForceTSOInformation(const IntervalList<uint64_t>& ValidRanges, fextl::set<uint64_t>&& Instructions) = 0;
+
+  FEX_DEFAULT_VISIBILITY virtual void RemoveForceTSOInformation(uint64_t Address, uint64_t Size) = 0;
 private:
 };