Do HuntForOriginalReference for the value in WBE

This allows to better detect when we have a null value at compile time.

Bug: 260843353
Test: art/test/testrunner/testrunner.py --host --64 --optimizing -b
Change-Id: Idb2a37a52fb379ea6264be3258d5babd10cf1273

5 files changed, 22 insertions, 26 deletions

diff --git a/compiler/optimizing/code_generator_arm64.cc b/compiler/optimizing/code_generator_arm64.cc
index 8974429955..50b86489ee 100644
--- a/compiler/optimizing/code_generator_arm64.cc
+++ b/compiler/optimizing/code_generator_arm64.cc
@@ -2356,14 +2356,12 @@ void InstructionCodeGeneratorARM64::HandleFieldSet(HInstruction* instruction,
       codegen_->StoreNeedsWriteBarrier(field_type, instruction->InputAt(1), write_barrier_kind);
 
   if (needs_write_barrier) {
-    // TODO(solanes): If we do a `HuntForOriginalReference` call to the value in WBE, we will be
-    // able to DCHECK that the write_barrier_kind is kBeingReliedOn when Register(value).IsZero(),
-    // and we could remove the `!Register(value).IsZero()` from below.
-    codegen_->MaybeMarkGCCard(obj,
-                              Register(value),
-                              value_can_be_null &&
-                                  write_barrier_kind == WriteBarrierKind::kEmitNotBeingReliedOn &&
-                                  !Register(value).IsZero());
+    DCHECK_IMPLIES(Register(value).IsZero(),
+                   write_barrier_kind == WriteBarrierKind::kEmitBeingReliedOn);
+    codegen_->MaybeMarkGCCard(
+        obj,
+        Register(value),
+        value_can_be_null && write_barrier_kind == WriteBarrierKind::kEmitNotBeingReliedOn);
   } else if (codegen_->ShouldCheckGCCard(field_type, instruction->InputAt(1), write_barrier_kind)) {
     codegen_->CheckGCCardIsValid(obj);
   }
@@ -3052,9 +3050,8 @@ void InstructionCodeGeneratorARM64::VisitArraySet(HArraySet* instruction) {
     // write barrier when its value is null (without an extra cbz since we already checked if the
     // value is null for the type check). This will be done as a follow-up since it is a runtime
     // optimization that needs extra care.
-    // TODO(solanes): We can also skip it for known zero values which are not relied on i.e. when
-    // we have the Zero register as the value. If we do `HuntForOriginalReference` on the value
-    // we'll resolve this.
+    DCHECK_IMPLIES(Register(value).IsZero(),
+                   write_barrier_kind == WriteBarrierKind::kEmitBeingReliedOn);
     codegen_->MarkGCCard(array);
 
     UseScratchRegisterScope temps(masm);
diff --git a/compiler/optimizing/code_generator_riscv64.cc b/compiler/optimizing/code_generator_riscv64.cc
index 93bd35b618..abbd74ac65 100644
--- a/compiler/optimizing/code_generator_riscv64.cc
+++ b/compiler/optimizing/code_generator_riscv64.cc
@@ -2518,8 +2518,7 @@ void InstructionCodeGeneratorRISCV64::HandleFieldSet(HInstruction* instruction,
       codegen_->StoreNeedsWriteBarrier(type, instruction->InputAt(1), write_barrier_kind);
   if (needs_write_barrier) {
     if (value.IsConstant()) {
-      // TODO(solanes): If we do a `HuntForOriginalReference` call to the value in WBE, we will be
-      // able to DCHECK that the write_barrier_kind is kBeingReliedOn.
+      DCHECK_EQ(write_barrier_kind, WriteBarrierKind::kEmitBeingReliedOn);
       codegen_->MarkGCCard(obj);
     } else {
       codegen_->MaybeMarkGCCard(
@@ -2947,9 +2946,8 @@ void InstructionCodeGeneratorRISCV64::VisitArraySet(HArraySet* instruction) {
 
   if (needs_write_barrier) {
     DCHECK_EQ(value_type, DataType::Type::kReference);
+    DCHECK_IMPLIES(value.IsConstant(), value.GetConstant()->IsArithmeticZero());
     const bool storing_constant_zero = value.IsConstant();
-    // TODO(solanes): If we do a `HuntForOriginalReference` call to the value in WBE, we will be
-    // able to DCHECK that the write_barrier_kind is kBeingReliedOn when we have a constant.
     if (!storing_constant_zero) {
       Riscv64Label do_store;
 
@@ -3010,11 +3008,13 @@ void InstructionCodeGeneratorRISCV64::VisitArraySet(HArraySet* instruction) {
       }
     }
 
-    DCHECK_NE(instruction->GetWriteBarrierKind(), WriteBarrierKind::kDontEmit);
+    DCHECK_NE(write_barrier_kind, WriteBarrierKind::kDontEmit);
     // TODO(solanes): The WriteBarrierKind::kEmitNotBeingReliedOn case should be able to skip
     // this write barrier when its value is null (without an extra Beqz since we already checked
     // if the value is null for the type check). This will be done as a follow-up since it is a
     // runtime optimization that needs extra care.
+    DCHECK_IMPLIES(storing_constant_zero,
+                   write_barrier_kind == WriteBarrierKind::kEmitBeingReliedOn);
     codegen_->MarkGCCard(array);
   } else if (codegen_->ShouldCheckGCCard(value_type, instruction->GetValue(), write_barrier_kind)) {
     codegen_->CheckGCCardIsValid(array);
diff --git a/compiler/optimizing/graph_checker.cc b/compiler/optimizing/graph_checker.cc
index d60ec06097..c880d45d20 100644
--- a/compiler/optimizing/graph_checker.cc
+++ b/compiler/optimizing/graph_checker.cc
@@ -1291,7 +1291,7 @@ bool IsRemovedWriteBarrier(DataType::Type type,
                            HInstruction* value) {
   return write_barrier_kind == WriteBarrierKind::kDontEmit &&
          type == DataType::Type::kReference &&
-         !value->IsNullConstant();
+         !HuntForOriginalReference(value)->IsNullConstant();
 }
 
 void GraphChecker::VisitArraySet(HArraySet* instruction) {
diff --git a/compiler/optimizing/write_barrier_elimination.cc b/compiler/optimizing/write_barrier_elimination.cc
index 220eb78e02..537bc09f93 100644
--- a/compiler/optimizing/write_barrier_elimination.cc
+++ b/compiler/optimizing/write_barrier_elimination.cc
@@ -45,7 +45,7 @@ class WBEVisitor final : public HGraphVisitor {
     DCHECK(!instruction->GetSideEffects().Includes(SideEffects::CanTriggerGC()));
 
     if (instruction->GetFieldType() != DataType::Type::kReference ||
-        instruction->GetValue()->IsNullConstant()) {
+        HuntForOriginalReference(instruction->GetValue())->IsNullConstant()) {
       instruction->SetWriteBarrierKind(WriteBarrierKind::kDontEmit);
       return;
     }
@@ -72,7 +72,7 @@ class WBEVisitor final : public HGraphVisitor {
     DCHECK(!instruction->GetSideEffects().Includes(SideEffects::CanTriggerGC()));
 
     if (instruction->GetFieldType() != DataType::Type::kReference ||
-        instruction->GetValue()->IsNullConstant()) {
+        HuntForOriginalReference(instruction->GetValue())->IsNullConstant()) {
       instruction->SetWriteBarrierKind(WriteBarrierKind::kDontEmit);
       return;
     }
@@ -100,7 +100,7 @@ class WBEVisitor final : public HGraphVisitor {
     }
 
     if (instruction->GetComponentType() != DataType::Type::kReference ||
-        instruction->GetValue()->IsNullConstant()) {
+        HuntForOriginalReference(instruction->GetValue())->IsNullConstant()) {
       instruction->SetWriteBarrierKind(WriteBarrierKind::kDontEmit);
       return;
     }
diff --git a/test/2272-checker-codegen-honor-write-barrier-kind/src/Main.java b/test/2272-checker-codegen-honor-write-barrier-kind/src/Main.java
index f07286b5b9..7d67249c6d 100644
--- a/test/2272-checker-codegen-honor-write-barrier-kind/src/Main.java
+++ b/test/2272-checker-codegen-honor-write-barrier-kind/src/Main.java
@@ -42,21 +42,20 @@ public class Main {
     /// CHECK-START: java.lang.String[] Main.$noinline$testArraySetsHonorWriteBarrier(java.lang.String[], java.lang.String) prepare_for_register_allocation (before)
     /// CHECK: <<Null:l\d+>>   NullConstant
     /// CHECK: <<BT:l\d+>>     BoundType [<<Null>>]
-    /// CHECK: ArraySet [<<arr:l\d+>>,<<index:i\d+>>,<<BT>>] value_can_be_null:true needs_type_check:false can_trigger_gc:false write_barrier_kind:EmitBeingReliedOn
-    /// CHECK: ArraySet value_can_be_null:true needs_type_check:false can_trigger_gc:false write_barrier_kind:DontEmit
+    /// CHECK: ArraySet [<<arr:l\d+>>,<<index:i\d+>>,<<BT>>] value_can_be_null:true needs_type_check:false can_trigger_gc:false write_barrier_kind:DontEmit
+    /// CHECK: ArraySet value_can_be_null:true needs_type_check:false can_trigger_gc:false write_barrier_kind:EmitNotBeingReliedOn
 
     /// CHECK-START: java.lang.String[] Main.$noinline$testArraySetsHonorWriteBarrier(java.lang.String[], java.lang.String) prepare_for_register_allocation (after)
     /// CHECK: <<Null:l\d+>>   NullConstant
-    /// CHECK: ArraySet [<<arr:l\d+>>,<<index:i\d+>>,<<Null>>] value_can_be_null:true needs_type_check:false can_trigger_gc:false write_barrier_kind:EmitBeingReliedOn
-    /// CHECK: ArraySet value_can_be_null:true needs_type_check:false can_trigger_gc:false write_barrier_kind:DontEmit
+    /// CHECK: ArraySet [<<arr:l\d+>>,<<index:i\d+>>,<<Null>>] value_can_be_null:true needs_type_check:false can_trigger_gc:false write_barrier_kind:DontEmit
+    /// CHECK: ArraySet value_can_be_null:true needs_type_check:false can_trigger_gc:false write_barrier_kind:EmitNotBeingReliedOn
 
     /// CHECK-START: java.lang.String[] Main.$noinline$testArraySetsHonorWriteBarrier(java.lang.String[], java.lang.String) prepare_for_register_allocation (after)
     /// CHECK-NOT: BoundType
 
     /// CHECK-START: java.lang.String[] Main.$noinline$testArraySetsHonorWriteBarrier(java.lang.String[], java.lang.String) disassembly (after)
-    /// CHECK: ArraySet value_can_be_null:true needs_type_check:false can_trigger_gc:false write_barrier_kind:EmitBeingReliedOn
-    // / CHECK: ; card_table
     /// CHECK: ArraySet value_can_be_null:true needs_type_check:false can_trigger_gc:false write_barrier_kind:DontEmit
+    /// CHECK: ArraySet value_can_be_null:true needs_type_check:false can_trigger_gc:false write_barrier_kind:EmitNotBeingReliedOn
     private static java.lang.String[] $noinline$testArraySetsHonorWriteBarrier(
             String[] arr, String o2) {
         Object o = null;