diff --git a/test/004-ReferenceMap/stack_walk_refmap_jni.cc b/test/004-ReferenceMap/stack_walk_refmap_jni.cc
index 7929554..e914bd9 100644
--- a/test/004-ReferenceMap/stack_walk_refmap_jni.cc
+++ b/test/004-ReferenceMap/stack_walk_refmap_jni.cc
@@ -14,138 +14,57 @@
  * limitations under the License.
  */
 
-#include <stdio.h>
-#include <memory>
-
-#include "class_linker.h"
-#include "dex_file-inl.h"
-#include "gc_map.h"
-#include "mirror/art_method-inl.h"
-#include "mirror/class-inl.h"
-#include "mirror/object_array-inl.h"
-#include "mirror/object-inl.h"
-#include "scoped_thread_state_change.h"
-#include "thread.h"
+#include "check_reference_map_visitor.h"
 #include "jni.h"
-#include "verifier/method_verifier.h"
 
 namespace art {
 
-#define IS_IN_REF_BITMAP(ref_bitmap, reg) \
-    (((reg) < m->GetCodeItem()->registers_size_) && \
-     ((*((ref_bitmap) + (reg)/8) >> ((reg) % 8) ) & 0x01))
+#define CHECK_REGS_CONTAIN_REFS(native_pc_offset, ...) do { \
+  int t[] = {__VA_ARGS__}; \
+  int t_size = sizeof(t) / sizeof(*t); \
+  CheckReferences(t, t_size, m->NativePcOffset(m->ToNativePc(native_pc_offset))); \
+} while (false);
 
-#define CHECK_REGS_CONTAIN_REFS(...)     \
-  do {                                   \
-    int t[] = {__VA_ARGS__};             \
-    int t_size = sizeof(t) / sizeof(*t);      \
-    for (int i = 0; i < t_size; ++i)          \
-      CHECK(IS_IN_REF_BITMAP(ref_bitmap, t[i])) \
-          << "Error: Reg @ " << i << "-th argument is not in GC map"; \
-  } while (false)
-
-struct ReferenceMap2Visitor : public StackVisitor {
-  explicit ReferenceMap2Visitor(Thread* thread)
-      SHARED_LOCKS_REQUIRED(Locks::mutator_lock_)
-      : StackVisitor(thread, NULL) {
-  }
+struct ReferenceMap2Visitor : public CheckReferenceMapVisitor {
+  explicit ReferenceMap2Visitor(Thread* thread) SHARED_LOCKS_REQUIRED(Locks::mutator_lock_)
+      : CheckReferenceMapVisitor(thread) {}
 
   bool VisitFrame() SHARED_LOCKS_REQUIRED(Locks::mutator_lock_) {
+    if (CheckReferenceMapVisitor::VisitFrame()) {
+      return true;
+    }
     mirror::ArtMethod* m = GetMethod();
-    if (!m || m->IsNative() || m->IsRuntimeMethod() || IsShadowFrame()) {
-      return true;
-    }
-    LOG(INFO) << "At " << PrettyMethod(m, false);
-
-    NativePcOffsetToReferenceMap map(m->GetNativeGcMap());
-
-    if (m->IsCalleeSaveMethod()) {
-      LOG(WARNING) << "no PC for " << PrettyMethod(m);
-      return true;
-    }
-
-    const uint8_t* ref_bitmap = NULL;
     std::string m_name(m->GetName());
 
     // Given the method name and the number of times the method has been called,
     // we know the Dex registers with live reference values. Assert that what we
     // find is what is expected.
     if (m_name.compare("f") == 0) {
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x03U)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8);  // v8: this
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x06U)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 1);  // v8: this, v1: x
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x08U)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 3, 1);  // v8: this, v3: y, v1: x
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x0cU)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 3, 1);  // v8: this, v3: y, v1: x
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x0eU)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 3, 1);  // v8: this, v3: y, v1: x
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x10U)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 3, 1);  // v8: this, v3: y, v1: x
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x13U)));
-      CHECK(ref_bitmap);
+      CHECK_REGS_CONTAIN_REFS(0x03U, 8);  // v8: this
+      CHECK_REGS_CONTAIN_REFS(0x06U, 8, 1);  // v8: this, v1: x
+      CHECK_REGS_CONTAIN_REFS(0x08U, 8, 3, 1);  // v8: this, v3: y, v1: x
+      CHECK_REGS_CONTAIN_REFS(0x0cU, 8, 3, 1);  // v8: this, v3: y, v1: x
+      CHECK_REGS_CONTAIN_REFS(0x0eU, 8, 3, 1);  // v8: this, v3: y, v1: x
+      CHECK_REGS_CONTAIN_REFS(0x10U, 8, 3, 1);  // v8: this, v3: y, v1: x
       // v2 is added because of the instruction at DexPC 0024. Object merges with 0 is Object. See:
       //   0024: move-object v3, v2
       //   0025: goto 0013
       // Detaled dex instructions for ReferenceMap.java are at the end of this function.
       // CHECK_REGS_CONTAIN_REFS(8, 3, 2, 1);  // v8: this, v3: y, v2: y, v1: x
       // We eliminate the non-live registers at a return, so only v3 is live:
