riscv64: Fix heap poisoning.

Fix heap poisoning for stubs, nterp, codegen and some
intrinsics and disable a few intrinsics broken for the
heap poisoning configuration.

Note that this fixes heap poisoning only for the CC GC.
The heap poisoning is pobably still broken for CMC GC.

Test: Build with `ART_HEAP_POISONING=true`, then
      testrunner.py --target --64 --ndebug --optimizing --interpreter
Bug: 283082089
Change-Id: I386235f6379722e1b2355150ab7d14f2ffb0da85
diff --git a/compiler/optimizing/code_generator_riscv64.cc b/compiler/optimizing/code_generator_riscv64.cc
index 182c1d4..0c0b8a9 100644
--- a/compiler/optimizing/code_generator_riscv64.cc
+++ b/compiler/optimizing/code_generator_riscv64.cc
@@ -2884,6 +2884,12 @@
   locations->SetInAt(0, Location::RequiresRegister());
   locations->SetInAt(1, Location::RegisterOrConstant(instruction->InputAt(1)));
   locations->SetInAt(2, ValueLocationForStore(instruction->GetValue()));
+  if (kPoisonHeapReferences &&
+      instruction->GetComponentType() == DataType::Type::kReference &&
+      !locations->InAt(1).IsConstant() &&
+      !locations->InAt(2).IsConstant()) {
+    locations->AddTemp(Location::RequiresRegister());
+  }
 }
 
 void InstructionCodeGeneratorRISCV64::VisitArraySet(HArraySet* instruction) {
@@ -2972,7 +2978,11 @@
     Store(value, array, offset, value_type);
   } else {
     ScratchRegisterScope srs(GetAssembler());
-    XRegister tmp = srs.AllocateXRegister();
+    // Heap poisoning needs two scratch registers in `Store()`, except for null constants.
+    XRegister tmp =
+        (kPoisonHeapReferences && value_type == DataType::Type::kReference && !value.IsConstant())
+            ? locations->GetTemp(0).AsRegister<XRegister>()
+            : srs.AllocateXRegister();
     ShNAdd(tmp, index.AsRegister<XRegister>(), array, value_type);
     Store(value, tmp, data_offset, value_type);
   }
diff --git a/compiler/optimizing/intrinsics_riscv64.cc b/compiler/optimizing/intrinsics_riscv64.cc
index 7f99f91..8b85ea4 100644
--- a/compiler/optimizing/intrinsics_riscv64.cc
+++ b/compiler/optimizing/intrinsics_riscv64.cc
@@ -2022,6 +2022,9 @@
   locations->SetInAt(1, Location::RequiresRegister());
   locations->SetInAt(2, Location::RequiresRegister());
   locations->SetInAt(3, Location::RequiresRegister());
+  if (kPoisonHeapReferences && invoke->InputAt(3)->GetType() == DataType::Type::kReference) {
+    locations->AddTemp(Location::RequiresRegister());
+  }
 }
 
 static void GenUnsafePut(HInvoke* invoke,
@@ -2038,7 +2041,10 @@
     // We use a block to end the scratch scope before the write barrier, thus
     // freeing the temporary registers so they can be used in `MarkGCCard()`.
     ScratchRegisterScope srs(assembler);
-    XRegister address = srs.AllocateXRegister();
+    // Heap poisoning needs two scratch registers in `Store()`.
+    XRegister address = (kPoisonHeapReferences && type == DataType::Type::kReference)
+        ? locations->GetTemp(0).AsRegister<XRegister>()
+        : srs.AllocateXRegister();
     __ Add(address, base, offset);
     GenerateSet(codegen, order, value, address, /*offset=*/ 0, type);
   }
@@ -2411,6 +2417,11 @@
     return;
   }
 
+  // TODO(riscv64): Fix this intrinsic for heap poisoning configuration.
+  if (kPoisonHeapReferences) {
+    return;
+  }
+
   CreateUnsafeCASLocations(allocator_, invoke, codegen_);
   if (codegen_->EmitReadBarrier()) {
     DCHECK(kUseBakerReadBarrier);
@@ -2572,6 +2583,11 @@
 }
 
 void IntrinsicLocationsBuilderRISCV64::VisitJdkUnsafeGetAndSetReference(HInvoke* invoke) {
+  // TODO(riscv64): Fix this intrinsic for heap poisoning configuration.
+  if (kPoisonHeapReferences) {
+    return;
+  }
+
   CreateUnsafeGetAndUpdateLocations(allocator_, invoke, codegen_);
 }
 
@@ -3184,6 +3200,14 @@
   }
 
   CreateVarHandleCommonLocations(invoke, codegen);
