1 files changed, 28 insertions, 2 deletions

diff --git a/compiler/optimizing/optimizing_unit_test.h b/compiler/optimizing/optimizing_unit_test.h
index c30292c1f0..f4f62e67ed 100644
--- a/compiler/optimizing/optimizing_unit_test.h
+++ b/compiler/optimizing/optimizing_unit_test.h
@@ -472,6 +472,7 @@ class OptimizingUnitTestHelper {
   HLoadClass* MakeLoadClass(HBasicBlock* block,
                             std::optional<dex::TypeIndex> ti = std::nullopt,
                             std::optional<Handle<mirror::Class>> klass = std::nullopt,
+                            std::initializer_list<HInstruction*> env = {},
                             uint32_t dex_pc = kNoDexPc) {
     HLoadClass* load_class = new (GetAllocator()) HLoadClass(
         graph_->GetCurrentMethod(),
@@ -482,10 +483,14 @@ class OptimizingUnitTestHelper {
         dex_pc,
         /* needs_access_check= */ false);
     AddOrInsertInstruction(block, load_class);
+    ManuallyBuildEnvFor(load_class, env);
     return load_class;
   }
 
-  HNewInstance* MakeNewInstance(HBasicBlock* block, HInstruction* cls, uint32_t dex_pc = kNoDexPc) {
+  HNewInstance* MakeNewInstance(HBasicBlock* block,
+                                HInstruction* cls,
+                                std::initializer_list<HInstruction*> env = {},
+                                uint32_t dex_pc = kNoDexPc) {
     EXPECT_TRUE(cls->IsLoadClass() || cls->IsClinitCheck()) << *cls;
     HLoadClass* load =
         cls->IsLoadClass() ? cls->AsLoadClass() : cls->AsClinitCheck()->GetLoadClass();
@@ -497,6 +502,7 @@ class OptimizingUnitTestHelper {
         /* finalizable= */ false,
         QuickEntrypointEnum::kQuickAllocObjectInitialized);
     AddOrInsertInstruction(block, new_instance);
+    ManuallyBuildEnvFor(new_instance, env);
     return new_instance;
   }
 
@@ -553,10 +559,12 @@ class OptimizingUnitTestHelper {
                           HInstruction* cls,
                           HInstruction* length,
                           size_t component_size_shift = DataType::SizeShift(DataType::Type::kInt32),
+                          std::initializer_list<HInstruction*> env = {},
                           uint32_t dex_pc = kNoDexPc) {
     HNewArray* new_array =
         new (GetAllocator()) HNewArray(cls, length, dex_pc, component_size_shift);
     AddOrInsertInstruction(block, new_array);
+    ManuallyBuildEnvFor(new_array, env);
     return new_array;
   }
 
@@ -564,6 +572,15 @@ class OptimizingUnitTestHelper {
                           HInstruction* array,
                           HInstruction* index,
                           HInstruction* value,
+                          uint32_t dex_pc = kNoDexPc) {
+    CHECK(value != nullptr);
+    return MakeArraySet(block, array, index, value, value->GetType(), dex_pc);
+  }
+
+  HArraySet* MakeArraySet(HBasicBlock* block,
+                          HInstruction* array,
+                          HInstruction* index,
+                          HInstruction* value,
                           DataType::Type type,
                           uint32_t dex_pc = kNoDexPc) {
     HArraySet* array_set = new (GetAllocator()) HArraySet(array, index, value, type, dex_pc);
@@ -591,18 +608,22 @@ class OptimizingUnitTestHelper {
 
   HNullCheck* MakeNullCheck(HBasicBlock* block,
                             HInstruction* value,
+                            std::initializer_list<HInstruction*> env = {},
                             uint32_t dex_pc = kNoDexPc) {
     HNullCheck* null_check = new (GetAllocator()) HNullCheck(value, dex_pc);
     AddOrInsertInstruction(block, null_check);
+    ManuallyBuildEnvFor(null_check, env);
     return null_check;
   }
 
   HBoundsCheck* MakeBoundsCheck(HBasicBlock* block,
                                 HInstruction* index,
                                 HInstruction* length,
+                                std::initializer_list<HInstruction*> env = {},
                                 uint32_t dex_pc = kNoDexPc) {
     HBoundsCheck* bounds_check = new (GetAllocator()) HBoundsCheck(index, length, dex_pc);
     AddOrInsertInstruction(block, bounds_check);
+    ManuallyBuildEnvFor(bounds_check, env);
     return bounds_check;
   }
 
@@ -623,6 +644,7 @@ class OptimizingUnitTestHelper {
   HInvokeStaticOrDirect* MakeInvokeStatic(HBasicBlock* block,
                                           DataType::Type return_type,
                                           const std::vector<HInstruction*>& args,
+                                          std::initializer_list<HInstruction*> env = {},
                                           uint32_t dex_pc = kNoDexPc) {
     MethodReference method_reference{/* file= */ &graph_->GetDexFile(), /* index= */ method_idx_++};
     HInvokeStaticOrDirect* invoke = new (GetAllocator())
@@ -641,6 +663,7 @@ class OptimizingUnitTestHelper {
       invoke->SetRawInputAt(idx, ins);
     }
     AddOrInsertInstruction(block, invoke);
+    ManuallyBuildEnvFor(invoke, env);
     return invoke;
   }
 
@@ -676,9 +699,12 @@ class OptimizingUnitTestHelper {
     return select;
   }
 
-  HSuspendCheck* MakeSuspendCheck(HBasicBlock* block, uint32_t dex_pc = kNoDexPc) {
+  HSuspendCheck* MakeSuspendCheck(HBasicBlock* block,
+                                  std::initializer_list<HInstruction*> env = {},
+                                  uint32_t dex_pc = kNoDexPc) {
     HSuspendCheck* suspend_check = new (GetAllocator()) HSuspendCheck(dex_pc);
     AddOrInsertInstruction(block, suspend_check);
+    ManuallyBuildEnvFor(suspend_check, env);
     return suspend_check;
   }