JitArm64: Don't update dest reg when load triggers exception

Fixes a problem introduced in the previous commit.

4 files changed, 47 insertions, 18 deletions

diff --git a/Source/Core/Core/PowerPC/JitArm64/Jit.cpp b/Source/Core/Core/PowerPC/JitArm64/Jit.cpp
index 3bd330040f..f494f128f6 100644
--- a/Source/Core/Core/PowerPC/JitArm64/Jit.cpp
+++ b/Source/Core/Core/PowerPC/JitArm64/Jit.cpp
@@ -498,14 +498,15 @@ void JitArm64::WriteExceptionExit(ARM64Reg dest, bool only_external, bool always
   B(dispatcher);
 }
 
-void JitArm64::WriteConditionalExceptionExit(int exception)
+void JitArm64::WriteConditionalExceptionExit(int exception, u64 increment_sp_on_exit)
 {
   ARM64Reg WA = gpr.GetReg();
-  WriteConditionalExceptionExit(exception, WA);
+  WriteConditionalExceptionExit(exception, WA, increment_sp_on_exit);
   gpr.Unlock(WA);
 }
 
-void JitArm64::WriteConditionalExceptionExit(int exception, ARM64Reg temp_reg)
+void JitArm64::WriteConditionalExceptionExit(int exception, ARM64Reg temp_reg,
+                                             u64 increment_sp_on_exit)
 {
   LDR(IndexType::Unsigned, temp_reg, PPC_REG, PPCSTATE_OFF(Exceptions));
   FixupBranch no_exception = TBZ(temp_reg, IntLog2(exception));
@@ -519,6 +520,9 @@ void JitArm64::WriteConditionalExceptionExit(int exception, ARM64Reg temp_reg)
     SetJumpTarget(handle_exception);
   }
 
+  if (increment_sp_on_exit != 0)
+    ADDI2R(ARM64Reg::SP, ARM64Reg::SP, increment_sp_on_exit, temp_reg);
+
   gpr.Flush(FlushMode::MaintainState, temp_reg);
   fpr.Flush(FlushMode::MaintainState, ARM64Reg::INVALID_REG);
 
diff --git a/Source/Core/Core/PowerPC/JitArm64/Jit.h b/Source/Core/Core/PowerPC/JitArm64/Jit.h
index 2d63743300..a1368066cd 100644
--- a/Source/Core/Core/PowerPC/JitArm64/Jit.h
+++ b/Source/Core/Core/PowerPC/JitArm64/Jit.h
@@ -265,8 +265,9 @@ protected:
                           bool always_exception = false);
   void WriteExceptionExit(Arm64Gen::ARM64Reg dest, bool only_external = false,
                           bool always_exception = false);
-  void WriteConditionalExceptionExit(int exception);
-  void WriteConditionalExceptionExit(int exception, Arm64Gen::ARM64Reg temp_reg);
+  void WriteConditionalExceptionExit(int exception, u64 increment_sp_on_exit = 0);
+  void WriteConditionalExceptionExit(int exception, Arm64Gen::ARM64Reg temp_reg,
+                                     u64 increment_sp_on_exit = 0);
   void FakeLKExit(u32 exit_address_after_return);
   void WriteBLRExit(Arm64Gen::ARM64Reg dest);
 
