JNI: Keep previous LRT state in callee-save reg.

Golem results for art-opt-cc (higher is better):
linux-armv7 (Odroid-C2)          before after
NativeDowncallStaticFast         21.622 21.923 (+1.395%)
NativeDowncallStaticFast6        18.491 18.719 (+1.236%)
NativeDowncallStaticFastRefs6    15.347 15.504 (+1.025%)
NativeDowncallVirtualFast        20.741 21.319 (+2.787%)
NativeDowncallVirtualFast6       18.953 19.183 (+1.218%)
NativeDowncallVirtualFastRefs6   15.500 15.663 (+1.053%)
NativeDowncallStaticNormal       14.620 14.757 (0.9495%)
NativeDowncallStaticNormal6      13.120 13.235 (+0.8823%)
NativeDowncallStaticNormalRefs6  11.454 11.538 (+0.7258%)
NativeDowncallVirtualNormal      14.216 14.486 (+1.898%)
NativeDowncallVirtualNormal6     13.347 13.466 (+0.8978%)
NativeDowncallVirtualNormalRefs6 11.538 11.628 (+0.7752%)
linux-armv7 (Raspberry Pi 4)     before after
NativeDowncallStaticFast         43.305 42.331 (-2.250%)
NativeDowncallStaticFast6        35.608 37.369 (+4.945%)
NativeDowncallStaticFastRefs6    31.390 31.793 (+1.285%)
NativeDowncallVirtualFast        33.814 31.825 (-5.882%)
NativeDowncallVirtualFast6       34.311 36.445 (+6.220%)
NativeDowncallVirtualFastRefs6   31.762 32.419 (+2.069%)
NativeDowncallStaticNormal       13.848 14.244 (+2.859%)
NativeDowncallStaticNormal6      13.592 13.725 (+0.9804%)
NativeDowncallStaticNormalRefs6  12.671 12.536 (-1.061%)
NativeDowncallVirtualNormal      13.979 13.848 (-0.9397%)
NativeDowncallVirtualNormal6     13.242 13.592 (+2.647%)
NativeDowncallVirtualNormalRefs6 12.364 12.358 (-0.094%)
linux-armv8 (Odroid-C2)          before after
NativeDowncallStaticFast         24.752 25.160 (+1.648%)
NativeDowncallStaticFast6        22.571 22.908 (+1.494%)
NativeDowncallStaticFastRefs6    19.183 19.183 (unchanged)
NativeDowncallVirtualFast        21.622 22.244 (+2.879%)
NativeDowncallVirtualFast6       21.319 21.934 (+2.887%)
NativeDowncallVirtualFastRefs6   17.448 17.848 (+2.296%)
NativeDowncallStaticNormal       17.048 17.250 (+1.183%)
NativeDowncallStaticNormal6      15.992 16.161 (+1.054%)
NativeDowncallStaticNormalRefs6  14.085 14.216 (+0.9314%)
NativeDowncallVirtualNormal      15.504 15.826 (+2.077%)
NativeDowncallVirtualNormal6     15.347 15.663 (+2.064%)
NativeDowncallVirtualNormalRefs6 13.466 13.586 (+0.8859%)
linux-armv8 (Raspberry Pi 4)     before after
NativeDowncallStaticFast         38.366 40.796 (+6.335%)
NativeDowncallStaticFast6        38.347 40.419 (+5.405%)
NativeDowncallStaticFastRefs6    31.636 32.528 (+2.820%)
NativeDowncallVirtualFast        35.201 37.406 (+6.266%)
NativeDowncallVirtualFast6       34.000 35.626 (+4.782%)
NativeDowncallVirtualFastRefs6   27.201 27.201 (unchanged)
NativeDowncallStaticNormal       14.808 15.107 (+2.024%)
NativeDowncallStaticNormal6      14.955 14.428 (-3.526%)
NativeDowncallStaticNormalRefs6  14.174 13.855 (-2.254%)
NativeDowncallVirtualNormal      14.735 14.307 (-2.904%)
NativeDowncallVirtualNormal6     14.244 14.385 (+0.9921%)
NativeDowncallVirtualNormalRefs6 14.105 14.244 (+0.9823%)
linux-ia32                       before after
NativeDowncallStaticFast         223.66 233.77 (+4.516%)
NativeDowncallStaticFast6        159.76 163.92 (+2.602%)
NativeDowncallStaticFastRefs6    137.16 141.72 (+3.324%)
NativeDowncallVirtualFast        211.79 224.05 (+5.791%)
NativeDowncallVirtualFast6       149.85 154.00 (+2.769%)
NativeDowncallVirtualFastRefs6   132.17 136.93 (+3.603%)
NativeDowncallStaticNormal       51.091 51.091 (unchanged)
NativeDowncallStaticNormal6      45.680 45.703 (+0.0497%)
NativeDowncallStaticNormalRefs6  44.732 45.161 (+0.9606%)
NativeDowncallVirtualNormal      50.450 50.450 (unchanged)
NativeDowncallVirtualNormal6     45.161 45.161 (unchanged)
NativeDowncallVirtualNormalRefs6 44.125 44.147 (+0.496%)
linux-x64                        before after
NativeDowncallStaticFast         173.07 181.05 (+4.611%)
NativeDowncallStaticFast6        156.50 161.34 (+3.092%)
NativeDowncallStaticFastRefs6    130.37 131.61 (+0.9499%)
NativeDowncallVirtualFast        169.00 174.83 (+3.447%)
NativeDowncallVirtualFast6       148.13 149.35 (+0.8243%)
NativeDowncallVirtualFastRefs6   127.31 130.11 (+2.200%)
NativeDowncallStaticNormal       47.952 47.952 (unchanged)
NativeDowncallStaticNormal6      46.789 46.789 (unchanged)
NativeDowncallStaticNormalRefs6  44.643 44.643 (unchanged)
NativeDowncallVirtualNormal      47.358 47.358 (unchanged)
NativeDowncallVirtualNormal6     45.703 45.680 (-0.0497%)
NativeDowncallVirtualNormalRefs6 44.643 44.643 (unchanged)

