diff --git a/compiler/optimizing/stack_map_test.cc b/compiler/optimizing/stack_map_test.cc
index 87ac2e7..e7075c0 100644
--- a/compiler/optimizing/stack_map_test.cc
+++ b/compiler/optimizing/stack_map_test.cc
@@ -54,14 +54,14 @@
   StackMap stack_map = code_info.GetStackMapAt(0);
   ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForDexPc(0)));
   ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForNativePcOffset(64)));
-  ASSERT_EQ(0u, stack_map.GetDexPc());
-  ASSERT_EQ(64u, stack_map.GetNativePcOffset());
-  ASSERT_EQ(0x3u, stack_map.GetRegisterMask());
+  ASSERT_EQ(0u, stack_map.GetDexPc(code_info));
+  ASSERT_EQ(64u, stack_map.GetNativePcOffset(code_info));
+  ASSERT_EQ(0x3u, stack_map.GetRegisterMask(code_info));
 
-  MemoryRegion stack_mask = stack_map.GetStackMask();
+  MemoryRegion stack_mask = stack_map.GetStackMask(code_info);
   ASSERT_TRUE(SameBits(stack_mask, sp_mask));
 
-  ASSERT_TRUE(stack_map.HasDexRegisterMap());
+  ASSERT_TRUE(stack_map.HasDexRegisterMap(code_info));
   DexRegisterMap dex_registers = code_info.GetDexRegisterMapOf(stack_map, number_of_dex_registers);
   ASSERT_EQ(7u, dex_registers.Size());
   DexRegisterLocation location0 = dex_registers.GetLocationKindAndValue(0, number_of_dex_registers);
@@ -73,7 +73,7 @@
   ASSERT_EQ(0, location0.GetValue());
   ASSERT_EQ(-2, location1.GetValue());
 
-  ASSERT_FALSE(stack_map.HasInlineInfo());
+  ASSERT_FALSE(stack_map.HasInlineInfo(code_info));
 }
 
 TEST(StackMapTest, Test2) {
@@ -112,14 +112,14 @@
     StackMap stack_map = code_info.GetStackMapAt(0);
     ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForDexPc(0)));
     ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForNativePcOffset(64)));
-    ASSERT_EQ(0u, stack_map.GetDexPc());
-    ASSERT_EQ(64u, stack_map.GetNativePcOffset());
-    ASSERT_EQ(0x3u, stack_map.GetRegisterMask());
+    ASSERT_EQ(0u, stack_map.GetDexPc(code_info));
+    ASSERT_EQ(64u, stack_map.GetNativePcOffset(code_info));
+    ASSERT_EQ(0x3u, stack_map.GetRegisterMask(code_info));
 
-    MemoryRegion stack_mask = stack_map.GetStackMask();
+    MemoryRegion stack_mask = stack_map.GetStackMask(code_info);
     ASSERT_TRUE(SameBits(stack_mask, sp_mask1));
 
-    ASSERT_TRUE(stack_map.HasDexRegisterMap());
+    ASSERT_TRUE(stack_map.HasDexRegisterMap(code_info));
     DexRegisterMap dex_registers =
         code_info.GetDexRegisterMapOf(stack_map, number_of_dex_registers);
     ASSERT_EQ(7u, dex_registers.Size());
@@ -134,7 +134,7 @@
     ASSERT_EQ(0, location0.GetValue());
     ASSERT_EQ(-2, location1.GetValue());
 
-    ASSERT_TRUE(stack_map.HasInlineInfo());
+    ASSERT_TRUE(stack_map.HasInlineInfo(code_info));
     InlineInfo inline_info = code_info.GetInlineInfoOf(stack_map);
     ASSERT_EQ(2u, inline_info.GetDepth());
     ASSERT_EQ(42u, inline_info.GetMethodReferenceIndexAtDepth(0));
@@ -146,14 +146,14 @@
     StackMap stack_map = code_info.GetStackMapAt(1);
     ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForDexPc(1u)));
     ASSERT_TRUE(stack_map.Equals(code_info.GetStackMapForNativePcOffset(128u)));