diff --git a/Source/Core/Core/PowerPC/JitArm64/JitArm64_BackPatch.cpp b/Source/Core/Core/PowerPC/JitArm64/JitArm64_BackPatch.cpp
index 8dfcadcc09..4923e713cb 100644
--- a/Source/Core/Core/PowerPC/JitArm64/JitArm64_BackPatch.cpp
+++ b/Source/Core/Core/PowerPC/JitArm64/JitArm64_BackPatch.cpp
@@ -57,6 +57,8 @@ void JitArm64::EmitBackpatchRoutine(u32 flags, bool fastmem, bool do_farcode, AR
                                     ARM64Reg addr, BitSet32 gprs_to_push, BitSet32 fprs_to_push,
                                     bool emitting_routine)
 {
+  const u32 access_size = BackPatchInfo::GetFlagSize(flags);
+
   bool in_far_code = false;
   const u8* fastmem_start = GetCodePtr();
   std::optional<FixupBranch> slowmem_fixup;
@@ -76,11 +78,11 @@ void JitArm64::EmitBackpatchRoutine(u32 flags, bool fastmem, bool do_farcode, AR
       ARM64Reg temp = ARM64Reg::D0;
       temp = ByteswapBeforeStore(this, &m_float_emit, temp, EncodeRegToDouble(RS), flags, true);
 
-      m_float_emit.STR(BackPatchInfo::GetFlagSize(flags), temp, MEM_REG, addr);
+      m_float_emit.STR(access_size, temp, MEM_REG, addr);
     }
     else if ((flags & BackPatchInfo::FLAG_LOAD) && (flags & BackPatchInfo::FLAG_FLOAT))
     {
-      m_float_emit.LDR(BackPatchInfo::GetFlagSize(flags), EncodeRegToDouble(RS), MEM_REG, addr);
+      m_float_emit.LDR(access_size, EncodeRegToDouble(RS), MEM_REG, addr);
 
       ByteswapAfterLoad(this, &m_float_emit, EncodeRegToDouble(RS), EncodeRegToDouble(RS), flags,
                         true, false);
@@ -139,12 +141,22 @@ void JitArm64::EmitBackpatchRoutine(u32 flags, bool fastmem, bool do_farcode, AR
     if (slowmem_fixup)
       SetJumpTarget(*slowmem_fixup);
 
-    ABI_PushRegisters(gprs_to_push);
+    const ARM64Reg temp_reg = flags & BackPatchInfo::FLAG_LOAD ? ARM64Reg::W30 : ARM64Reg::W0;
+    const int temp_reg_index = DecodeReg(temp_reg);
+
+    if (memcheck && (flags & BackPatchInfo::FLAG_LOAD))
+    {
+      ABI_PushRegisters(gprs_to_push & BitSet32{temp_reg_index});
+      ABI_PushRegisters(gprs_to_push & ~BitSet32{temp_reg_index});
+    }
+    else
+    {
+      ABI_PushRegisters(gprs_to_push);
+    }
     m_float_emit.ABI_PushRegisters(fprs_to_push, ARM64Reg::X30);
 
     if (flags & BackPatchInfo::FLAG_STORE)
     {
-      const u32 access_size = BackPatchInfo::GetFlagSize(flags);
       ARM64Reg src_reg = RS;
       const ARM64Reg dst_reg = access_size == 64 ? ARM64Reg::X0 : ARM64Reg::W0;
 
@@ -188,8 +200,6 @@ void JitArm64::EmitBackpatchRoutine(u32 flags, bool fastmem, bool do_farcode, AR
     }
     else
     {
-      const u32 access_size = BackPatchInfo::GetFlagSize(flags);
-
       if (access_size == 64)
         MOVP2R(ARM64Reg::X8, &PowerPC::Read_U64);
       else if (access_size == 32)
@@ -200,8 +210,24 @@ void JitArm64::EmitBackpatchRoutine(u32 flags, bool fastmem, bool do_farcode, AR
         MOVP2R(ARM64Reg::X8, &PowerPC::Read_U8);
 
       BLR(ARM64Reg::X8);
+    }
+
+    m_float_emit.ABI_PopRegisters(fprs_to_push, ARM64Reg::X30);
+    if (memcheck && (flags & BackPatchInfo::FLAG_LOAD))
+      ABI_PopRegisters(gprs_to_push & ~BitSet32{temp_reg_index});
+    else
+      ABI_PopRegisters(gprs_to_push);
+
+    if (memcheck)
+    {
+      const u64 early_push_size = flags & BackPatchInfo::FLAG_LOAD ? 16 : 0;
+      WriteConditionalExceptionExit(EXCEPTION_DSI, temp_reg, early_push_size);
+    }
 
+    if (flags & BackPatchInfo::FLAG_LOAD)
+    {
       ARM64Reg src_reg = access_size == 64 ? ARM64Reg::X0 : ARM64Reg::W0;
+      ASSERT(!gprs_to_push[DecodeReg(src_reg)]);
 
       if (flags & BackPatchInfo::FLAG_PAIR)
       {
@@ -222,13 +248,10 @@ void JitArm64::EmitBackpatchRoutine(u32 flags, bool fastmem, bool do_farcode, AR
       }
 
       ByteswapAfterLoad(this, &m_float_emit, RS, src_reg, flags, false, false);
-    }
-
-    m_float_emit.ABI_PopRegisters(fprs_to_push, ARM64Reg::X30);
-    ABI_PopRegisters(gprs_to_push);
 
-    if (memcheck)
-      WriteConditionalExceptionExit(EXCEPTION_DSI, ARM64Reg::W0);
+      if (memcheck)
+        ABI_PopRegisters(gprs_to_push & BitSet32{temp_reg_index});
+    }
   }
 
   if (in_far_code)
diff --git a/Source/Core/Core/PowerPC/JitArm64/JitArm64_LoadStore.cpp b/Source/Core/Core/PowerPC/JitArm64/JitArm64_LoadStore.cpp
index 8398d1b677..38e693ca07 100644
--- a/Source/Core/Core/PowerPC/JitArm64/JitArm64_LoadStore.cpp
+++ b/Source/Core/Core/PowerPC/JitArm64/JitArm64_LoadStore.cpp
@@ -113,7 +113,8 @@ void JitArm64::SafeLoadToReg(u32 dest, s32 addr, s32 offsetReg, u32 flags, s32 o
   BitSet32 regs_in_use = gpr.GetCallerSavedUsed();
   BitSet32 fprs_in_use = fpr.GetCallerSavedUsed();
   regs_in_use[DecodeReg(ARM64Reg::W0)] = 0;
-  regs_in_use[DecodeReg(dest_reg)] = 0;
+  if (!jo.memcheck)
+    regs_in_use[DecodeReg(dest_reg)] = 0;
 
   u32 access_size = BackPatchInfo::GetFlagSize(flags);
   u32 mmio_address = 0;