Test: m test-art-host-gtest
Test: testrunner.py --host
Test: run-gtests.sh
Test: testrunner.py --target --optimizing
Bug: 172332525
Change-Id: I9606412c658cae8b7583308facf5ba095a982349

11 files changed, 136 insertions, 130 deletions

diff --git a/compiler/jni/quick/calling_convention.h b/compiler/jni/quick/calling_convention.h
index 21336edb44..b8b4cc14b1 100644
--- a/compiler/jni/quick/calling_convention.h
+++ b/compiler/jni/quick/calling_convention.h
@@ -324,9 +324,11 @@ class JniCallingConvention : public CallingConvention {
   virtual ArrayRef<const ManagedRegister> CalleeSaveRegisters() const = 0;
 
   // Subset of core callee save registers that can be used for arbitrary purposes after
-  // constructing the JNI transition frame. These should be managed callee-saves as well.
+  // constructing the JNI transition frame. These should be both managed and native callee-saves.
   // These should not include special purpose registers such as thread register.
-  // JNI compiler currently requires at least 3 callee save scratch registers.
+  // JNI compiler currently requires at least 4 callee save scratch registers, except for x86
+  // where we have only 3 such registers but all args are passed on stack, so the method register
+  // is never clobbered by argument moves and does not need to be preserved elsewhere.
   virtual ArrayRef<const ManagedRegister> CalleeSaveScratchRegisters() const = 0;
 
   // Subset of core argument registers that can be used for arbitrary purposes after
diff --git a/compiler/jni/quick/jni_compiler.cc b/compiler/jni/quick/jni_compiler.cc
index 1c077a1497..c721825683 100644
--- a/compiler/jni/quick/jni_compiler.cc
+++ b/compiler/jni/quick/jni_compiler.cc
@@ -37,6 +37,7 @@
 #include "driver/compiler_options.h"
 #include "entrypoints/quick/quick_entrypoints.h"
 #include "instrumentation.h"
+#include "jni/jni_env_ext.h"
 #include "jni/local_reference_table.h"
 #include "runtime.h"
 #include "thread.h"
@@ -293,7 +294,8 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
   // 3. Push local reference frame.
   // Skip this for @CriticalNative methods, they cannot use any references.
   ManagedRegister jni_env_reg = ManagedRegister::NoRegister();
-  ManagedRegister saved_cookie_reg = ManagedRegister::NoRegister();
+  ManagedRegister previous_state_reg = ManagedRegister::NoRegister();
+  ManagedRegister current_state_reg = ManagedRegister::NoRegister();
   ManagedRegister callee_save_temp = ManagedRegister::NoRegister();
   if (LIKELY(!is_critical_native)) {
     // To pop the local reference frame later, we shall need the JNI environment pointer
@@ -301,14 +303,21 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
     CHECK_GE(callee_save_scratch_regs.size(), 3u);  // At least 3 for each supported architecture.
     jni_env_reg = callee_save_scratch_regs[0];
     constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
-    saved_cookie_reg = __ CoreRegisterWithSize(callee_save_scratch_regs[1], kLRTSegmentStateSize);
-    callee_save_temp = __ CoreRegisterWithSize(callee_save_scratch_regs[2], kLRTSegmentStateSize);
+    previous_state_reg = __ CoreRegisterWithSize(callee_save_scratch_regs[1], kLRTSegmentStateSize);
+    current_state_reg = __ CoreRegisterWithSize(callee_save_scratch_regs[2], kLRTSegmentStateSize);
+    if (callee_save_scratch_regs.size() >= 4) {
+      callee_save_temp = callee_save_scratch_regs[3];
+    }
+    const MemberOffset previous_state_offset = JNIEnvExt::LrtPreviousStateOffset(kPointerSize);
 
     // Load the JNI environment pointer.
     __ LoadRawPtrFromThread(jni_env_reg, Thread::JniEnvOffset<kPointerSize>());
 
-    // Push the local reference frame.
-    __ PushLocalReferenceFrame(jni_env_reg, saved_cookie_reg, callee_save_temp);
+    // Load the local reference frame states.
+    __ LoadLocalReferenceTableStates(jni_env_reg, previous_state_reg, current_state_reg);
+
+    // Store the current state as the previous state (push the LRT frame).
+    __ Store(jni_env_reg, previous_state_offset, current_state_reg, kLRTSegmentStateSize);
   }
 
   // 4. Make the main native call.
@@ -331,6 +340,7 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
   refs.clear();
   mr_conv->ResetIterator(FrameOffset(current_frame_size));
   main_jni_conv->ResetIterator(FrameOffset(main_out_arg_size));
+  bool check_method_not_clobbered = false;
   if (UNLIKELY(is_critical_native)) {
     // Move the method pointer to the hidden argument register.
     // TODO: Pass this as the last argument, not first. Change ARM assembler
@@ -341,10 +351,16 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
   } else {
     main_jni_conv->Next();    // Skip JNIEnv*.
     FrameOffset method_offset(current_out_arg_size + mr_conv->MethodStackOffset().SizeValue());
-    if (!is_static || main_jni_conv->IsCurrentParamOnStack()) {
+    if (main_jni_conv->IsCurrentParamOnStack()) {
+      // This is for x86 only. The method shall not be clobbered by argument moves
+      // because all arguments are passed on the stack to the native method.
+      check_method_not_clobbered = true;
+      DCHECK(callee_save_temp.IsNoRegister());
+    } else if (!is_static) {
       // The method shall not be available in the `jclass` argument register.
       // Make sure it is available in `callee_save_temp` for the call below.
       // (The old method register can be clobbered by argument moves.)
+      DCHECK(!callee_save_temp.IsNoRegister());
       ManagedRegister new_method_reg = __ CoreRegisterWithSize(callee_save_temp, kRawPointerSize);
       DCHECK(!method_register.IsNoRegister());
       __ Move(new_method_reg, method_register, kRawPointerSize);
@@ -386,6 +402,10 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
     refs.push_back(is_reference ? mr_conv->CurrentParamStackOffset() : kInvalidReferenceOffset);
   }
   DCHECK(!main_jni_conv->HasNext());
+  DCHECK_IMPLIES(check_method_not_clobbered,
+                 std::all_of(dest_args.begin(),
+                             dest_args.end(),
+                             [](const ArgumentLocation& loc) { return !loc.IsRegister(); }));
   __ MoveArguments(ArrayRef<ArgumentLocation>(dest_args),
                    ArrayRef<ArgumentLocation>(src_args),
                    ArrayRef<FrameOffset>(refs));
@@ -512,7 +532,10 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
 
   // 6. Pop local reference frame.
   if (LIKELY(!is_critical_native)) {
-    __ PopLocalReferenceFrame(jni_env_reg, saved_cookie_reg, callee_save_temp);
+    __ StoreLocalReferenceTableStates(jni_env_reg, previous_state_reg, current_state_reg);
+    // For x86, the `callee_save_temp` is not valid, so let's simply change it to one
+    // of the callee save registers that we don't need anymore for all architectures.
+    callee_save_temp = current_state_reg;
   }
 
   // 7. Return from the JNI stub.
@@ -644,7 +667,7 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
         jni_asm->cfi().AdjustCFAOffset(main_out_arg_size);
         __ DecreaseFrameSize(main_out_arg_size);
       }
-      __ PopLocalReferenceFrame(jni_env_reg, saved_cookie_reg, callee_save_temp);
+      __ StoreLocalReferenceTableStates(jni_env_reg, previous_state_reg, current_state_reg);
     }
     DCHECK_EQ(jni_asm->cfi().GetCurrentCFAOffset(), static_cast<int>(current_frame_size));
     __ DeliverPendingException();
diff --git a/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc b/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc
index 851d33acfd..9930e39b4b 100644
--- a/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc
+++ b/compiler/utils/arm/jni_macro_assembler_arm_vixl.cc
@@ -1110,42 +1110,34 @@ void ArmVIXLJNIMacroAssembler::Load(ArmManagedRegister dest,
   }
 }
 
