riscv64: Fix wrong sign-extension for references.

Test: Modify kPreferredAllocSpaceBegin = 0x90000000, then
      testrunner.py --target --64 --ndebug --optimizing
Bug: 283082089
Change-Id: Ifb82d616a0d9664a2e7f5f96a1a79ddce5862cdf

7 files changed, 46 insertions, 8 deletions

diff --git a/compiler/jni/quick/jni_compiler.cc b/compiler/jni/quick/jni_compiler.cc
index b125d2ef7c..3357a5f8d7 100644
--- a/compiler/jni/quick/jni_compiler.cc
+++ b/compiler/jni/quick/jni_compiler.cc
@@ -550,10 +550,10 @@ static JniCompiledMethod ArtJniCompileMethodInternal(const CompilerOptions& comp
       FrameOffset method_offset = mr_conv->MethodStackOffset();
       __ Load(temp, method_offset, kRawPointerSize);
       DCHECK_EQ(ArtMethod::DeclaringClassOffset().SizeValue(), 0u);
-      __ Load(to_lock, temp, MemberOffset(0u), kObjectReferenceSize);
+      __ LoadGcRootWithoutReadBarrier(to_lock, temp, MemberOffset(0u));
     } else {
       // Pass the `this` argument from its spill slot.
-      __ Load(to_lock, mr_conv->CurrentParamStackOffset(), kObjectReferenceSize);
+      __ LoadStackReference(to_lock, mr_conv->CurrentParamStackOffset());
     }
     __ CallFromThread(QUICK_ENTRYPOINT_OFFSET(kPointerSize, pJniUnlockObject));
   }
diff --git a/compiler/utils/jni_macro_assembler.cc b/compiler/utils/jni_macro_assembler.cc
index dc7ec60032..7a90a46f51 100644
--- a/compiler/utils/jni_macro_assembler.cc
+++ b/compiler/utils/jni_macro_assembler.cc
@@ -38,6 +38,7 @@
 #include "base/globals.h"
 #include "base/memory_region.h"
 #include "gc_root.h"
+#include "stack_reference.h"
 
 namespace art HIDDEN {
 
@@ -115,4 +116,17 @@ void JNIMacroAssembler<PointerSize::k64>::LoadGcRootWithoutReadBarrier(ManagedRe
                                                                        ManagedRegister base,
                                                                        MemberOffset offs);
 
+template <PointerSize kPointerSize>
+void JNIMacroAssembler<kPointerSize>::LoadStackReference(ManagedRegister dest, FrameOffset offs) {
+  static_assert(sizeof(uint32_t) == sizeof(StackReference<mirror::Object>));
+  Load(dest, offs, sizeof(uint32_t));
+}
+
+template
+void JNIMacroAssembler<PointerSize::k32>::LoadStackReference(ManagedRegister dest,
+                                                             FrameOffset offs);
+template
+void JNIMacroAssembler<PointerSize::k64>::LoadStackReference(ManagedRegister dest,
+                                                             FrameOffset offs);
+
 }  // namespace art
diff --git a/compiler/utils/jni_macro_assembler.h b/compiler/utils/jni_macro_assembler.h
index 2d51439ee8..2d52eada08 100644
--- a/compiler/utils/jni_macro_assembler.h
+++ b/compiler/utils/jni_macro_assembler.h
@@ -137,6 +137,10 @@ class JNIMacroAssembler : public DeletableArenaObject<kArenaAllocAssembler> {
                                             ManagedRegister base,
                                             MemberOffset offs);
 
+  // Load reference from a `StackReference<>`. The default is to load as `jint`. Some architectures
+  // (say, RISC-V) override this to provide a different sign- or zero-extension.
+  virtual void LoadStackReference(ManagedRegister dest, FrameOffset offs);
+
   // Copying routines
 
   // Move arguments from `srcs` locations to `dests` locations.
diff --git a/compiler/utils/riscv64/jni_macro_assembler_riscv64.cc b/compiler/utils/riscv64/jni_macro_assembler_riscv64.cc
index 3aeee8a154..9d3a29d252 100644
--- a/compiler/utils/riscv64/jni_macro_assembler_riscv64.cc
+++ b/compiler/utils/riscv64/jni_macro_assembler_riscv64.cc
@@ -19,10 +19,12 @@
 #include "base/bit_utils_iterator.h"
 #include "dwarf/register.h"
 #include "entrypoints/quick/quick_entrypoints.h"
+#include "gc_root.h"
 #include "indirect_reference_table.h"
 #include "lock_word.h"
 #include "managed_register_riscv64.h"
 #include "offsets.h"
