From 23ddc0b9f79130fd002f3b72c7632274e4e77d22 Mon Sep 17 00:00:00 2001
From: Simone <91993281+SimoneN64@users.noreply.github.com>
Date: Tue, 14 Jan 2025 17:21:04 +0100
Subject: [PATCH] [JIT]: Specialize register write handlers

---
 .vscode/settings.json                    |  5 ++
 src/backend/core/JIT.cpp                 |  2 +-
 src/backend/core/JIT.hpp                 | 12 ++++
 src/backend/core/jit/helpers.hpp         |  1 -
 src/backend/core/jit/instructions.cpp    |  9 ++-
 src/backend/core/registers/Registers.cpp | 78 +++++++++++++++++++++++-
 src/backend/core/registers/Registers.hpp |  9 ++-
 7 files changed, 110 insertions(+), 6 deletions(-)
 create mode 100644 .vscode/settings.json

diff --git a/.vscode/settings.json b/.vscode/settings.json
new file mode 100644
index 00000000..86e7c8c9
--- /dev/null
+++ b/.vscode/settings.json
@@ -0,0 +1,5 @@
+{
+    "githubPullRequests.ignoredPullRequestBranches": [
+        "dev"
+    ]
+}
\ No newline at end of file
diff --git a/src/backend/core/JIT.cpp b/src/backend/core/JIT.cpp
index 8ac4f76c..fcf549f9 100644
--- a/src/backend/core/JIT.cpp
+++ b/src/backend/core/JIT.cpp
@@ -2,7 +2,7 @@
 #include <jit/helpers.hpp>
 
 namespace n64 {
-JIT::JIT(ParallelRDP &parallel) : mem(regs, parallel) {}
+JIT::JIT(ParallelRDP &parallel) : regs(this), mem(regs, parallel) {}
 
 bool JIT::ShouldServiceInterrupt() const {
   const bool interrupts_pending = (regs.cop0.status.im & regs.cop0.cause.interruptPending) != 0;
diff --git a/src/backend/core/JIT.hpp b/src/backend/core/JIT.hpp
index 90f0a8c2..ac252d40 100644
--- a/src/backend/core/JIT.hpp
+++ b/src/backend/core/JIT.hpp
@@ -40,6 +40,18 @@ struct JIT : BaseCPU {
   u64 cop2Latch{};
   friend struct Cop1;
 
+  template <typename T>
+  AddressFrame GPR(size_t index) {
+    if constexpr(sizeof(T) == 1) {
+      return code.byte[offsetof(JIT, regs) + offsetof(Registers, gpr) + 8 * x];
+    } else if constexpr(sizeof(T) == 2) {
+      return code.word[offsetof(JIT, regs) + offsetof(Registers, gpr) + 8 * x];
+    } else if constexpr(sizeof(T) == 4) {
+      return code.dword[offsetof(JIT, regs) + offsetof(Registers, gpr) + 8 * x];
+    } else if constexpr(sizeof(T) == 8) {
+      return code.qword[offsetof(JIT, regs) + offsetof(Registers, gpr) + 8 * x];
+    }
+  }
 
   // Credits to PCSX-Redux: https://github.com/grumpycoders/pcsx-redux
   // Sets dest to "pointer"
diff --git a/src/backend/core/jit/helpers.hpp b/src/backend/core/jit/helpers.hpp
index b8f4d661..10bee975 100644
--- a/src/backend/core/jit/helpers.hpp
+++ b/src/backend/core/jit/helpers.hpp
@@ -42,6 +42,5 @@ static bool InstrEndsBlock(const u32 instr) {
 }
 
 #define REG(acc, x) code.acc[offsetof(JIT, regs) + offsetof(Registers, x)]
-#define GPR(x) code.qword[offsetof(JIT, regs) + offsetof(Registers, gpr) + 8 * x]
 #define GPR_constant_marker(x) code.byte[offsetof(JIT, regs) + offsetof(Registers, gprIsConstant) + x]
 } // namespace n64
diff --git a/src/backend/core/jit/instructions.cpp b/src/backend/core/jit/instructions.cpp
index 97acc7d7..5849c017 100644
--- a/src/backend/core/jit/instructions.cpp
+++ b/src/backend/core/jit/instructions.cpp
@@ -76,8 +76,13 @@ void JIT::addu(u32 instr) {
     code.mov(code.eax, result);
     code.movsxd(code.rax, code.eax);
     code.mov(GPR(RD(instr)), code.rax);
-  } else {
-    Util::panic("[JIT]: Implement non constant ADDI");
+    code.mov(REG(byte, gprIsConstant), 1);
+    return;
+  }
+  
+  if (regs.IsRegConstant(RS(instr))) {
+    const s32 rs = regs.Read<s32>(RS(instr));
+    return;
   }
 }
 
diff --git a/src/backend/core/registers/Registers.cpp b/src/backend/core/registers/Registers.cpp
index cc7fcc85..40497ec0 100644
--- a/src/backend/core/registers/Registers.cpp
+++ b/src/backend/core/registers/Registers.cpp
@@ -1,7 +1,8 @@
 #include <core/registers/Registers.hpp>
+#include <core/JIT.hpp>
 
 namespace n64 {
-Registers::Registers() : cop0(*this), cop1(*this) { Reset(); }
+Registers::Registers(JIT* jit) : jit(jit), cop0(*this), cop1(*this) { Reset(); }
 
 void Registers::Reset() {
   hi = 0;
@@ -74,6 +75,15 @@ template <>
 void Registers::Write<bool>(size_t idx, bool v) {
   if (idx == 0)
     return;
+
+  if(jit) {
+    jit->code.mov(jit->code.rax, v);
+    jit->code.mov(GPR(idx), jit->code.rax);
+    jit->code.mov(jit->code.rax, 1);
+    jit->code.mov(GPR_constant_marker(idx), jit->code.rax);
+    return;
+  }
+
   gpr[idx] = v;
   gprIsConstant[idx] = true;
 }
@@ -82,6 +92,15 @@ template <>
 void Registers::Write<u64>(size_t idx, u64 v) {
   if (idx == 0)
     return;
+
+  if(jit) {
+    jit->code.mov(jit->code.rax, v);
+    jit->code.mov(GPR(idx), jit->code.rax);
+    jit->code.mov(jit->code.rax, 1);
+    jit->code.mov(GPR_constant_marker(idx), jit->code.rax);
+    return;
+  }
+
   gpr[idx] = v;
   gprIsConstant[idx] = true;
 }
@@ -95,6 +114,15 @@ template <>
 void Registers::Write<u32>(size_t idx, u32 v) {
   if (idx == 0)
     return;
+
+  if(jit) {
+    jit->code.mov(jit->code.rax, v);
+    jit->code.mov(GPR(idx), jit->code.rax);
+    jit->code.mov(jit->code.rax, 1);
+    jit->code.mov(GPR_constant_marker(idx), jit->code.rax);
+    return;
+  }
+
   gpr[idx] = (u32)v;
   gprIsConstant[idx] = true;
 }
@@ -103,6 +131,16 @@ template <>
 void Registers::Write<s32>(size_t idx, s32 v) {
   if (idx == 0)
     return;
+
+  if(jit) {
+    jit->code.mov(jit->code.eax, v);
+    jit->code.movsxd(jit->code.rax, jit->code.eax);
+    jit->code.mov(GPR(idx), jit->code.rax);
+    jit->code.mov(jit->code.rax, 1);
+    jit->code.mov(GPR_constant_marker(idx), jit->code.rax);
+    return;
+  }
+
   gpr[idx] = v;
   gprIsConstant[idx] = true;
 }
@@ -111,6 +149,15 @@ template <>
 void Registers::Write<u16>(size_t idx, u16 v) {
   if (idx == 0)
     return;
+
+  if(jit) {
+    jit->code.mov(jit->code.rax, v);
+    jit->code.mov(GPR(idx), jit->code.rax);
+    jit->code.mov(jit->code.rax, 1);
+    jit->code.mov(GPR_constant_marker(idx), jit->code.rax);
+    return;
+  }
+
   gpr[idx] = (u16)v;
   gprIsConstant[idx] = true;
 }
@@ -119,6 +166,16 @@ template <>
 void Registers::Write<s16>(size_t idx, s16 v) {
   if (idx == 0)
     return;
+
+  if(jit) {
+    jit->code.mov(jit->code.ax, v);
+    jit->code.movsx(jit->code.rax, jit->code.ax);
+    jit->code.mov(GPR(idx), jit->code.rax);
+    jit->code.mov(jit->code.rax, 1);
+    jit->code.mov(GPR_constant_marker(idx), jit->code.rax);
+    return;
+  }
+
   gpr[idx] = v;
   gprIsConstant[idx] = true;
 }
@@ -127,6 +184,15 @@ template <>
 void Registers::Write<u8>(size_t idx, u8 v) {
   if (idx == 0)
     return;
+  
+  if(jit) {
+    jit->code.mov(jit->code.rax, v);
+    jit->code.mov(GPR(idx), jit->code.rax);
+    jit->code.mov(jit->code.rax, 1);
+    jit->code.mov(GPR_constant_marker(idx), jit->code.rax);
+    return;
+  }
+
   gpr[idx] = (u8)v;
   gprIsConstant[idx] = true;
 }
@@ -135,6 +201,16 @@ template <>
 void Registers::Write<s8>(size_t idx, s8 v) {
   if (idx == 0)
     return;
+  
+  if(jit) {
+    jit->code.mov(jit->code.al, v);
+    jit->code.movsx(jit->code.rax, jit->code.al);
+    jit->code.mov(GPR(idx), jit->code.rax);
+    jit->code.mov(jit->code.rax, 1);
+    jit->code.mov(GPR_constant_marker(idx), jit->code.rax);
+    return;
+  }
+  
   gpr[idx] = v;
   gprIsConstant[idx] = true;
 }
diff --git a/src/backend/core/registers/Registers.hpp b/src/backend/core/registers/Registers.hpp
index 82a283de..48402bae 100644
--- a/src/backend/core/registers/Registers.hpp
+++ b/src/backend/core/registers/Registers.hpp
@@ -3,8 +3,9 @@
 #include <backend/core/registers/Cop1.hpp>
 
 namespace n64 {
+struct JIT;
 struct Registers {
-  Registers();
+  Registers(JIT* jit = nullptr);
   void Reset();
   void SetPC64(s64);
   void SetPC32(s32);
@@ -19,6 +20,8 @@ struct Registers {
     return IsRegConstant(first) && IsRegConstant(second);
   }
 
+  JIT* jit = nullptr;
+
   std::array<bool, 32> gprIsConstant{};
   bool loIsConstant = false, hiIsConstant = false;
   Cop0 cop0;
@@ -41,6 +44,10 @@ struct Registers {
   T Read(size_t);
   template <typename T>
   void Write(size_t, T);
+
   std::array<s64, 32> gpr{};
+private:
+  template <typename T>
+  void WriteJIT(size_t, T);
 };
 } // namespace n64