-void ArmVIXLJNIMacroAssembler::PushLocalReferenceFrame(ManagedRegister jni_env_reg,
-                                                       ManagedRegister saved_cookie_reg,
-                                                       ManagedRegister temp_reg) {
+void ArmVIXLJNIMacroAssembler::LoadLocalReferenceTableStates(ManagedRegister jni_env_reg,
+                                                             ManagedRegister previous_state_reg,
+                                                             ManagedRegister current_state_reg) {
   constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
   DCHECK_EQ(kLRTSegmentStateSize, kRegSizeInBytes);
-  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kArmPointerSize);
-  const MemberOffset jni_env_segment_state_offset = JNIEnvExt::SegmentStateOffset(kArmPointerSize);
-  DCHECK_EQ(jni_env_cookie_offset.SizeValue() + kLRTSegmentStateSize,
-            jni_env_segment_state_offset.SizeValue());
-
-  // Load the old cookie that we shall need to restore together with the current segment state.
-  ___ Ldrd(AsVIXLRegister(saved_cookie_reg.AsArm()),
-           AsVIXLRegister(temp_reg.AsArm()),
-           MemOperand(AsVIXLRegister(jni_env_reg.AsArm()), jni_env_cookie_offset.Int32Value()));
-
-  // Set the cookie in JNI environment to the current segment state.
-  Store(jni_env_reg, jni_env_cookie_offset, temp_reg, kLRTSegmentStateSize);
+  const MemberOffset previous_state_offset = JNIEnvExt::LrtPreviousStateOffset(kArmPointerSize);
+  const MemberOffset current_state_offset = JNIEnvExt::LrtSegmentStateOffset(kArmPointerSize);
+  DCHECK_EQ(previous_state_offset.SizeValue() + kLRTSegmentStateSize,
+            current_state_offset.SizeValue());
+
+  ___ Ldrd(AsVIXLRegister(previous_state_reg.AsArm()),
+           AsVIXLRegister(current_state_reg.AsArm()),
+           MemOperand(AsVIXLRegister(jni_env_reg.AsArm()), previous_state_offset.Int32Value()));
 }
 
