JNI: Use callee-save register for IRT cookie.

Test: m test-art-host-gtest
Test: testrunner.py --host --optimizing
Test: run-gtests.sh
Test: testrunner.py --target --optimizing
Bug: 172332525
Change-Id: I6a91d86fd31ff33882b41646aae9fcccc157d638
diff --git a/compiler/jni/quick/arm/calling_convention_arm.cc b/compiler/jni/quick/arm/calling_convention_arm.cc
index bc1e866..d849c28 100644
--- a/compiler/jni/quick/arm/calling_convention_arm.cc
+++ b/compiler/jni/quick/arm/calling_convention_arm.cc
@@ -385,6 +385,13 @@
   return is_critical_native_ ? 0u : kFpCalleeSpillMask;
 }
 
+ManagedRegister ArmJniCallingConvention::SavedLocalReferenceCookieRegister() const {
+  // The r5 is callee-save register in both managed and native ABIs.
+  // It is saved in the stack frame and it has no special purpose like `tr`.
+  static_assert((kCoreCalleeSpillMask & (1u << R5)) != 0u);  // Managed callee save register.
+  return ArmManagedRegister::FromCoreRegister(R5);
+}
+
 ManagedRegister ArmJniCallingConvention::ReturnScratchRegister() const {
   return ArmManagedRegister::FromCoreRegister(R2);
 }
@@ -404,14 +411,18 @@
   size_t total_size = method_ptr_size + callee_save_area_size;
 
   DCHECK(HasLocalReferenceSegmentState());
-  const size_t cookie_size = SavedLocalReferenceCookieSize();
-  total_size += cookie_size;
+  // Cookie is saved in one of the spilled registers.
 
   // Plus return value spill area size
   if (SpillsReturnValue()) {
-    // No padding between cookie and return value on arm.
-    DCHECK_EQ(ReturnValueSaveLocation().SizeValue(),
-              SavedLocalReferenceCookieOffset().SizeValue() + cookie_size);
+    // For 64-bit return values there shall be a 4B alignment gap between
+    // the method pointer and the saved return value.
+    size_t padding = ReturnValueSaveLocation().SizeValue() - method_ptr_size;
+    DCHECK_EQ(padding,
+              (GetReturnType() == Primitive::kPrimLong || GetReturnType() == Primitive::kPrimDouble)
+                  ? 4u
+                  : 0u);
+    total_size += padding;
     total_size += SizeOfReturnValue();
   }
 
diff --git a/compiler/jni/quick/arm/calling_convention_arm.h b/compiler/jni/quick/arm/calling_convention_arm.h
index 38f7184..985d971 100644
--- a/compiler/jni/quick/arm/calling_convention_arm.h
+++ b/compiler/jni/quick/arm/calling_convention_arm.h
@@ -67,6 +67,7 @@
   size_t FrameSize() const override;
   size_t OutFrameSize() const override;
   ArrayRef<const ManagedRegister> CalleeSaveRegisters() const override;
+  ManagedRegister SavedLocalReferenceCookieRegister() const override;
   ManagedRegister ReturnScratchRegister() const override;
   uint32_t CoreSpillMask() const override;
   uint32_t FpSpillMask() const override;
diff --git a/compiler/jni/quick/arm64/calling_convention_arm64.cc b/compiler/jni/quick/arm64/calling_convention_arm64.cc
index 8d40f2e..1a13689 100644
--- a/compiler/jni/quick/arm64/calling_convention_arm64.cc
+++ b/compiler/jni/quick/arm64/calling_convention_arm64.cc
@@ -232,6 +232,13 @@
   return is_critical_native_ ? 0u : kFpCalleeSpillMask;
 }
 
+ManagedRegister Arm64JniCallingConvention::SavedLocalReferenceCookieRegister() const {
+  // The w21 is callee-save register in both managed and native ABIs.
+  // It is saved in the stack frame and it has no special purpose like `tr`.
+  static_assert((kCoreCalleeSpillMask & (1u << W21)) != 0u);  // Managed callee save register.
+  return Arm64ManagedRegister::FromWRegister(W21);
+}
+
 ManagedRegister Arm64JniCallingConvention::ReturnScratchRegister() const {
   return ManagedRegister::NoRegister();
 }
@@ -251,17 +258,12 @@
   size_t total_size = method_ptr_size + callee_save_area_size;
 
   DCHECK(HasLocalReferenceSegmentState());