+#include "stack_reference.h"
 #include "thread.h"
 
 namespace art HIDDEN {
@@ -242,9 +244,20 @@ void Riscv64JNIMacroAssembler::LoadGcRootWithoutReadBarrier(ManagedRegister m_de
                                                             MemberOffset offs) {
   Riscv64ManagedRegister base = m_base.AsRiscv64();
   Riscv64ManagedRegister dest = m_dest.AsRiscv64();
+  static_assert(sizeof(uint32_t) == sizeof(GcRoot<mirror::Object>));
   __ Loadwu(dest.AsXRegister(), base.AsXRegister(), offs.Int32Value());
 }
 
+void Riscv64JNIMacroAssembler::LoadStackReference(ManagedRegister m_dest, FrameOffset offs) {
+  // `StackReference<>` and `GcRoot<>` have the same underlying representation, namely
+  // `CompressedReference<>`. And `StackReference<>` does not need a read barrier.
+  static_assert(sizeof(uint32_t) == sizeof(mirror::CompressedReference<mirror::Object>));
+  static_assert(sizeof(uint32_t) == sizeof(StackReference<mirror::Object>));
+  static_assert(sizeof(uint32_t) == sizeof(GcRoot<mirror::Object>));
+  LoadGcRootWithoutReadBarrier(
+      m_dest, Riscv64ManagedRegister::FromXRegister(SP), MemberOffset(offs.Int32Value()));
+}
+
 void Riscv64JNIMacroAssembler::MoveArguments(ArrayRef<ArgumentLocation> dests,
                                              ArrayRef<ArgumentLocation> srcs,
                                              ArrayRef<FrameOffset> refs) {
@@ -407,7 +420,7 @@ void Riscv64JNIMacroAssembler::DecodeJNITransitionOrLocalJObject(ManagedRegister
   __ Andi(TMP, reg, kGlobalOrWeakGlobalMask);
   __ Bnez(TMP, Riscv64JNIMacroLabel::Cast(slow_path)->AsRiscv64());
   __ Andi(reg, reg, ~kIndirectRefKindMask);
-  __ Loadw(reg, reg, 0);
+  __ Loadwu(reg, reg, 0);
 }
 
 void Riscv64JNIMacroAssembler::VerifyObject([[maybe_unused]] ManagedRegister m_src,
diff --git a/compiler/utils/riscv64/jni_macro_assembler_riscv64.h b/compiler/utils/riscv64/jni_macro_assembler_riscv64.h
index 79618e2c8e..3cbed0d53b 100644
--- a/compiler/utils/riscv64/jni_macro_assembler_riscv64.h
+++ b/compiler/utils/riscv64/jni_macro_assembler_riscv64.h
@@ -71,6 +71,7 @@ class Riscv64JNIMacroAssembler  : public JNIMacroAssemblerFwd<Riscv64Assembler,
   void LoadGcRootWithoutReadBarrier(ManagedRegister dest,
                                     ManagedRegister base,
                                     MemberOffset offs) override;
+  void LoadStackReference(ManagedRegister dest, FrameOffset offs) override;
 
   // Copying routines.
   void MoveArguments(ArrayRef<ArgumentLocation> dests,
diff --git a/compiler/utils/riscv64/jni_macro_assembler_riscv64_test.cc b/compiler/utils/riscv64/jni_macro_assembler_riscv64_test.cc
index 004ba9bb7f..be6feeb9de 100644
--- a/compiler/utils/riscv64/jni_macro_assembler_riscv64_test.cc
+++ b/compiler/utils/riscv64/jni_macro_assembler_riscv64_test.cc
@@ -272,6 +272,12 @@ TEST_F(JniMacroAssemblerRiscv64Test, Load) {
   expected += "addi t6, s2, 0x7f8\n"
               "lwu t1, 8(t6)\n";
 
+  __ LoadStackReference(AsManaged(T0), FrameOffset(0));
+  expected += "lwu t0, 0(sp)\n";
+  __ LoadStackReference(AsManaged(T1), FrameOffset(0x800));
+  expected += "addi t6, sp, 0x7f8\n"
+              "lwu t1, 8(t6)\n";
+
   DriverStr(expected, "Load");
 }
 
@@ -693,7 +699,7 @@ TEST_F(JniMacroAssemblerRiscv64Test, DecodeJNITransitionOrLocalJObject) {
               "andi t6, a0, " + std::to_string(kGlobalOrWeakGlobalMask) + "\n"
               "bnez t6, 2f\n"
               "andi a0, a0, ~" + std::to_string(kIndirectRefKindMask) + "\n"
-              "lw a0, (a0)\n";
+              "lwu a0, (a0)\n";
 
   __ Bind(resume.get());
   expected += "1:\n";
diff --git a/runtime/arch/riscv64/quick_entrypoints_riscv64.S b/runtime/arch/riscv64/quick_entrypoints_riscv64.S
index d8d6ef29da..7d5443918e 100644
--- a/runtime/arch/riscv64/quick_entrypoints_riscv64.S
+++ b/runtime/arch/riscv64/quick_entrypoints_riscv64.S
@@ -168,7 +168,7 @@ END art_deliver_pending_exception
 
     // Load this (if instance method) and record the number of GPRs to fill.
     .ifc \sfx, _instance
-        lw   a1, (t0)                  // Load "this" parameter,
+        lwu  a1, (t0)                  // Load "this" parameter,
         addi t0, t0, 4                 // and increment arg pointer.
         .equ NUM_GPRS_TO_FILL, 6
     .else
@@ -1389,9 +1389,9 @@ ENTRY \name
     and t6, t6, t5
     CFI_REMEMBER_STATE
     bgez t6, .Lrb_full_\name
-    // Note: The mark bit which is shifted to the sign bit and sign-extended is
-    // always zero in the forwarding address state. No zero-extension is needed.
-    slliw \reg, t5, LOCK_WORD_STATE_FORWARDING_ADDRESS_SHIFT
+    // Extract and zero-extend the forwarding address.
+    slli \reg, t5, (LOCK_WORD_STATE_FORWARDING_ADDRESS_SHIFT + 32)
+    srli \reg, \reg, 32
     .ifc \reg, t5
       sd t5, (8*0)(sp)
     .endif