-void ArmVIXLJNIMacroAssembler::PopLocalReferenceFrame(ManagedRegister jni_env_reg,
-                                                      ManagedRegister saved_cookie_reg,
-                                                      ManagedRegister temp_reg) {
+void ArmVIXLJNIMacroAssembler::StoreLocalReferenceTableStates(ManagedRegister jni_env_reg,
+                                                              ManagedRegister previous_state_reg,
+                                                              ManagedRegister current_state_reg) {
   constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
   DCHECK_EQ(kLRTSegmentStateSize, kRegSizeInBytes);
-  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kArmPointerSize);
-  const MemberOffset jni_env_segment_state_offset = JNIEnvExt::SegmentStateOffset(kArmPointerSize);
-  DCHECK_EQ(jni_env_cookie_offset.SizeValue() + kLRTSegmentStateSize,
-            jni_env_segment_state_offset.SizeValue());
-
-  // Load the current cookie.
-  Load(temp_reg, jni_env_reg, jni_env_cookie_offset, kLRTSegmentStateSize);
-
-  // Set the current segment state together with restoring the cookie.
-  ___ Strd(AsVIXLRegister(saved_cookie_reg.AsArm()),
-           AsVIXLRegister(temp_reg.AsArm()),
-           MemOperand(AsVIXLRegister(jni_env_reg.AsArm()), jni_env_cookie_offset.Int32Value()));
+  const MemberOffset previous_state_offset = JNIEnvExt::LrtPreviousStateOffset(kArmPointerSize);
+  const MemberOffset current_state_offset = JNIEnvExt::LrtSegmentStateOffset(kArmPointerSize);
+  DCHECK_EQ(previous_state_offset.SizeValue() + kLRTSegmentStateSize,
+            current_state_offset.SizeValue());
+
+  ___ Strd(AsVIXLRegister(previous_state_reg.AsArm()),
+           AsVIXLRegister(current_state_reg.AsArm()),
+           MemOperand(AsVIXLRegister(jni_env_reg.AsArm()), previous_state_offset.Int32Value()));
 }
 
 }  // namespace arm
diff --git a/compiler/utils/arm/jni_macro_assembler_arm_vixl.h b/compiler/utils/arm/jni_macro_assembler_arm_vixl.h
index 9d6b092105..3f6512c2a8 100644
--- a/compiler/utils/arm/jni_macro_assembler_arm_vixl.h
+++ b/compiler/utils/arm/jni_macro_assembler_arm_vixl.h
@@ -91,13 +91,13 @@ class ArmVIXLJNIMacroAssembler final
   void GetCurrentThread(ManagedRegister dest) override;
   void GetCurrentThread(FrameOffset dest_offset) override;
 
