Fix test to search for a populated StackMap.

With the baseline compiler, the test could have hit a stack map for a
NullCheck, where we don't populate the reference StackMask in the stack
map. To fix this, look for another stack map starting at the same dex
pc.

Test: 004-ReferenceMap

Change-Id: I2628ded5c89c6fb0059a18341bffda78ef06f197

3 files changed, 38 insertions, 15 deletions

diff --git a/runtime/check_reference_map_visitor.h b/runtime/check_reference_map_visitor.h
index f4bda90fc1..a7c3e45ae7 100644
--- a/runtime/check_reference_map_visitor.h
+++ b/runtime/check_reference_map_visitor.h
@@ -59,18 +59,35 @@ class CheckReferenceMapVisitor : public StackVisitor {
     return false;
   }
 
-  void CheckReferences(int* registers, int number_of_references, uint32_t native_pc_offset)
+  void CheckReferences(int* registers,
+                       int number_of_references,
+                       uint32_t dex_pc,
+                       uint32_t native_pc_offset,
+                       bool search_for_valid_stack_map)
       REQUIRES_SHARED(Locks::mutator_lock_) {
     CHECK(GetCurrentOatQuickMethodHeader()->IsOptimized());
-    CheckOptimizedMethod(registers, number_of_references, native_pc_offset);
+    CheckOptimizedMethod(
+        registers, number_of_references, dex_pc, native_pc_offset, search_for_valid_stack_map);
   }
 
  private:
-  void CheckOptimizedMethod(int* registers, int number_of_references, uint32_t native_pc_offset)
+  void CheckOptimizedMethod(int* registers,
+                            int number_of_references,
+                            uint32_t dex_pc,
+                            uint32_t native_pc_offset,
+                            bool search_for_valid_stack_map)
       REQUIRES_SHARED(Locks::mutator_lock_) {
     ArtMethod* m = GetMethod();
     CodeInfo code_info(GetCurrentOatQuickMethodHeader());
     StackMap stack_map = code_info.GetStackMapForNativePcOffset(native_pc_offset);
+    if (search_for_valid_stack_map && !code_info.GetStackMaskOf(stack_map).IsValid()) {
+      for (StackMap map : code_info.GetStackMaps()) {
+        if (map.GetDexPc() == dex_pc && code_info.GetStackMaskOf(map).IsValid()) {
+          stack_map = map;
+          break;
+        }
+      }
+    }
     CodeItemDataAccessor accessor(m->DexInstructionData());
     uint16_t number_of_dex_registers = accessor.RegistersSize();
 
@@ -93,6 +110,7 @@ class CheckReferenceMapVisitor : public StackVisitor {
           CHECK(false);
           break;
         case DexRegisterLocation::Kind::kInStack:
+          CHECK(stack_mask.IsValid());
           DCHECK_EQ(location.GetValue() % kFrameSlotSize, 0);
           CHECK(stack_mask.LoadBit(location.GetValue() / kFrameSlotSize));
           break;
diff --git a/test/004-ReferenceMap/stack_walk_refmap_jni.cc b/test/004-ReferenceMap/stack_walk_refmap_jni.cc
index 4c344a3a5f..4317b5d0ba 100644
--- a/test/004-ReferenceMap/stack_walk_refmap_jni.cc
+++ b/test/004-ReferenceMap/stack_walk_refmap_jni.cc
@@ -20,17 +20,21 @@
 
 namespace art {
 
-#define CHECK_REGS_CONTAIN_REFS(dex_pc, abort_if_not_found, ...) do {                 \
-  int t[] = {__VA_ARGS__};                                                            \
-  int t_size = sizeof(t) / sizeof(*t);                                                \
-  const OatQuickMethodHeader* method_header = GetCurrentOatQuickMethodHeader();       \
-  uintptr_t native_quick_pc = method_header->ToNativeQuickPc(GetMethod(),             \
-                                                 dex_pc,                              \
-                                                 /* is_catch_handler */ false,        \
-                                                 abort_if_not_found);                 \
-  if (native_quick_pc != UINTPTR_MAX) {                                               \
-    CheckReferences(t, t_size, method_header->NativeQuickPcOffset(native_quick_pc));  \
-  }                                                                                   \
+#define CHECK_REGS_CONTAIN_REFS(dex_pc, abort_if_not_found, ...) do {                         \
+  int t[] = {__VA_ARGS__};                                                                    \
+  int t_size = sizeof(t) / sizeof(*t);                                                        \
+  const OatQuickMethodHeader* method_header = GetCurrentOatQuickMethodHeader();               \
+  uintptr_t native_quick_pc = method_header->ToNativeQuickPc(GetMethod(),                     \
+                                                 dex_pc,                                      \
+                                                 /* is_catch_handler */ false,                \
+                                                 abort_if_not_found);                         \
+  if (native_quick_pc != UINTPTR_MAX) {                                                       \
+    CheckReferences(t,                                                                        \
+                    t_size,                                                                   \
+                    dex_pc,                                                                   \
+                    method_header->NativeQuickPcOffset(native_quick_pc),                      \
+                    /* search_for_valid_stack_map= */ true);                                  \
+  }                                                                                           \
 } while (false);
 
 struct ReferenceMap2Visitor : public CheckReferenceMapVisitor {
diff --git a/test/004-StackWalk/stack_walk_jni.cc b/test/004-StackWalk/stack_walk_jni.cc
index 6e53f30d66..3ec0fa3291 100644
--- a/test/004-StackWalk/stack_walk_jni.cc
+++ b/test/004-StackWalk/stack_walk_jni.cc
@@ -25,7 +25,8 @@ namespace art {
 #define CHECK_REGS_ARE_REFERENCES(...) do { \
   int t[] = {__VA_ARGS__}; \
   int t_size = sizeof(t) / sizeof(*t); \
-  CheckReferences(t, t_size, GetNativePcOffset()); \
+  CheckReferences( \
+      t, t_size, GetDexPc(), GetNativePcOffset(), /* search_for_valid_sack_map= */ false); \
 } while (false);
 
 static int gJava_StackWalk_refmap_calls = 0;