From 4612fb5a678d9d08aced2ceeba4832733b6a79be Mon Sep 17 00:00:00 2001
From: Ryan Houdek <Sonicadvance1@gmail.com>
Date: Thu, 2 Jul 2020 05:56:08 +0100
Subject: [PATCH] Implements support for cache clearing in the ARM64 JIT

Previously the ARM64 JIT didn't support cache clearing at all.
This adds support for clearing the JIT cache while keeping dispatcher
helpers around and the ASM dispatcher alive.
---
 .../Source/Interface/Core/JIT/Arm64/JIT.cpp   | 30 +++++++++++++++++--
 .../Interface/Core/JIT/Arm64/JITClass.h       | 11 +++++++
 2 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/External/FEXCore/Source/Interface/Core/JIT/Arm64/JIT.cpp b/External/FEXCore/Source/Interface/Core/JIT/Arm64/JIT.cpp
index a0adcb0ef..336184d52 100644
--- a/External/FEXCore/Source/Interface/Core/JIT/Arm64/JIT.cpp
+++ b/External/FEXCore/Source/Interface/Core/JIT/Arm64/JIT.cpp
@@ -61,7 +61,7 @@ void JITCore::Op_NoOp(FEXCore::IR::IROp_Header *IROp, uint32_t Node) {
 }
 
 JITCore::JITCore(FEXCore::Context::Context *ctx, FEXCore::Core::InternalThreadState *Thread)
-  : vixl::aarch64::MacroAssembler(1024 * 1024 * 128, vixl::aarch64::PositionDependentCode)
+  : vixl::aarch64::MacroAssembler(MAX_CODE_SIZE, vixl::aarch64::PositionDependentCode)
   , CTX {ctx}
   , State {Thread}
 #if _M_X86_64
@@ -97,6 +97,9 @@ JITCore::JITCore(FEXCore::Context::Context *ctx, FEXCore::Core::InternalThreadSt
   CPU.SetUp();
   SetAllowAssembler(true);
   CreateCustomDispatch(Thread);
+  GenerateDispatchHelpers();
+
+  ConstantCodeCacheOffset = GetCursorOffset();
 
   uint32_t NumUsedGPRs = NumGPRs;
   uint32_t NumUsedGPRPairs = NumGPRPairs;
@@ -138,6 +141,17 @@ JITCore::JITCore(FEXCore::Context::Context *ctx, FEXCore::Core::InternalThreadSt
   RegisterVectorHandlers();
 }
 
+void JITCore::ClearCache() {
+  // Get the backing code buffer
+  auto Buffer = GetBuffer();
+
+  // Rewind the code buffer's cursor back to just after the helpers
+  Buffer->Rewind(ConstantCodeCacheOffset);
+
+  // Make sure to set internal state as clean
+  Buffer->SetClean();
+}
+
 JITCore::~JITCore() {
 }
 
@@ -203,7 +217,7 @@ void *JITCore::CompileCode([[maybe_unused]] FEXCore::IR::IRListView<true> const
   using namespace aarch64;
   JumpTargets.clear();
   CurrentIR = IR;
-
+  uint32_t SSACount = CurrentIR->GetSSACount();
   uintptr_t ListBegin = CurrentIR->GetListData();
   uintptr_t DataBegin = CurrentIR->GetData();
 
@@ -217,6 +231,12 @@ void *JITCore::CompileCode([[maybe_unused]] FEXCore::IR::IRListView<true> const
     return State->IntBackend->CompileCode(IR, DebugData);
   }
 
+  // Fairly excessive buffer range to make sure we don't overflow
+  uint32_t BufferRange = SSACount * 16;
+  if ((GetCursorOffset() + BufferRange) > MAX_CODE_SIZE) {
+    State->CTX->ClearCodeCache(State, HeaderOp->Entry);
+  }
+
   LogMan::Throw::A(RAPass->HasFullRA(), "Arm64 JIT only works with RA");
 
   SpillSlots = RAPass->SpillSlots();
@@ -298,7 +318,6 @@ void *JITCore::CompileCode([[maybe_unused]] FEXCore::IR::IRListView<true> const
 
   FinalizeCode();
 
-
   auto CodeEnd = Buffer->GetOffsetAddress<uint64_t>(GetCursorOffset());
   CPU.EnsureIAndDCacheCoherency(reinterpret_cast<void*>(Entry), Buffer->GetOffsetAddress<uint64_t>(GetCursorOffset()) - reinterpret_cast<uint64_t>(Entry));
 #if _M_X86_64
@@ -480,6 +499,11 @@ void JITCore::CreateCustomDispatch(FEXCore::Core::InternalThreadState *Thread) {
   //CustomDispatchGenerated = true;
 }
 
+
+void JITCore::GenerateDispatchHelpers() {
+  // Nothing here yet
+}
+
 FEXCore::CPU::CPUBackend *CreateJITCore(FEXCore::Context::Context *ctx, FEXCore::Core::InternalThreadState *Thread) {
   return new JITCore(ctx, Thread);
 }
diff --git a/External/FEXCore/Source/Interface/Core/JIT/Arm64/JITClass.h b/External/FEXCore/Source/Interface/Core/JIT/Arm64/JITClass.h
index 3e6e14f42..2dc244a2d 100644
--- a/External/FEXCore/Source/Interface/Core/JIT/Arm64/JITClass.h
+++ b/External/FEXCore/Source/Interface/Core/JIT/Arm64/JITClass.h
@@ -100,6 +100,8 @@ public:
   }
 #endif
 
+  void ClearCache() override;
+
 private:
   FEXCore::Context::Context *CTX;
   FEXCore::Core::InternalThreadState *State;
@@ -128,6 +130,8 @@ private:
   constexpr static uint8_t RA_64 = 1;
   constexpr static uint8_t RA_FPR = 2;
 
+  static constexpr uint32_t MAX_CODE_SIZE = 1024 * 1024 * 128;
+
   template<uint8_t RAType>
   aarch64::Register GetReg(uint32_t Node);
 
@@ -169,6 +173,13 @@ private:
   void LoadConstant(vixl::aarch64::Register Reg, uint64_t Constant);
 
   void CreateCustomDispatch(FEXCore::Core::InternalThreadState *Thread);
+  void GenerateDispatchHelpers();
+  ptrdiff_t ConstantCodeCacheOffset{};
+  /**
+   * @name Dispatch Helper functions
+   * @{ */
+  /**  @} */
+
   bool CustomDispatchGenerated {false};
   using CustomDispatch = void(*)(FEXCore::Core::InternalThreadState *Thread);
   CustomDispatch DispatchPtr{};