-  // Manipulating local reference frames.
-  void PushLocalReferenceFrame(ManagedRegister jni_env_reg,
-                               ManagedRegister saved_cookie_reg,
-                               ManagedRegister temp_reg) override;
-  void PopLocalReferenceFrame(ManagedRegister jni_env_reg,
-                              ManagedRegister saved_cookie_reg,
-                              ManagedRegister temp_reg) override;
+  // Manipulating local reference table states.
+  void LoadLocalReferenceTableStates(ManagedRegister jni_env_reg,
+                                     ManagedRegister previous_state_reg,
+                                     ManagedRegister current_state_reg) override;
+  void StoreLocalReferenceTableStates(ManagedRegister jni_env_reg,
+                                      ManagedRegister previous_state_reg,
+                                      ManagedRegister current_state_reg) override;
 
   // Decode JNI transition or local `jobject`. For (weak) global `jobject`, jump to slow path.
   void DecodeJNITransitionOrLocalJObject(ManagedRegister reg,
diff --git a/compiler/utils/arm64/jni_macro_assembler_arm64.cc b/compiler/utils/arm64/jni_macro_assembler_arm64.cc
index 3e7091b60b..2d1c003981 100644
--- a/compiler/utils/arm64/jni_macro_assembler_arm64.cc
+++ b/compiler/utils/arm64/jni_macro_assembler_arm64.cc
@@ -981,46 +981,37 @@ void Arm64JNIMacroAssembler::RemoveFrame(size_t frame_size,
   cfi().DefCFAOffset(frame_size);
 }
 
-void Arm64JNIMacroAssembler::PushLocalReferenceFrame(ManagedRegister jni_env_reg,
-                                                     ManagedRegister saved_cookie_reg,
-                                                     ManagedRegister temp_reg) {
+void Arm64JNIMacroAssembler::LoadLocalReferenceTableStates(ManagedRegister jni_env_reg,
+                                                           ManagedRegister previous_state_reg,
+                                                           ManagedRegister current_state_reg) {
   constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
   DCHECK_EQ(kLRTSegmentStateSize, kWRegSizeInBytes);
-  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kArm64PointerSize);
-  const MemberOffset jni_env_segment_state_offset =
-      JNIEnvExt::SegmentStateOffset(kArm64PointerSize);
-  DCHECK_EQ(jni_env_cookie_offset.SizeValue() + kLRTSegmentStateSize,
-            jni_env_segment_state_offset.SizeValue());
+  const MemberOffset previous_state_offset = JNIEnvExt::LrtPreviousStateOffset(kArm64PointerSize);
+  const MemberOffset current_state_offset = JNIEnvExt::LrtSegmentStateOffset(kArm64PointerSize);
+  DCHECK_EQ(previous_state_offset.SizeValue() + kLRTSegmentStateSize,
+            current_state_offset.SizeValue());
 
-  // Load the old cookie that we shall need to restore together with the current segment state.
   ___ Ldp(
-      reg_w(saved_cookie_reg.AsArm64().AsWRegister()),
-      reg_w(temp_reg.AsArm64().AsWRegister()),
-      MemOperand(reg_x(jni_env_reg.AsArm64().AsXRegister()), jni_env_cookie_offset.Int32Value()));
-
-  // Set the cookie in JNI environment to the current segment state.
-  Store(jni_env_reg, jni_env_cookie_offset, temp_reg, kLRTSegmentStateSize);
+      reg_w(previous_state_reg.AsArm64().AsWRegister()),
+      reg_w(current_state_reg.AsArm64().AsWRegister()),
+      MemOperand(reg_x(jni_env_reg.AsArm64().AsXRegister()), previous_state_offset.Int32Value()));
 }
 