+  if (kPoisonHeapReferences && invoke->GetLocations() != nullptr) {
+    LocationSummary* locations = invoke->GetLocations();
+    uint32_t value_index = invoke->GetNumberOfArguments() - 1;
+    DataType::Type value_type = GetDataTypeFromShorty(invoke, value_index);
+    if (value_type == DataType::Type::kReference && !locations->InAt(value_index).IsConstant()) {
+      locations->AddTemp(Location::RequiresRegister());
+    }
+  }
 }
 
 static void GenerateVarHandleSet(HInvoke* invoke,
@@ -3208,7 +3232,11 @@
 
   {
     ScratchRegisterScope srs(assembler);
-    XRegister address = srs.AllocateXRegister();
+    // Heap poisoning needs two scratch registers in `Store()`, except for null constants.
+    XRegister address =
+        (kPoisonHeapReferences && value_type == DataType::Type::kReference && !value.IsConstant())
+            ? invoke->GetLocations()->GetTemp(0).AsRegister<XRegister>()
+            : srs.AllocateXRegister();
     __ Add(address, target.object, target.offset);
 
     if (byte_swap) {
@@ -3302,6 +3330,11 @@
     return;
   }
 
+  // TODO(riscv64): Fix this intrinsic for heap poisoning configuration.
+  if (kPoisonHeapReferences && value_type == DataType::Type::kReference) {
+    return;
+  }
+
   LocationSummary* locations = CreateVarHandleCommonLocations(invoke, codegen);
   DCHECK_EQ(expected_index, 1u + GetExpectedVarHandleCoordinatesCount(invoke));
 
@@ -3716,6 +3749,11 @@
     return;
   }
 
+  // TODO(riscv64): Fix this intrinsic for heap poisoning configuration.
+  if (kPoisonHeapReferences && invoke->GetType() == DataType::Type::kReference) {
+    return;
+  }
+
   LocationSummary* locations = CreateVarHandleCommonLocations(invoke, codegen);
   uint32_t arg_index = invoke->GetNumberOfArguments() - 1;
   DCHECK_EQ(arg_index, 1u + GetExpectedVarHandleCoordinatesCount(invoke));
diff --git a/runtime/arch/riscv64/asm_support_riscv64.S b/runtime/arch/riscv64/asm_support_riscv64.S
index 7b3e36a..22acbfb 100644
--- a/runtime/arch/riscv64/asm_support_riscv64.S
+++ b/runtime/arch/riscv64/asm_support_riscv64.S
@@ -94,6 +94,7 @@
 .macro POISON_HEAP_REF ref
 #ifdef USE_HEAP_POISONING
     neg \ref, \ref
+    zext.w \ref, \ref
 #endif  // USE_HEAP_POISONING
 .endm
 
@@ -102,6 +103,7 @@
 .macro UNPOISON_HEAP_REF ref
 #ifdef USE_HEAP_POISONING
     neg \ref, \ref
+    zext.w \ref, \ref
 #endif  // USE_HEAP_POISONING
 .endm
 
diff --git a/runtime/arch/riscv64/quick_entrypoints_riscv64.S b/runtime/arch/riscv64/quick_entrypoints_riscv64.S
index 3c1b073..bf00278 100644
--- a/runtime/arch/riscv64/quick_entrypoints_riscv64.S
+++ b/runtime/arch/riscv64/quick_entrypoints_riscv64.S
@@ -1795,7 +1795,7 @@
     mv    a0, \temp0
     add   \temp0, \temp0, \temp1
     sd    \temp0, THREAD_LOCAL_POS_OFFSET(xSELF)            // Store new thread_local_pos.
-    POISON_HEAP_REF class
+    POISON_HEAP_REF \class
     sw    \class, MIRROR_OBJECT_CLASS_OFFSET(a0)            // Store the class pointer.
     sw    \count, MIRROR_ARRAY_LENGTH_OFFSET(a0)            // Store the array length.
 // new-array is special. The class is loaded and immediately goes to the Initialized state
diff --git a/runtime/interpreter/mterp/riscv64/array.S b/runtime/interpreter/mterp/riscv64/array.S
index 14ff93b..81171ea 100644
--- a/runtime/interpreter/mterp/riscv64/array.S
+++ b/runtime/interpreter/mterp/riscv64/array.S
@@ -150,6 +150,7 @@
    sh2add t0, t1, t0
    lwu a0, MIRROR_OBJECT_ARRAY_DATA_OFFSET(t0)
        // a0 := *(array obj + data offset + idx * elem_size)
+   UNPOISON_HEAP_REF a0
    TEST_IF_MARKING t1, .L${opcode}_mark
 .L${opcode}_mark_resume:
    srliw t1, xINST, 8  // t1 := AA
diff --git a/runtime/interpreter/mterp/riscv64/invoke.S b/runtime/interpreter/mterp/riscv64/invoke.S
index b04a6f7..f606f95 100644
--- a/runtime/interpreter/mterp/riscv64/invoke.S
+++ b/runtime/interpreter/mterp/riscv64/invoke.S
@@ -65,6 +65,7 @@
    // Note: null case handled by SEGV handler.
    lwu t0, MIRROR_OBJECT_CLASS_OFFSET(a1)
                                     // t0 := klass object (32-bit addr)
+   UNPOISON_HEAP_REF t0
    // Entry address = entry's byte offset in vtable + vtable's byte offset in klass object.
    sh3add a0, a0, t0                // a0 := entry's byte offset
    ld a0, MIRROR_CLASS_VTABLE_OFFSET_64(a0)
@@ -169,6 +170,7 @@
    // Note: null case handled by SEGV handler.
    lwu t0, MIRROR_OBJECT_CLASS_OFFSET(a1)
                           // t0 := klass object (32-bit addr)
+   UNPOISON_HEAP_REF t0
    slliw t1, a0, 30       // test LSB #0 and #1
    bltz t1, 3f            // LSB #1 is set; handle default method
    bgtz t1, 4f            // LSB #0 is set; handle object method
diff --git a/runtime/interpreter/mterp/riscv64/object.S b/runtime/interpreter/mterp/riscv64/object.S
index 75ddcfd..f2de309 100644
--- a/runtime/interpreter/mterp/riscv64/object.S
+++ b/runtime/interpreter/mterp/riscv64/object.S
@@ -125,6 +125,7 @@
    beqz t0, .L${opcode}_null
    add t0, a0, t0          // t0 := field addr
    lwu a0, (t0)            // a0 := object
+   UNPOISON_HEAP_REF a0
    TEST_IF_MARKING t1, .L${opcode}_mark
 .L${opcode}_mark_resume:
 
@@ -165,6 +166,7 @@
    fence rw, rw
    lwu a0, ($z0)              // a0 := object
    fence r, rw
+   UNPOISON_HEAP_REF a0
    TEST_IF_MARKING t1, .L${opcode}_volatile_mark
 .L${opcode}_volatile_mark_resume:
 
@@ -296,6 +298,7 @@
    GET_VREG_OBJECT t0, s7  // t0 := holder
    beqz t0, .L${opcode}_null
    add t1, a0, t0          // t1 := field addr
+   POISON_HEAP_REF s9      // Poisoning maps null to null for the null check in write barrier.
    sw s9, (t1)
 %  object_write_barrier(value="s9", holder="t0", z0="t1", z1="t2", uniq=f"{opcode}")
 
@@ -331,6 +334,7 @@
    beqz $z0, .L${opcode}_volatile_null
    sub $z1, $z0, a0              // z1 := field addr (holder - (-offset))
    // Ensure the volatile store is released.
+   POISON_HEAP_REF $value  // Poisoning maps null to null for the null check in write barrier.
    amoswap.w.rl zero, $value, ($z1)
 %  object_write_barrier(value=value, holder=z0, z0=z1, z1=z2, uniq=f"slow_{opcode}")
 
@@ -461,6 +465,7 @@
 
    add t0, t0, a0      // t0 := field addr
    lwu a0, (t0)        // a0 := value (ref)
+   UNPOISON_HEAP_REF a0
 
 .L${opcode}_mark_resume:
    FETCH_ADVANCE_INST 2
@@ -472,6 +477,7 @@
    call art_quick_read_barrier_mark_reg10  // a0, holder
    add t0, t0, a0      // t0 := field addr, after a0 update
    lwu a0, (t0)        // a0 := value (ref)
+   UNPOISON_HEAP_REF a0
    call art_quick_read_barrier_mark_reg10  // a0, object
    j .L${opcode}_mark_resume
 
@@ -507,6 +513,7 @@
    fence rw, rw
    lwu a0, ($z0)  // Atomic ref load: "fence rw,rw, ; LOAD ; fence r,rw"
    fence r, rw
+   UNPOISON_HEAP_REF a0
 
 .L${opcode}_volatile_mark_resume:
    FETCH_ADVANCE_INST 2
@@ -520,6 +527,7 @@
    fence rw, rw
    lwu a0, ($z0)  // Atomic ref load: "fence rw,rw, ; LOAD ; fence r,rw"
    fence r, rw
+   UNPOISON_HEAP_REF a0
    call art_quick_read_barrier_mark_reg10  // a0, object
    j .L${opcode}_volatile_mark_resume
 
@@ -652,6 +660,7 @@
 .L${opcode}_mark_resume:
 
    add t0, t0, a0          // t0 := field addr, after possible a0 update
+   POISON_HEAP_REF s8      // Poisoning maps null to null for the null check in write barrier.
    sw s8, (t0)             // store reference
 %  object_write_barrier(value="s8", holder="a0", z0="t0", z1="t1", uniq=f"{opcode}")
    FETCH_ADVANCE_INST 2
@@ -699,6 +708,7 @@
    // Ensure the volatile store is released.
    // \value must NOT be the destination register, the destination gets clobbered!
    // For refs, \value's original value is used in the write barrier below.
+   POISON_HEAP_REF $value  // Poisoning maps null to null for the null check in write barrier.
    amoswap.w.rl zero, $value, ($z0)
 %  object_write_barrier(value=value, holder="a0", z0=z0, z1=z1, uniq=f"slow_{opcode}")