[vm/ffi] Fix safepoint implementation in ARM64 assembler.

Fixes #37180

Change-Id: I4be04fbd98be2c724218da89fbd2bea78cda0c26
Reviewed-on: https://dart-review.googlesource.com/c/sdk/+/106081
Auto-Submit: Samir Jindel <sjindel@google.com>
Commit-Queue: Daco Harkes <dacoharkes@google.com>
Reviewed-by: Daco Harkes <dacoharkes@google.com>
diff --git a/runtime/vm/compiler/backend/il_arm64.cc b/runtime/vm/compiler/backend/il_arm64.cc
index 154be8b..2fb105f 100644
--- a/runtime/vm/compiler/backend/il_arm64.cc
+++ b/runtime/vm/compiler/backend/il_arm64.cc
@@ -887,9 +887,6 @@
   __ set_constant_pool_allowed(false);
   __ EnterDartFrame(0, PP);
 
-  // Save the stack limit address.
-  __ PushRegister(CSP);
-
   // Make space for arguments and align the frame.
   __ ReserveAlignedFrameSpace(compiler::ffi::NumStackSlots(arg_locations_) *
                               kWordSize);
@@ -911,22 +908,21 @@
 
   __ StoreToOffset(temp, FPREG, kSavedCallerPcSlotFromFp * kWordSize);
 
+  // Update information in the thread object and enter a safepoint.
+  __ TransitionGeneratedToNative(branch, FPREG, temp);
+
   // We are entering runtime code, so the C stack pointer must be restored from
   // the stack limit to the top of the stack.
   __ mov(CSP, SP);
 
-  // Update information in the thread object and enter a safepoint.
-  __ TransitionGeneratedToNative(branch, FPREG, temp);
-
   __ blr(branch);
 
+  // Restore the Dart stack pointer.
+  __ mov(SP, CSP);
+
   // Update information in the thread object and leave the safepoint.
   __ TransitionNativeToGenerated(temp);
 
-  // Restore the Dart stack pointer and the saved C stack pointer.
-  __ mov(SP, CSP);
-  __ LoadFromOffset(CSP, FPREG, kFirstLocalSlotFromFp * kWordSize);
-
   // Refresh write barrier mask.
   __ ldr(BARRIER_MASK,
          Address(THR, compiler::target::Thread::write_barrier_mask_offset()));
diff --git a/runtime/vm/compiler/stub_code_compiler_arm64.cc b/runtime/vm/compiler/stub_code_compiler_arm64.cc
index 97dcff4..2f3b1b0 100644
--- a/runtime/vm/compiler/stub_code_compiler_arm64.cc
+++ b/runtime/vm/compiler/stub_code_compiler_arm64.cc
@@ -208,34 +208,40 @@
 void StubCodeCompiler::GenerateEnterSafepointStub(Assembler* assembler) {
   RegisterSet all_registers;
   all_registers.AddAllGeneralRegisters();
-  __ PushRegisters(all_registers);
 
   __ EnterFrame(0);
+  __ PushRegisters(all_registers);
+
+  __ mov(CALLEE_SAVED_TEMP, SP);
   __ ReserveAlignedFrameSpace(0);
+
   __ mov(CSP, SP);
   __ ldr(R0, Address(THR, kEnterSafepointRuntimeEntry.OffsetFromThread()));
   __ blr(R0);
-  __ LeaveFrame();
+  __ mov(SP, CALLEE_SAVED_TEMP);
 
   __ PopRegisters(all_registers);
-  __ mov(CSP, SP);
+  __ LeaveFrame();
   __ Ret();
 }
 
 void StubCodeCompiler::GenerateExitSafepointStub(Assembler* assembler) {
   RegisterSet all_registers;
   all_registers.AddAllGeneralRegisters();
-  __ PushRegisters(all_registers);
 
   __ EnterFrame(0);
+  __ PushRegisters(all_registers);
+
+  __ mov(CALLEE_SAVED_TEMP, SP);
   __ ReserveAlignedFrameSpace(0);
+
   __ mov(CSP, SP);
   __ ldr(R0, Address(THR, kExitSafepointRuntimeEntry.OffsetFromThread()));
   __ blr(R0);
-  __ LeaveFrame();
+  __ mov(SP, CALLEE_SAVED_TEMP);
 
   __ PopRegisters(all_registers);
-  __ mov(CSP, SP);
+  __ LeaveFrame();
   __ Ret();
 }