-void Arm64JNIMacroAssembler::PopLocalReferenceFrame(ManagedRegister jni_env_reg,
-                                                    ManagedRegister saved_cookie_reg,
-                                                    ManagedRegister temp_reg) {
+void Arm64JNIMacroAssembler::StoreLocalReferenceTableStates(ManagedRegister jni_env_reg,
+                                                            ManagedRegister previous_state_reg,
+                                                            ManagedRegister current_state_reg) {
   constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
   DCHECK_EQ(kLRTSegmentStateSize, kWRegSizeInBytes);
-  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kArm64PointerSize);
-  const MemberOffset jni_env_segment_state_offset =
-      JNIEnvExt::SegmentStateOffset(kArm64PointerSize);
-  DCHECK_EQ(jni_env_cookie_offset.SizeValue() + kLRTSegmentStateSize,
-            jni_env_segment_state_offset.SizeValue());
-
-  // Load the current cookie.
-  Load(temp_reg, jni_env_reg, jni_env_cookie_offset, kLRTSegmentStateSize);
+  const MemberOffset previous_state_offset = JNIEnvExt::LrtPreviousStateOffset(kArm64PointerSize);
+  const MemberOffset current_state_offset = JNIEnvExt::LrtSegmentStateOffset(kArm64PointerSize);
+  DCHECK_EQ(previous_state_offset.SizeValue() + kLRTSegmentStateSize,
+            current_state_offset.SizeValue());
 
   // Set the current segment state together with restoring the cookie.
   ___ Stp(
-      reg_w(saved_cookie_reg.AsArm64().AsWRegister()),
-      reg_w(temp_reg.AsArm64().AsWRegister()),
-      MemOperand(reg_x(jni_env_reg.AsArm64().AsXRegister()), jni_env_cookie_offset.Int32Value()));
+      reg_w(previous_state_reg.AsArm64().AsWRegister()),
+      reg_w(current_state_reg.AsArm64().AsWRegister()),
+      MemOperand(reg_x(jni_env_reg.AsArm64().AsXRegister()), previous_state_offset.Int32Value()));
 }
 
 #undef ___
diff --git a/compiler/utils/arm64/jni_macro_assembler_arm64.h b/compiler/utils/arm64/jni_macro_assembler_arm64.h
index 1482012688..0750c2655f 100644
--- a/compiler/utils/arm64/jni_macro_assembler_arm64.h
+++ b/compiler/utils/arm64/jni_macro_assembler_arm64.h
@@ -93,13 +93,13 @@ class Arm64JNIMacroAssembler final : public JNIMacroAssemblerFwd<Arm64Assembler,
   void GetCurrentThread(ManagedRegister dest) override;
   void GetCurrentThread(FrameOffset dest_offset) override;
 