-  const size_t cookie_size = SavedLocalReferenceCookieSize();
-  total_size += cookie_size;
+  // Cookie is saved in one of the spilled registers.
 
   // Plus return value spill area size
   if (SpillsReturnValue()) {
-    // For 64-bit return values there shall be a 4B alignment gap between the cookie
-    // and the saved return value. However, we do not need to round the intermediate
-    // `total_size` here as the final rounding below shall add sufficient padding.
-    DCHECK_ALIGNED(total_size, 4u);
-    DCHECK(!IsAligned<8u>(total_size));
-    static_assert(IsAligned<8u>(kStackAlignment));
+    // No padding between the method pointer and the return value on arm64.
+    DCHECK_EQ(ReturnValueSaveLocation().SizeValue(), method_ptr_size);
     total_size += SizeOfReturnValue();
   }
 
diff --git a/compiler/jni/quick/arm64/calling_convention_arm64.h b/compiler/jni/quick/arm64/calling_convention_arm64.h
index d381d9d..e1e9407 100644
--- a/compiler/jni/quick/arm64/calling_convention_arm64.h
+++ b/compiler/jni/quick/arm64/calling_convention_arm64.h
@@ -58,6 +58,7 @@
   size_t FrameSize() const override;
   size_t OutFrameSize() const override;
   ArrayRef<const ManagedRegister> CalleeSaveRegisters() const override;
+  ManagedRegister SavedLocalReferenceCookieRegister() const override;
   ManagedRegister ReturnScratchRegister() const override;
   uint32_t CoreSpillMask() const override;
   uint32_t FpSpillMask() const override;
diff --git a/compiler/jni/quick/calling_convention.cc b/compiler/jni/quick/calling_convention.cc
index 2127f73..fd05941 100644
--- a/compiler/jni/quick/calling_convention.cc
+++ b/compiler/jni/quick/calling_convention.cc
@@ -173,18 +173,10 @@
   return NumReferenceArgs() + (IsStatic() ? 1 : 0);
 }
 