-    ASSERT_EQ(1u, stack_map.GetDexPc());
-    ASSERT_EQ(128u, stack_map.GetNativePcOffset());
-    ASSERT_EQ(0xFFu, stack_map.GetRegisterMask());
+    ASSERT_EQ(1u, stack_map.GetDexPc(code_info));
+    ASSERT_EQ(128u, stack_map.GetNativePcOffset(code_info));
+    ASSERT_EQ(0xFFu, stack_map.GetRegisterMask(code_info));
 
-    MemoryRegion stack_mask = stack_map.GetStackMask();
+    MemoryRegion stack_mask = stack_map.GetStackMask(code_info);
     ASSERT_TRUE(SameBits(stack_mask, sp_mask2));
 
-    ASSERT_TRUE(stack_map.HasDexRegisterMap());
+    ASSERT_TRUE(stack_map.HasDexRegisterMap(code_info));
     DexRegisterMap dex_registers =
         code_info.GetDexRegisterMapOf(stack_map, number_of_dex_registers);
     ASSERT_EQ(3u, dex_registers.Size());
@@ -168,7 +168,7 @@
     ASSERT_EQ(18, location0.GetValue());
     ASSERT_EQ(3, location1.GetValue());
 
-    ASSERT_FALSE(stack_map.HasInlineInfo());
+    ASSERT_FALSE(stack_map.HasInlineInfo(code_info));
   }
 }
 
@@ -190,14 +190,45 @@
 
   CodeInfo code_info(region);
   StackMap stack_map = code_info.GetStackMapAt(0);
-  ASSERT_TRUE(stack_map.HasDexRegisterMap());
+  ASSERT_TRUE(stack_map.HasDexRegisterMap(code_info));
   DexRegisterMap dex_registers = code_info.GetDexRegisterMapOf(stack_map, 2);
   ASSERT_EQ(DexRegisterLocation::Kind::kNone,
             dex_registers.GetLocationKind(0, number_of_dex_registers));
   ASSERT_EQ(DexRegisterLocation::Kind::kConstant,
             dex_registers.GetLocationKind(1, number_of_dex_registers));
   ASSERT_EQ(-2, dex_registers.GetConstant(1, number_of_dex_registers));
-  ASSERT_FALSE(stack_map.HasInlineInfo());
+  ASSERT_FALSE(stack_map.HasInlineInfo(code_info));
+}
+
+// Generate a stack map whose dex register offset is
+// StackMap::kNoDexRegisterMapSmallEncoding, and ensure we do
+// not treat it as kNoDexRegisterMap.
+TEST(StackMapTest, DexRegisterMapOffsetOverflow) {
+  ArenaPool pool;
+  ArenaAllocator arena(&pool);
+  StackMapStream stream(&arena);
+
+  ArenaBitVector sp_mask(&arena, 0, false);
+  uint32_t number_of_dex_registers = 0xEA;
+  stream.AddStackMapEntry(0, 64, 0x3, &sp_mask, number_of_dex_registers, 0);
+  for (uint32_t i = 0; i < number_of_dex_registers - 9; ++i) {
+    stream.AddDexRegisterEntry(i, DexRegisterLocation::Kind::kConstant, 0);
+  }
+  stream.AddStackMapEntry(0, 64, 0x3, &sp_mask, number_of_dex_registers, 0);
+  for (uint32_t i = 0; i < number_of_dex_registers; ++i) {
+    stream.AddDexRegisterEntry(i, DexRegisterLocation::Kind::kConstant, 0);
+  }
+
+  size_t size = stream.ComputeNeededSize();
+  void* memory = arena.Alloc(size, kArenaAllocMisc);
+  MemoryRegion region(memory, size);
+  stream.FillIn(region);
+
+  CodeInfo code_info(region);
+  StackMap stack_map = code_info.GetStackMapAt(1);
+  ASSERT_TRUE(stack_map.HasDexRegisterMap(code_info));
+  ASSERT_NE(stack_map.GetDexRegisterMapOffset(code_info), StackMap::kNoDexRegisterMap);
+  ASSERT_EQ(stack_map.GetDexRegisterMapOffset(code_info), StackMap::kNoDexRegisterMapSmallEncoding);
 }
 
 }  // namespace art