-  // Manipulating local reference frames.
-  void PushLocalReferenceFrame(ManagedRegister jni_env_reg,
-                               ManagedRegister saved_cookie_reg,
-                               ManagedRegister temp_reg) override;
-  void PopLocalReferenceFrame(ManagedRegister jni_env_reg,
-                              ManagedRegister saved_cookie_reg,
-                              ManagedRegister temp_reg) override;
+  // Manipulating local reference table states.
+  void LoadLocalReferenceTableStates(ManagedRegister jni_env_reg,
+                                     ManagedRegister previous_state_reg,
+                                     ManagedRegister current_state_reg) override;
+  void StoreLocalReferenceTableStates(ManagedRegister jni_env_reg,
+                                      ManagedRegister previous_state_reg,
+                                      ManagedRegister current_state_reg) override;
 
   // Decode JNI transition or local `jobject`. For (weak) global `jobject`, jump to slow path.
   void DecodeJNITransitionOrLocalJObject(ManagedRegister reg,
diff --git a/compiler/utils/jni_macro_assembler.cc b/compiler/utils/jni_macro_assembler.cc
index e17d915ad7..1806180980 100644
--- a/compiler/utils/jni_macro_assembler.cc
+++ b/compiler/utils/jni_macro_assembler.cc
@@ -132,53 +132,51 @@ void JNIMacroAssembler<PointerSize::k64>::LoadStackReference(ManagedRegister des
                                                              FrameOffset offs);
 
 template <PointerSize kPointerSize>
-void JNIMacroAssembler<kPointerSize>::PushLocalReferenceFrame(ManagedRegister jni_env_reg,
-                                                              ManagedRegister saved_cookie_reg,
-                                                              ManagedRegister temp_reg) {
+void JNIMacroAssembler<kPointerSize>::LoadLocalReferenceTableStates(
+    ManagedRegister jni_env_reg,
+    ManagedRegister previous_state_reg,
+    ManagedRegister current_state_reg) {
   constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
-  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kPointerSize);
-  const MemberOffset jni_env_segment_state_offset = JNIEnvExt::SegmentStateOffset(kPointerSize);
+  const MemberOffset previous_state_offset = JNIEnvExt::LrtPreviousStateOffset(kPointerSize);
+  const MemberOffset current_state_offset = JNIEnvExt::LrtSegmentStateOffset(kPointerSize);
 
-  // Load the old cookie that we shall need to restore.
-  Load(saved_cookie_reg, jni_env_reg, jni_env_cookie_offset, kLRTSegmentStateSize);
-
-  // Set the cookie to the current segment state.
-  Load(temp_reg, jni_env_reg, jni_env_segment_state_offset, kLRTSegmentStateSize);
-  Store(jni_env_reg, jni_env_cookie_offset, temp_reg, kLRTSegmentStateSize);
+  Load(previous_state_reg, jni_env_reg, previous_state_offset, kLRTSegmentStateSize);
+  Load(current_state_reg, jni_env_reg, current_state_offset, kLRTSegmentStateSize);
 }
 
 template
-void JNIMacroAssembler<PointerSize::k32>::PushLocalReferenceFrame(ManagedRegister jni_env_reg,
-                                                                  ManagedRegister saved_cookie_reg,
-                                                                  ManagedRegister temp_reg);
+void JNIMacroAssembler<PointerSize::k32>::LoadLocalReferenceTableStates(
+    ManagedRegister jni_env_reg,
+    ManagedRegister previous_state_reg,
+    ManagedRegister current_state_reg);
 template
-void JNIMacroAssembler<PointerSize::k64>::PushLocalReferenceFrame(ManagedRegister jni_env_reg,
-                                                                  ManagedRegister saved_cookie_reg,
-                                                                  ManagedRegister temp_reg);
+void JNIMacroAssembler<PointerSize::k64>::LoadLocalReferenceTableStates(
+    ManagedRegister jni_env_reg,
+    ManagedRegister previous_state_reg,
+    ManagedRegister current_state_reg);
 
 template <PointerSize kPointerSize>
-void JNIMacroAssembler<kPointerSize>::PopLocalReferenceFrame(ManagedRegister jni_env_reg,
-                                                             ManagedRegister saved_cookie_reg,
-                                                             ManagedRegister temp_reg) {
+void JNIMacroAssembler<kPointerSize>::StoreLocalReferenceTableStates(
+    ManagedRegister jni_env_reg,
+    ManagedRegister previous_state_reg,
+    ManagedRegister current_state_reg) {
   constexpr size_t kLRTSegmentStateSize = sizeof(jni::LRTSegmentState);
-  const MemberOffset jni_env_cookie_offset = JNIEnvExt::LocalRefCookieOffset(kPointerSize);
-  const MemberOffset jni_env_segment_state_offset = JNIEnvExt::SegmentStateOffset(kPointerSize);
-
-  // Set the current segment state to the current cookie.
-  Load(temp_reg, jni_env_reg, jni_env_cookie_offset, kLRTSegmentStateSize);
-  Store(jni_env_reg, jni_env_segment_state_offset, temp_reg, kLRTSegmentStateSize);
+  const MemberOffset previous_state_offset = JNIEnvExt::LrtPreviousStateOffset(kPointerSize);
+  const MemberOffset segment_state_offset = JNIEnvExt::LrtSegmentStateOffset(kPointerSize);
 
-  // Restore the cookie to the saved value.
-  Store(jni_env_reg, jni_env_cookie_offset, saved_cookie_reg, kLRTSegmentStateSize);
+  Store(jni_env_reg, previous_state_offset, previous_state_reg, kLRTSegmentStateSize);
+  Store(jni_env_reg, segment_state_offset, current_state_reg, kLRTSegmentStateSize);
 }
 
 template
-void JNIMacroAssembler<PointerSize::k32>::PopLocalReferenceFrame(ManagedRegister jni_env_reg,
-                                                                 ManagedRegister saved_cookie_reg,
-                                                                 ManagedRegister temp_reg);
+void JNIMacroAssembler<PointerSize::k32>::StoreLocalReferenceTableStates(
+    ManagedRegister jni_env_reg,
+    ManagedRegister previous_state_reg,
+    ManagedRegister current_state_reg);
 template
-void JNIMacroAssembler<PointerSize::k64>::PopLocalReferenceFrame(ManagedRegister jni_env_reg,
-                                                                 ManagedRegister saved_cookie_reg,
-                                                                 ManagedRegister temp_reg);
+void JNIMacroAssembler<PointerSize::k64>::StoreLocalReferenceTableStates(
+    ManagedRegister jni_env_reg,
+    ManagedRegister previous_state_reg,
+    ManagedRegister current_state_reg);
 
 }  // namespace art
diff --git a/compiler/utils/jni_macro_assembler.h b/compiler/utils/jni_macro_assembler.h
index a415141fbe..9d32071285 100644
--- a/compiler/utils/jni_macro_assembler.h
+++ b/compiler/utils/jni_macro_assembler.h
@@ -167,16 +167,16 @@ class JNIMacroAssembler : public DeletableArenaObject<kArenaAllocAssembler> {
   virtual void GetCurrentThread(ManagedRegister dest) = 0;
   virtual void GetCurrentThread(FrameOffset dest_offset) = 0;
 
-  // Manipulating local reference frames.
+  // Manipulating local reference table states.
   //
   // These have a default implementation but they can be overridden to use register pair
   // load/store instructions on architectures that support them (arm, arm64).
-  virtual void PushLocalReferenceFrame(ManagedRegister jni_env_reg,
-                                       ManagedRegister saved_cookie_reg,
-                                       ManagedRegister temp_reg);
-  virtual void PopLocalReferenceFrame(ManagedRegister jni_env_reg,
-                                      ManagedRegister saved_cookie_reg,
-                                      ManagedRegister temp_reg);
+  virtual void LoadLocalReferenceTableStates(ManagedRegister jni_env_reg,
+                                             ManagedRegister previous_state_reg,
+                                             ManagedRegister current_state_reg);
+  virtual void StoreLocalReferenceTableStates(ManagedRegister jni_env_reg,
+                                              ManagedRegister previous_state_reg,
+                                              ManagedRegister current_state_reg);
 
   // Decode JNI transition or local `jobject`. For (weak) global `jobject`, jump to slow path.
   virtual void DecodeJNITransitionOrLocalJObject(ManagedRegister reg,
diff --git a/runtime/jni/jni_env_ext.cc b/runtime/jni/jni_env_ext.cc
index 30f2c1e604..3b10ce29c9 100644
--- a/runtime/jni/jni_env_ext.cc
+++ b/runtime/jni/jni_env_ext.cc
@@ -152,12 +152,12 @@ inline MemberOffset JNIEnvExt::LocalReferenceTableOffset(PointerSize pointer_siz
                       2 * static_cast<size_t>(pointer_size));  // Thread* self + JavaVMExt* vm
 }
 
-MemberOffset JNIEnvExt::SegmentStateOffset(PointerSize pointer_size) {
+MemberOffset JNIEnvExt::LrtSegmentStateOffset(PointerSize pointer_size) {
   return MemberOffset(LocalReferenceTableOffset(pointer_size).SizeValue() +
                       jni::LocalReferenceTable::SegmentStateOffset().SizeValue());
 }
 
-MemberOffset JNIEnvExt::LocalRefCookieOffset(PointerSize pointer_size) {
+MemberOffset JNIEnvExt::LrtPreviousStateOffset(PointerSize pointer_size) {
   return MemberOffset(LocalReferenceTableOffset(pointer_size).SizeValue() +
                       jni::LocalReferenceTable::PreviousStateOffset().SizeValue());
 }
diff --git a/runtime/jni/jni_env_ext.h b/runtime/jni/jni_env_ext.h
index 7172acc7b3..ca5cc0a7ba 100644
--- a/runtime/jni/jni_env_ext.h
+++ b/runtime/jni/jni_env_ext.h
@@ -42,8 +42,8 @@ class JNIEnvExt : public JNIEnv {
   // Creates a new JNIEnvExt. Returns null on error, in which case error_msg
   // will contain a description of the error.
   static JNIEnvExt* Create(Thread* self, JavaVMExt* vm, std::string* error_msg);
-  static MemberOffset SegmentStateOffset(PointerSize pointer_size);
-  static MemberOffset LocalRefCookieOffset(PointerSize pointer_size);
+  static MemberOffset LrtSegmentStateOffset(PointerSize pointer_size);
+  static MemberOffset LrtPreviousStateOffset(PointerSize pointer_size);
   static MemberOffset SelfOffset(PointerSize pointer_size);
   static jint GetEnvHandler(JavaVMExt* vm, /*out*/void** out, jint version);
 
diff --git a/runtime/jni/jni_internal_test.cc b/runtime/jni/jni_internal_test.cc
index 1c878b079a..ed97e4d4c8 100644
--- a/runtime/jni/jni_internal_test.cc
+++ b/runtime/jni/jni_internal_test.cc
@@ -2643,13 +2643,13 @@ TEST_F(JniInternalTest, JNIEnvExtOffsets) {
       OFFSETOF_MEMBER(JNIEnvExt, locals_) +
       jni::LocalReferenceTable::PreviousStateOffset().Uint32Value();
   uint32_t previous_state_computed =
-      JNIEnvExt::LocalRefCookieOffset(kRuntimePointerSize).Uint32Value();
+      JNIEnvExt::LrtPreviousStateOffset(kRuntimePointerSize).Uint32Value();
   EXPECT_EQ(previous_state_now, previous_state_computed);
   uint32_t segment_state_now =
       OFFSETOF_MEMBER(JNIEnvExt, locals_) +
       jni::LocalReferenceTable::SegmentStateOffset().Uint32Value();
   uint32_t segment_state_computed =
-      JNIEnvExt::SegmentStateOffset(kRuntimePointerSize).Uint32Value();
+      JNIEnvExt::LrtSegmentStateOffset(kRuntimePointerSize).Uint32Value();
   EXPECT_EQ(segment_state_now, segment_state_computed);
 }