-      CHECK_REGS_CONTAIN_REFS(3);  // v3: y
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x18U)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 2, 1, 0);  // v8: this, v2: y, v1: x, v0: ex
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x1aU)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 5, 2, 1, 0);  // v8: this, v5: x[1], v2: y, v1: x, v0: ex
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x1dU)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 5, 2, 1, 0);  // v8: this, v5: x[1], v2: y, v1: x, v0: ex
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x1fU)));
-      CHECK(ref_bitmap);
+      CHECK_REGS_CONTAIN_REFS(0x13U);  // v3: y
+      CHECK_REGS_CONTAIN_REFS(0x18U, 8, 2, 1, 0);  // v8: this, v2: y, v1: x, v0: ex
+      CHECK_REGS_CONTAIN_REFS(0x1aU, 8, 5, 2, 1, 0);  // v8: this, v5: x[1], v2: y, v1: x, v0: ex
+      CHECK_REGS_CONTAIN_REFS(0x1dU, 8, 5, 2, 1, 0);  // v8: this, v5: x[1], v2: y, v1: x, v0: ex
       // v5 is removed from the root set because there is a "merge" operation.
       // See 0015: if-nez v2, 001f.
-      CHECK_REGS_CONTAIN_REFS(8, 2, 1, 0);  // v8: this, v2: y, v1: x, v0: ex
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x21U)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 2, 1, 0);  // v8: this, v2: y, v1: x, v0: ex
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x27U)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 4, 2, 1);  // v8: this, v4: ex, v2: y, v1: x
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x29U)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 4, 2, 1);  // v8: this, v4: ex, v2: y, v1: x
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x2cU)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 4, 2, 1);  // v8: this, v4: ex, v2: y, v1: x
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x2fU)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 4, 3, 2, 1);  // v8: this, v4: ex, v3: y, v2: y, v1: x
-
-      ref_bitmap = map.FindBitMap(m->NativePcOffset(m->ToNativePc(0x32U)));
-      CHECK(ref_bitmap);
-      CHECK_REGS_CONTAIN_REFS(8, 3, 2, 1, 0);  // v8: this, v3: y, v2: y, v1: x, v0: ex
+      CHECK_REGS_CONTAIN_REFS(0x1fU, 8, 2, 1, 0);  // v8: this, v2: y, v1: x, v0: ex
+      CHECK_REGS_CONTAIN_REFS(0x21U, 8, 2, 1, 0);  // v8: this, v2: y, v1: x, v0: ex
+      CHECK_REGS_CONTAIN_REFS(0x27U, 8, 4, 2, 1);  // v8: this, v4: ex, v2: y, v1: x
+      CHECK_REGS_CONTAIN_REFS(0x29U, 8, 4, 2, 1);  // v8: this, v4: ex, v2: y, v1: x
+      CHECK_REGS_CONTAIN_REFS(0x2cU, 8, 4, 2, 1);  // v8: this, v4: ex, v2: y, v1: x
+      CHECK_REGS_CONTAIN_REFS(0x2fU, 8, 4, 3, 2, 1);  // v8: this, v4: ex, v3: y, v2: y, v1: x
+      CHECK_REGS_CONTAIN_REFS(0x32U, 8, 3, 2, 1, 0);  // v8: this, v3: y, v2: y, v1: x, v0: ex
     }
 
     return true;