-FrameOffset JniCallingConvention::SavedLocalReferenceCookieOffset() const {
-  // The cookie goes after the method pointer.
-  DCHECK_EQ(SavedLocalReferenceCookieSize(), sizeof(IRTSegmentState));
-  DCHECK(HasLocalReferenceSegmentState());
-  return FrameOffset(displacement_.SizeValue() + static_cast<size_t>(frame_pointer_size_));
-}
-
 FrameOffset JniCallingConvention::ReturnValueSaveLocation() const {
-  // The saved return value goes at a properly aligned slot after the cookie.
+  // The saved return value goes at a properly aligned slot after the method pointer.
   DCHECK(SpillsReturnValue());
-  size_t cookie_offset = SavedLocalReferenceCookieOffset().SizeValue() - displacement_.SizeValue();
-  size_t return_value_offset = cookie_offset + SavedLocalReferenceCookieSize();
+  size_t return_value_offset = static_cast<size_t>(frame_pointer_size_);
   const size_t return_value_size = SizeOfReturnValue();
   DCHECK(return_value_size == 4u || return_value_size == 8u) << return_value_size;
   DCHECK_ALIGNED(return_value_offset, 4u);
diff --git a/compiler/jni/quick/calling_convention.h b/compiler/jni/quick/calling_convention.h
index 5679263..c11e09d 100644
--- a/compiler/jni/quick/calling_convention.h
+++ b/compiler/jni/quick/calling_convention.h
@@ -304,8 +304,6 @@
   virtual size_t OutFrameSize() const = 0;
   // Number of references in stack indirect reference table
   size_t ReferenceCount() const;
-  // Location where the segment state of the local indirect reference table is saved
-  FrameOffset SavedLocalReferenceCookieOffset() const;
   // Location where the return value of a call can be squirreled if another
   // call is made following the native call
   FrameOffset ReturnValueSaveLocation() const;
@@ -317,13 +315,17 @@
   // Callee save registers to spill prior to native code (which may clobber)
   virtual ArrayRef<const ManagedRegister> CalleeSaveRegisters() const = 0;
 
-  // Spill mask values
-  virtual uint32_t CoreSpillMask() const = 0;
-  virtual uint32_t FpSpillMask() const = 0;
+  // Register where the segment state of the local indirect reference table is saved.
+  // This must be a native callee-save register without another special purpose.
+  virtual ManagedRegister SavedLocalReferenceCookieRegister() const = 0;
 
   // An extra scratch register live after the call
   virtual ManagedRegister ReturnScratchRegister() const = 0;
 
+  // Spill mask values
+  virtual uint32_t CoreSpillMask() const = 0;
+  virtual uint32_t FpSpillMask() const = 0;
+
   // Iterator interface
   bool HasNext();
   virtual void Next();
@@ -343,7 +345,7 @@
 
   virtual ~JniCallingConvention() {}
 
-  size_t SavedLocalReferenceCookieSize() const {
+  static constexpr size_t SavedLocalReferenceCookieSize() {
     return 4u;
   }
 
diff --git a/compiler/jni/quick/jni_compiler.cc b/compiler/jni/quick/jni_compiler.cc
index a4be7ef..cdd0263 100644
--- a/compiler/jni/quick/jni_compiler.cc
+++ b/compiler/jni/quick/jni_compiler.cc
@@ -278,8 +278,8 @@
   //    can occur. The result is the saved JNI local state that is restored by the exit call. We
   //    abuse the JNI calling convention here, that is guaranteed to support passing 2 pointer
   //    arguments.
-  FrameOffset saved_cookie_offset(
-      FrameOffset(0xDEADBEEFu));  // @CriticalNative - use obviously bad value for debugging
+  constexpr size_t cookie_size = JniCallingConvention::SavedLocalReferenceCookieSize();
+  ManagedRegister saved_cookie_register = ManagedRegister::NoRegister();
   if (LIKELY(!is_critical_native)) {
     // Skip this for @CriticalNative methods. They do not call JniMethodStart.
     ThreadOffset<kPointerSize> jni_start(
@@ -324,8 +324,8 @@
     }
 
     // Store into stack_frame[saved_cookie_offset] the return value of JniMethodStart.
-    saved_cookie_offset = main_jni_conv->SavedLocalReferenceCookieOffset();
-    __ Store(saved_cookie_offset, main_jni_conv->IntReturnRegister(), 4 /* sizeof cookie */);
+    saved_cookie_register = main_jni_conv->SavedLocalReferenceCookieRegister();
+    __ Move(saved_cookie_register, main_jni_conv->IntReturnRegister(), cookie_size);
   }
 
   // 6. Fill arguments.
@@ -500,7 +500,6 @@
       current_out_arg_size = end_out_arg_size;
       __ IncreaseFrameSize(out_arg_size_diff);
       current_frame_size += out_arg_size_diff;
-      saved_cookie_offset = FrameOffset(saved_cookie_offset.SizeValue() + out_arg_size_diff);
       return_save_location = FrameOffset(return_save_location.SizeValue() + out_arg_size_diff);
     }
     end_jni_conv->ResetIterator(FrameOffset(end_out_arg_size));
@@ -519,10 +518,10 @@
     // Pass saved local reference state.
     if (end_jni_conv->IsCurrentParamOnStack()) {
       FrameOffset out_off = end_jni_conv->CurrentParamStackOffset();
-      __ Copy(out_off, saved_cookie_offset, 4);
+      __ Store(out_off, saved_cookie_register, cookie_size);
     } else {
       ManagedRegister out_reg = end_jni_conv->CurrentParamRegister();
-      __ Load(out_reg, saved_cookie_offset, 4);
+      __ Move(out_reg, saved_cookie_register, cookie_size);
     }
     end_jni_conv->Next();
     if (is_synchronized) {
diff --git a/compiler/jni/quick/x86/calling_convention_x86.cc b/compiler/jni/quick/x86/calling_convention_x86.cc
index d624831..1baffc5 100644
--- a/compiler/jni/quick/x86/calling_convention_x86.cc
+++ b/compiler/jni/quick/x86/calling_convention_x86.cc
@@ -71,6 +71,13 @@
 
 // Calling convention
 
+ManagedRegister X86JniCallingConvention::SavedLocalReferenceCookieRegister() const {
+  // The EBP is callee-save register in both managed and native ABIs.
+  // It is saved in the stack frame and it has no special purpose like `tr` on arm/arm64.
+  static_assert((kCoreCalleeSpillMask & (1u << EBP)) != 0u);  // Managed callee save register.
+  return X86ManagedRegister::FromCpuRegister(EBP);
+}
+
 ManagedRegister X86JniCallingConvention::ReturnScratchRegister() const {
   return ManagedRegister::NoRegister();  // No free regs, so assembler uses push/pop
 }
@@ -206,14 +213,18 @@
   size_t total_size = method_ptr_size + pc_return_addr_size + callee_save_area_size;
 
   DCHECK(HasLocalReferenceSegmentState());
-  const size_t cookie_size = SavedLocalReferenceCookieSize();
-  total_size += cookie_size;
+  // Cookie is saved in one of the spilled registers.
 
   // Plus return value spill area size
   if (SpillsReturnValue()) {
-    // No padding between cookie and return value on x86.
-    DCHECK_EQ(ReturnValueSaveLocation().SizeValue(),
-              SavedLocalReferenceCookieOffset().SizeValue() + cookie_size);
+    // For 64-bit return values there shall be a 4B alignment gap between
+    // the method pointer and the saved return value.
+    size_t padding = ReturnValueSaveLocation().SizeValue() - method_ptr_size;
+    DCHECK_EQ(padding,
+              (GetReturnType() == Primitive::kPrimLong || GetReturnType() == Primitive::kPrimDouble)
+                  ? 4u
+                  : 0u);
+    total_size += padding;
     total_size += SizeOfReturnValue();
   }
 
diff --git a/compiler/jni/quick/x86/calling_convention_x86.h b/compiler/jni/quick/x86/calling_convention_x86.h
index 81f617d..cbb362c 100644
--- a/compiler/jni/quick/x86/calling_convention_x86.h
+++ b/compiler/jni/quick/x86/calling_convention_x86.h
@@ -63,6 +63,7 @@
   size_t FrameSize() const override;
   size_t OutFrameSize() const override;
   ArrayRef<const ManagedRegister> CalleeSaveRegisters() const override;
+  ManagedRegister SavedLocalReferenceCookieRegister() const override;
   ManagedRegister ReturnScratchRegister() const override;
   uint32_t CoreSpillMask() const override;
   uint32_t FpSpillMask() const override;
diff --git a/compiler/jni/quick/x86_64/calling_convention_x86_64.cc b/compiler/jni/quick/x86_64/calling_convention_x86_64.cc
index bb01371..33a921b 100644
--- a/compiler/jni/quick/x86_64/calling_convention_x86_64.cc
+++ b/compiler/jni/quick/x86_64/calling_convention_x86_64.cc
@@ -91,6 +91,13 @@
 
 // Calling convention
 
+ManagedRegister X86_64JniCallingConvention::SavedLocalReferenceCookieRegister() const {
+  // The RBX is callee-save register in both managed and native ABIs.
+  // It is saved in the stack frame and it has no special purpose like `tr` on arm/arm64.
+  static_assert((kCoreCalleeSpillMask & (1u << RBX)) != 0u);  // Managed callee save register.
+  return X86_64ManagedRegister::FromCpuRegister(RBX);
+}
+
 ManagedRegister X86_64JniCallingConvention::ReturnScratchRegister() const {
   return ManagedRegister::NoRegister();  // No free regs, so assembler uses push/pop
 }
@@ -194,17 +201,12 @@
   size_t total_size = method_ptr_size + pc_return_addr_size + callee_save_area_size;
 
   DCHECK(HasLocalReferenceSegmentState());
-  const size_t cookie_size = SavedLocalReferenceCookieSize();
-  total_size += cookie_size;
+  // Cookie is saved in one of the spilled registers.
 
   // Plus return value spill area size
   if (SpillsReturnValue()) {
-    // For 64-bit return values there shall be a 4B alignment gap between the cookie
-    // and the saved return value. However, we do not need to round the intermediate
-    // `total_size` here as the final rounding below shall add sufficient padding.
-    DCHECK_ALIGNED(total_size, 4u);
-    DCHECK(!IsAligned<8u>(total_size));
-    static_assert(IsAligned<8u>(kStackAlignment));
+    // No padding between the method pointer and the return value on arm64.
+    DCHECK_EQ(ReturnValueSaveLocation().SizeValue(), method_ptr_size);
     total_size += SizeOfReturnValue();
   }
 
diff --git a/compiler/jni/quick/x86_64/calling_convention_x86_64.h b/compiler/jni/quick/x86_64/calling_convention_x86_64.h
index 5bde766..f9d6fc0 100644
--- a/compiler/jni/quick/x86_64/calling_convention_x86_64.h
+++ b/compiler/jni/quick/x86_64/calling_convention_x86_64.h
@@ -58,6 +58,7 @@
   size_t FrameSize() const override;
   size_t OutFrameSize() const override;
   ArrayRef<const ManagedRegister> CalleeSaveRegisters() const override;
+  ManagedRegister SavedLocalReferenceCookieRegister() const override;
   ManagedRegister ReturnScratchRegister() const override;
   uint32_t CoreSpillMask() const override;
   uint32_t FpSpillMask() const override;