diff --git a/test/004-StackWalk/stack_walk_jni.cc b/test/004-StackWalk/stack_walk_jni.cc
index 30a0d59..c40de7e 100644
--- a/test/004-StackWalk/stack_walk_jni.cc
+++ b/test/004-StackWalk/stack_walk_jni.cc
@@ -14,54 +14,29 @@
  * limitations under the License.
  */
 
-#include <stdio.h>
-#include <memory>
-
-#include "class_linker.h"
-#include "gc_map.h"
-#include "mirror/art_method-inl.h"
-#include "mirror/class-inl.h"
-#include "mirror/object_array-inl.h"
-#include "mirror/object-inl.h"
+#include "check_reference_map_visitor.h"
 #include "jni.h"
-#include "scoped_thread_state_change.h"
 
 namespace art {
 
-#define REG(reg_bitmap, reg) \
-    (((reg) < m->GetCodeItem()->registers_size_) && \
-     ((*((reg_bitmap) + (reg)/8) >> ((reg) % 8) ) & 0x01))
-
-#define CHECK_REGS(...) if (!IsShadowFrame()) { \
-    int t[] = {__VA_ARGS__}; \
-    int t_size = sizeof(t) / sizeof(*t); \
-    for (int i = 0; i < t_size; ++i) \
-      CHECK(REG(reg_bitmap, t[i])) << "Error: Reg " << i << " is not in RegisterMap"; \
-  }
+#define CHECK_REGS(...) do { \
+  int t[] = {__VA_ARGS__}; \
+  int t_size = sizeof(t) / sizeof(*t); \
+  CheckReferences(t, t_size, GetNativePcOffset()); \
+} while (false);
 
 static int gJava_StackWalk_refmap_calls = 0;
 
-struct TestReferenceMapVisitor : public StackVisitor {
-  explicit TestReferenceMapVisitor(Thread* thread)
-      SHARED_LOCKS_REQUIRED(Locks::mutator_lock_)
-      : StackVisitor(thread, NULL) {
-  }
+class TestReferenceMapVisitor : public CheckReferenceMapVisitor {
+ public:
+  explicit TestReferenceMapVisitor(Thread* thread) SHARED_LOCKS_REQUIRED(Locks::mutator_lock_)
+      : CheckReferenceMapVisitor(thread) {}
 
   bool VisitFrame() SHARED_LOCKS_REQUIRED(Locks::mutator_lock_) {
-    mirror::ArtMethod* m = GetMethod();
-    CHECK(m != NULL);
-    LOG(INFO) << "At " << PrettyMethod(m, false);
-
-    if (m->IsCalleeSaveMethod() || m->IsNative()) {
-      LOG(WARNING) << "no PC for " << PrettyMethod(m);
-      CHECK_EQ(GetDexPc(), DexFile::kDexNoIndex);
+    if (CheckReferenceMapVisitor::VisitFrame()) {
       return true;
     }
-    const uint8_t* reg_bitmap = NULL;
-    if (!IsShadowFrame()) {
-      NativePcOffsetToReferenceMap map(m->GetNativeGcMap());
-      reg_bitmap = map.FindBitMap(GetNativePcOffset());
-    }
+    mirror::ArtMethod* m = GetMethod();
     StringPiece m_name(m->GetName());
 
     // Given the method name and the number of times the method has been called,
diff --git a/test/401-optimizing-compiler/src/Main.java b/test/401-optimizing-compiler/src/Main.java
index 2c6d1c2..07c407b 100644
--- a/test/401-optimizing-compiler/src/Main.java
+++ b/test/401-optimizing-compiler/src/Main.java
@@ -97,6 +97,11 @@
     if (exception == null) {
       throw new Error("Missing NullPointerException");
     }
+
+    result = $opt$InvokeVirtualMethod();
+    if (result != 42) {
+      throw new Error("Unexpected result: " + result);
+    }
   }
 
   public static void invokePrivate() {
@@ -205,5 +210,13 @@
     m.o = new Main();
   }
 
+  public static int $opt$InvokeVirtualMethod() {
+    return new Main().virtualMethod();
+  }
+
+  public int virtualMethod() {
+    return 42;
+  }
+
   Object o;
 }