Implement invoke virtual in optimizing compiler.

Also refactor 004 tests to make them work with both Quick and
Optimizing.

Change-Id: I87e275cb0ae0258fc3bb32b612140000b1d2adf8

15 files changed, 419 insertions, 224 deletions

diff --git a/compiler/driver/compiler_driver.h b/compiler/driver/compiler_driver.h
index 3c76098109..e7bd35776a 100644
--- a/compiler/driver/compiler_driver.h
+++ b/compiler/driver/compiler_driver.h
@@ -260,7 +260,7 @@ class CompilerDriver {
       uint16_t* declaring_class_idx, uint16_t* declaring_method_idx)
     SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
-  // Get declaration location of a resolved field.
+  // Get the index in the vtable of the method.
   uint16_t GetResolvedMethodVTableIndex(
       mirror::ArtMethod* resolved_method, InvokeType type)
     SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
diff --git a/compiler/optimizing/builder.cc b/compiler/optimizing/builder.cc
index a03588f4fd..33b00d2ac9 100644
--- a/compiler/optimizing/builder.cc
+++ b/compiler/optimizing/builder.cc
@@ -331,18 +331,61 @@ bool HGraphBuilder::BuildInvoke(const Instruction& instruction,
                                 bool is_range,
                                 uint32_t* args,
                                 uint32_t register_index) {
+  Instruction::Code opcode = instruction.Opcode();
+  InvokeType invoke_type;
+  switch (opcode) {
+    case Instruction::INVOKE_STATIC:
+    case Instruction::INVOKE_STATIC_RANGE:
+      invoke_type = kStatic;
+      break;
+    case Instruction::INVOKE_DIRECT:
+    case Instruction::INVOKE_DIRECT_RANGE:
+      invoke_type = kDirect;
+      break;
+    case Instruction::INVOKE_VIRTUAL:
+    case Instruction::INVOKE_VIRTUAL_RANGE:
+      invoke_type = kVirtual;
+      break;
+    case Instruction::INVOKE_INTERFACE:
+    case Instruction::INVOKE_INTERFACE_RANGE:
+      invoke_type = kInterface;
+      break;
+    case Instruction::INVOKE_SUPER_RANGE:
+    case Instruction::INVOKE_SUPER:
+      invoke_type = kSuper;
+      break;
+    default:
+      LOG(FATAL) << "Unexpected invoke op: " << opcode;
+      return false;
+  }
+
   const DexFile::MethodId& method_id = dex_file_->GetMethodId(method_idx);
   const DexFile::ProtoId& proto_id = dex_file_->GetProtoId(method_id.proto_idx_);
   const char* descriptor = dex_file_->StringDataByIdx(proto_id.shorty_idx_);
   Primitive::Type return_type = Primitive::GetType(descriptor[0]);
-  bool is_instance_call =
-      instruction.Opcode() != Instruction::INVOKE_STATIC
-      && instruction.Opcode() != Instruction::INVOKE_STATIC_RANGE;
+  bool is_instance_call = invoke_type != kStatic;
   const size_t number_of_arguments = strlen(descriptor) - (is_instance_call ? 0 : 1);
 
-  // Treat invoke-direct like static calls for now.
-  HInvoke* invoke = new (arena_) HInvokeStatic(
-      arena_, number_of_arguments, return_type, dex_offset, method_idx);
+  HInvoke* invoke = nullptr;
+  if (invoke_type == kVirtual) {
+    MethodReference target_method(dex_file_, method_idx);
+    uintptr_t direct_code;
+    uintptr_t direct_method;
+    int vtable_index;
+    // TODO: Add devirtualization support.
+    compiler_driver_->ComputeInvokeInfo(dex_compilation_unit_, dex_offset, true, true,
+                                        &invoke_type, &target_method, &vtable_index,
+                                        &direct_code, &direct_method);
+    if (vtable_index == -1) {
+      return false;
+    }
+    invoke = new (arena_) HInvokeVirtual(
+        arena_, number_of_arguments, return_type, dex_offset, vtable_index);
+  } else {
+    // Treat invoke-direct like static calls for now.
+    invoke = new (arena_) HInvokeStatic(
+        arena_, number_of_arguments, return_type, dex_offset, method_idx);
+  }
 
   size_t start_index = 0;
   Temporaries temps(graph_, is_instance_call ? 1 : 0);
@@ -620,7 +663,8 @@ bool HGraphBuilder::AnalyzeDexInstruction(const Instruction& instruction, uint32
     }
 
     case Instruction::INVOKE_STATIC:
-    case Instruction::INVOKE_DIRECT: {
+    case Instruction::INVOKE_DIRECT:
+    case Instruction::INVOKE_VIRTUAL: {
       uint32_t method_idx = instruction.VRegB_35c();
       uint32_t number_of_vreg_arguments = instruction.VRegA_35c();
       uint32_t args[5];
@@ -632,7 +676,8 @@ bool HGraphBuilder::AnalyzeDexInstruction(const Instruction& instruction, uint32
     }
 
     case Instruction::INVOKE_STATIC_RANGE:
-    case Instruction::INVOKE_DIRECT_RANGE: {
+    case Instruction::INVOKE_DIRECT_RANGE:
+    case Instruction::INVOKE_VIRTUAL_RANGE: {
       uint32_t method_idx = instruction.VRegB_3rc();
       uint32_t number_of_vreg_arguments = instruction.VRegA_3rc();
       uint32_t register_index = instruction.VRegC();
diff --git a/compiler/optimizing/code_generator_arm.cc b/compiler/optimizing/code_generator_arm.cc
index 99030922a7..ad622798a6 100644
--- a/compiler/optimizing/code_generator_arm.cc
+++ b/compiler/optimizing/code_generator_arm.cc
@@ -20,6 +20,7 @@
 #include "gc/accounting/card_table.h"
 #include "mirror/array.h"
 #include "mirror/art_method.h"
+#include "mirror/class.h"
 #include "thread.h"
 #include "utils/assembler.h"
 #include "utils/arm/assembler_arm.h"
@@ -818,6 +819,47 @@ void InstructionCodeGeneratorARM::VisitReturn(HReturn* ret) {
 }
 
 void LocationsBuilderARM::VisitInvokeStatic(HInvokeStatic* invoke) {
+  HandleInvoke(invoke);
+}
+
+void InstructionCodeGeneratorARM::LoadCurrentMethod(Register reg) {
+  __ ldr(reg, Address(SP, kCurrentMethodStackOffset));
+}
+
+void InstructionCodeGeneratorARM::VisitInvokeStatic(HInvokeStatic* invoke) {
+  Register temp = invoke->GetLocations()->GetTemp(0).AsArm().AsCoreRegister();
+  uint32_t heap_reference_size = sizeof(mirror::HeapReference<mirror::Object>);
+  size_t index_in_cache = mirror::Array::DataOffset(heap_reference_size).Int32Value() +
+      invoke->GetIndexInDexCache() * kArmWordSize;
+
+  // TODO: Implement all kinds of calls:
+  // 1) boot -> boot
+  // 2) app -> boot
+  // 3) app -> app
+  //
+  // Currently we implement the app -> app logic, which looks up in the resolve cache.
+
+  // temp = method;
+  LoadCurrentMethod(temp);
+  // temp = temp->dex_cache_resolved_methods_;
+  __ ldr(temp, Address(temp, mirror::ArtMethod::DexCacheResolvedMethodsOffset().Int32Value()));
+  // temp = temp[index_in_cache]
+  __ ldr(temp, Address(temp, index_in_cache));
+  // LR = temp[offset_of_quick_compiled_code]
+  __ ldr(LR, Address(temp,
+                     mirror::ArtMethod::EntryPointFromQuickCompiledCodeOffset().Int32Value()));
+  // LR()
+  __ blx(LR);
+
+  codegen_->RecordPcInfo(invoke, invoke->GetDexPc());
+  DCHECK(!codegen_->IsLeafMethod());
+}
+
+void LocationsBuilderARM::VisitInvokeVirtual(HInvokeVirtual* invoke) {
+  HandleInvoke(invoke);
+}
+
+void LocationsBuilderARM::HandleInvoke(HInvoke* invoke) {
   LocationSummary* locations =
       new (GetGraph()->GetArena()) LocationSummary(invoke, LocationSummary::kCall);
   locations->AddTemp(ArmCoreLocation(R0));
@@ -852,37 +894,30 @@ void LocationsBuilderARM::VisitInvokeStatic(HInvokeStatic* invoke) {
   }
 }
 
-void InstructionCodeGeneratorARM::LoadCurrentMethod(Register reg) {
-  __ ldr(reg, Address(SP, kCurrentMethodStackOffset));
-}
 
-void InstructionCodeGeneratorARM::VisitInvokeStatic(HInvokeStatic* invoke) {
+void InstructionCodeGeneratorARM::VisitInvokeVirtual(HInvokeVirtual* invoke) {
   Register temp = invoke->GetLocations()->GetTemp(0).AsArm().AsCoreRegister();
-  uint32_t heap_reference_size = sizeof(mirror::HeapReference<mirror::Object>);
-  size_t index_in_cache = mirror::Array::DataOffset(heap_reference_size).Int32Value() +
-      invoke->GetIndexInDexCache() * kArmWordSize;
-
-  // TODO: Implement all kinds of calls:
-  // 1) boot -> boot
-  // 2) app -> boot
-  // 3) app -> app
-  //
-  // Currently we implement the app -> app logic, which looks up in the resolve cache.
-
-  // temp = method;
-  LoadCurrentMethod(temp);
-  // temp = temp->dex_cache_resolved_methods_;
-  __ ldr(temp, Address(temp, mirror::ArtMethod::DexCacheResolvedMethodsOffset().Int32Value()));
-  // temp = temp[index_in_cache]
-  __ ldr(temp, Address(temp, index_in_cache));
-  // LR = temp[offset_of_quick_compiled_code]
-  __ ldr(LR, Address(temp,
-                     mirror::ArtMethod::EntryPointFromQuickCompiledCodeOffset().Int32Value()));
-  // LR()
+  uint32_t method_offset = mirror::Class::EmbeddedVTableOffset().Uint32Value() +
+          invoke->GetVTableIndex() * sizeof(mirror::Class::VTableEntry);
+  LocationSummary* locations = invoke->GetLocations();
+  Location receiver = locations->InAt(0);
+  uint32_t class_offset = mirror::Object::ClassOffset().Int32Value();
+  // temp = object->GetClass();
+  if (receiver.IsStackSlot()) {
+    __ ldr(temp, Address(SP, receiver.GetStackIndex()));
+    __ ldr(temp, Address(temp, class_offset));
+  } else {
+    __ ldr(temp, Address(receiver.AsArm().AsCoreRegister(), class_offset));
+  }
+  // temp = temp->GetMethodAt(method_offset);
+  uint32_t entry_point = mirror::ArtMethod::EntryPointFromQuickCompiledCodeOffset().Int32Value();
+  __ ldr(temp, Address(temp, method_offset));
+  // LR = temp->GetEntryPoint();
+  __ ldr(LR, Address(temp, entry_point));
+  // LR();
   __ blx(LR);
-
-  codegen_->RecordPcInfo(invoke, invoke->GetDexPc());
   DCHECK(!codegen_->IsLeafMethod());
+  codegen_->RecordPcInfo(invoke, invoke->GetDexPc());
 }
 
 void LocationsBuilderARM::VisitAdd(HAdd* add) {
diff --git a/compiler/optimizing/code_generator_arm.h b/compiler/optimizing/code_generator_arm.h
index 660294b147..2480960f32 100644
--- a/compiler/optimizing/code_generator_arm.h
+++ b/compiler/optimizing/code_generator_arm.h
@@ -93,6 +93,8 @@ class LocationsBuilderARM : public HGraphVisitor {
 
 #undef DECLARE_VISIT_INSTRUCTION
 
+  void HandleInvoke(HInvoke* invoke);
+
  private:
   CodeGeneratorARM* const codegen_;
   InvokeDexCallingConventionVisitor parameter_visitor_;
diff --git a/compiler/optimizing/code_generator_x86.cc b/compiler/optimizing/code_generator_x86.cc
index 3dd9b37158..3383cb2117 100644
--- a/compiler/optimizing/code_generator_x86.cc
+++ b/compiler/optimizing/code_generator_x86.cc
@@ -20,6 +20,7 @@
 #include "gc/accounting/card_table.h"
 #include "mirror/array.h"
 #include "mirror/art_method.h"
+#include "mirror/class.h"
 #include "thread.h"
 #include "utils/assembler.h"
 #include "utils/stack_checks.h"
@@ -763,6 +764,40 @@ void InstructionCodeGeneratorX86::VisitReturn(HReturn* ret) {
 }
 
 void LocationsBuilderX86::VisitInvokeStatic(HInvokeStatic* invoke) {
+  HandleInvoke(invoke);
+}
+
+void InstructionCodeGeneratorX86::VisitInvokeStatic(HInvokeStatic* invoke) {
+  Register temp = invoke->GetLocations()->GetTemp(0).AsX86().AsCpuRegister();
+  uint32_t heap_reference_size = sizeof(mirror::HeapReference<mirror::Object>);
+  size_t index_in_cache = mirror::Array::DataOffset(heap_reference_size).Int32Value() +
+      invoke->GetIndexInDexCache() * kX86WordSize;
+
+  // TODO: Implement all kinds of calls:
+  // 1) boot -> boot
+  // 2) app -> boot
+  // 3) app -> app
+  //
+  // Currently we implement the app -> app logic, which looks up in the resolve cache.
+
+  // temp = method;
+  LoadCurrentMethod(temp);
+  // temp = temp->dex_cache_resolved_methods_;
+  __ movl(temp, Address(temp, mirror::ArtMethod::DexCacheResolvedMethodsOffset().Int32Value()));
+  // temp = temp[index_in_cache]
+  __ movl(temp, Address(temp, index_in_cache));
+  // (temp + offset_of_quick_compiled_code)()
+  __ call(Address(temp, mirror::ArtMethod::EntryPointFromQuickCompiledCodeOffset().Int32Value()));
+
+  DCHECK(!codegen_->IsLeafMethod());
+  codegen_->RecordPcInfo(invoke, invoke->GetDexPc());
+}
+
+void LocationsBuilderX86::VisitInvokeVirtual(HInvokeVirtual* invoke) {
+  HandleInvoke(invoke);
+}
+
+void LocationsBuilderX86::HandleInvoke(HInvoke* invoke) {
   LocationSummary* locations =
       new (GetGraph()->GetArena()) LocationSummary(invoke, LocationSummary::kCall);
   locations->AddTemp(X86CpuLocation(EAX));
@@ -799,26 +834,23 @@ void LocationsBuilderX86::VisitInvokeStatic(HInvokeStatic* invoke) {
   invoke->SetLocations(locations);
 }
 
-void InstructionCodeGeneratorX86::VisitInvokeStatic(HInvokeStatic* invoke) {
+void InstructionCodeGeneratorX86::VisitInvokeVirtual(HInvokeVirtual* invoke) {
   Register temp = invoke->GetLocations()->GetTemp(0).AsX86().AsCpuRegister();
-  uint32_t heap_reference_size = sizeof(mirror::HeapReference<mirror::Object>);
-  size_t index_in_cache = mirror::Array::DataOffset(heap_reference_size).Int32Value() +
-      invoke->GetIndexInDexCache() * kX86WordSize;
-
-  // TODO: Implement all kinds of calls:
-  // 1) boot -> boot
-  // 2) app -> boot
-  // 3) app -> app
-  //
-  // Currently we implement the app -> app logic, which looks up in the resolve cache.
-
-  // temp = method;
-  LoadCurrentMethod(temp);
-  // temp = temp->dex_cache_resolved_methods_;
-  __ movl(temp, Address(temp, mirror::ArtMethod::DexCacheResolvedMethodsOffset().Int32Value()));
-  // temp = temp[index_in_cache]
-  __ movl(temp, Address(temp, index_in_cache));
-  // (temp + offset_of_quick_compiled_code)()
+  uint32_t method_offset = mirror::Class::EmbeddedVTableOffset().Uint32Value() +
+          invoke->GetVTableIndex() * sizeof(mirror::Class::VTableEntry);
+  LocationSummary* locations = invoke->GetLocations();
+  Location receiver = locations->InAt(0);
+  uint32_t class_offset = mirror::Object::ClassOffset().Int32Value();
+  // temp = object->GetClass();
+  if (receiver.IsStackSlot()) {
+    __ movl(temp, Address(ESP, receiver.GetStackIndex()));
+    __ movl(temp, Address(temp, class_offset));
+  } else {
+    __ movl(temp, Address(receiver.AsX86().AsCpuRegister(), class_offset));
+  }
+  // temp = temp->GetMethodAt(method_offset);
+  __ movl(temp, Address(temp, method_offset));
+  // call temp->GetEntryPoint();
   __ call(Address(temp, mirror::ArtMethod::EntryPointFromQuickCompiledCodeOffset().Int32Value()));
 
   DCHECK(!codegen_->IsLeafMethod());
diff --git a/compiler/optimizing/code_generator_x86.h b/compiler/optimizing/code_generator_x86.h
index 7c502049d8..f1be0ad5b7 100644
--- a/compiler/optimizing/code_generator_x86.h
+++ b/compiler/optimizing/code_generator_x86.h
@@ -94,6 +94,8 @@ class LocationsBuilderX86 : public HGraphVisitor {
 
 #undef DECLARE_VISIT_INSTRUCTION
 
+  void HandleInvoke(HInvoke* invoke);
+
  private:
   CodeGeneratorX86* const codegen_;
   InvokeDexCallingConventionVisitor parameter_visitor_;
diff --git a/compiler/optimizing/code_generator_x86_64.cc b/compiler/optimizing/code_generator_x86_64.cc
index 2f352e0838..ca03af8e9f 100644
--- a/compiler/optimizing/code_generator_x86_64.cc
+++ b/compiler/optimizing/code_generator_x86_64.cc
@@ -20,6 +20,7 @@
 #include "gc/accounting/card_table.h"
 #include "mirror/array.h"
 #include "mirror/art_method.h"
+#include "mirror/class.h"
 #include "mirror/object_reference.h"
 #include "thread.h"
 #include "utils/assembler.h"
@@ -709,12 +710,46 @@ Location InvokeDexCallingConventionVisitor::GetNextLocation(Primitive::Type type
 }
 
 void LocationsBuilderX86_64::VisitInvokeStatic(HInvokeStatic* invoke) {
+  HandleInvoke(invoke);
+}
+
+void InstructionCodeGeneratorX86_64::VisitInvokeStatic(HInvokeStatic* invoke) {
+  CpuRegister temp = invoke->GetLocations()->GetTemp(0).AsX86_64().AsCpuRegister();
+  uint32_t heap_reference_size = sizeof(mirror::HeapReference<mirror::Object>);
+  size_t index_in_cache = mirror::Array::DataOffset(heap_reference_size).SizeValue() +
+      invoke->GetIndexInDexCache() * heap_reference_size;
+
+  // TODO: Implement all kinds of calls:
+  // 1) boot -> boot
+  // 2) app -> boot
+  // 3) app -> app
+  //
+  // Currently we implement the app -> app logic, which looks up in the resolve cache.
+
+  // temp = method;
+  LoadCurrentMethod(temp);
+  // temp = temp->dex_cache_resolved_methods_;
+  __ movl(temp, Address(temp, mirror::ArtMethod::DexCacheResolvedMethodsOffset().SizeValue()));
+  // temp = temp[index_in_cache]
+  __ movl(temp, Address(temp, index_in_cache));
+  // (temp + offset_of_quick_compiled_code)()
+  __ call(Address(temp, mirror::ArtMethod::EntryPointFromQuickCompiledCodeOffset().SizeValue()));
+
+  DCHECK(!codegen_->IsLeafMethod());
+  codegen_->RecordPcInfo(invoke, invoke->GetDexPc());
+}
+
+void LocationsBuilderX86_64::VisitInvokeVirtual(HInvokeVirtual* invoke) {
+  HandleInvoke(invoke);
+}
+
+void LocationsBuilderX86_64::HandleInvoke(HInvoke* invoke) {
   LocationSummary* locations =
       new (GetGraph()->GetArena()) LocationSummary(invoke, LocationSummary::kCall);
   locations->AddTemp(X86_64CpuLocation(RDI));
 
   InvokeDexCallingConventionVisitor calling_convention_visitor;
-  for (size_t i = 0; i < invoke->InputCount(); ++i) {
+  for (size_t i = 0; i < invoke->InputCount(); i++) {
     HInstruction* input = invoke->InputAt(i);
     locations->SetInAt(i, calling_convention_visitor.GetNextLocation(input->GetType()));
   }
@@ -740,26 +775,23 @@ void LocationsBuilderX86_64::VisitInvokeStatic(HInvokeStatic* invoke) {
   }
 }
 
-void InstructionCodeGeneratorX86_64::VisitInvokeStatic(HInvokeStatic* invoke) {
+void InstructionCodeGeneratorX86_64::VisitInvokeVirtual(HInvokeVirtual* invoke) {
   CpuRegister temp = invoke->GetLocations()->GetTemp(0).AsX86_64().AsCpuRegister();
-  uint32_t heap_reference_size = sizeof(mirror::HeapReference<mirror::Object>);
-  size_t index_in_cache = mirror::Array::DataOffset(heap_reference_size).SizeValue() +
-      invoke->GetIndexInDexCache() * heap_reference_size;
-
-  // TODO: Implement all kinds of calls:
-  // 1) boot -> boot
-  // 2) app -> boot
-  // 3) app -> app
-  //
-  // Currently we implement the app -> app logic, which looks up in the resolve cache.
-
-  // temp = method;
-  LoadCurrentMethod(temp);
-  // temp = temp->dex_cache_resolved_methods_;
-  __ movl(temp, Address(temp, mirror::ArtMethod::DexCacheResolvedMethodsOffset().SizeValue()));
-  // temp = temp[index_in_cache]
-  __ movl(temp, Address(temp, index_in_cache));
-  // (temp + offset_of_quick_compiled_code)()
+  size_t method_offset = mirror::Class::EmbeddedVTableOffset().SizeValue() +
+          invoke->GetVTableIndex() * sizeof(mirror::Class::VTableEntry);
+  LocationSummary* locations = invoke->GetLocations();
+  Location receiver = locations->InAt(0);
+  size_t class_offset = mirror::Object::ClassOffset().SizeValue();
+  // temp = object->GetClass();
+  if (receiver.IsStackSlot()) {
+    __ movq(temp, Address(CpuRegister(RSP), receiver.GetStackIndex()));
+    __ movq(temp, Address(temp, class_offset));
+  } else {
+    __ movq(temp, Address(receiver.AsX86_64().AsCpuRegister(), class_offset));
+  }
+  // temp = temp->GetMethodAt(method_offset);
+  __ movl(temp, Address(temp, method_offset));
+  // call temp->GetEntryPoint();
   __ call(Address(temp, mirror::ArtMethod::EntryPointFromQuickCompiledCodeOffset().SizeValue()));
 
   DCHECK(!codegen_->IsLeafMethod());
diff --git a/compiler/optimizing/code_generator_x86_64.h b/compiler/optimizing/code_generator_x86_64.h
index 44552ea465..78b60fe93c 100644
--- a/compiler/optimizing/code_generator_x86_64.h
+++ b/compiler/optimizing/code_generator_x86_64.h
@@ -91,6 +91,8 @@ class LocationsBuilderX86_64 : public HGraphVisitor {
 
 #undef DECLARE_VISIT_INSTRUCTION
 
+  void HandleInvoke(HInvoke* invoke);
+
  private:
   CodeGeneratorX86_64* const codegen_;
   InvokeDexCallingConventionVisitor parameter_visitor_;
diff --git a/compiler/optimizing/nodes.h b/compiler/optimizing/nodes.h
index ed6dd939de..d6dfeaede8 100644
--- a/compiler/optimizing/nodes.h
+++ b/compiler/optimizing/nodes.h
@@ -422,6 +422,7 @@ class HBasicBlock : public ArenaObject {
   M(If)                                                    \
   M(IntConstant)                                           \
   M(InvokeStatic)                                          \
+  M(InvokeVirtual)                                         \
   M(LoadLocal)                                             \
   M(Local)                                                 \
   M(LongConstant)                                          \
@@ -1272,6 +1273,26 @@ class HInvokeStatic : public HInvoke {
   DISALLOW_COPY_AND_ASSIGN(HInvokeStatic);
 };
 
+class HInvokeVirtual : public HInvoke {
+ public:
+  HInvokeVirtual(ArenaAllocator* arena,
+                 uint32_t number_of_arguments,
+                 Primitive::Type return_type,
+                 uint32_t dex_pc,
+                 uint32_t vtable_index)
+      : HInvoke(arena, number_of_arguments, return_type, dex_pc),
+        vtable_index_(vtable_index) {}
+
+  uint32_t GetVTableIndex() const { return vtable_index_; }
+
+  DECLARE_INSTRUCTION(InvokeVirtual);
+
+ private:
+  const uint32_t vtable_index_;
+
+  DISALLOW_COPY_AND_ASSIGN(HInvokeVirtual);
+};
+
 class HNewInstance : public HExpression<0> {
  public:
   HNewInstance(uint32_t dex_pc, uint16_t type_index)
diff --git a/runtime/check_reference_map_visitor.h b/runtime/check_reference_map_visitor.h
new file mode 100644
index 0000000000..166e23f3de
--- /dev/null
+++ b/runtime/check_reference_map_visitor.h
@@ -0,0 +1,111 @@
+/*
+ * Copyright (C) 2011 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef ART_RUNTIME_CHECK_REFERENCE_MAP_VISITOR_H_
+#define ART_RUNTIME_CHECK_REFERENCE_MAP_VISITOR_H_
+
+#include "gc_map.h"
+#include "mirror/art_method-inl.h"
+#include "scoped_thread_state_change.h"
+#include "stack_map.h"
+
+namespace art {
+
+// Helper class for tests checking that the compiler keeps track of dex registers
+// holding references.
+class CheckReferenceMapVisitor : public StackVisitor {
+ public:
+  explicit CheckReferenceMapVisitor(Thread* thread) SHARED_LOCKS_REQUIRED(Locks::mutator_lock_)
+      : StackVisitor(thread, nullptr) {}
+
+  bool VisitFrame() SHARED_LOCKS_REQUIRED(Locks::mutator_lock_) {
+    mirror::ArtMethod* m = GetMethod();
+    if (m->IsCalleeSaveMethod() || m->IsNative()) {
+      CHECK_EQ(GetDexPc(), DexFile::kDexNoIndex);
+    }
+
+    if (!m || m->IsNative() || m->IsRuntimeMethod() || IsShadowFrame()) {
+      return true;
+    }
+
+    LOG(INFO) << "At " << PrettyMethod(m, false);
+
+    if (m->IsCalleeSaveMethod()) {
+      LOG(WARNING) << "no PC for " << PrettyMethod(m);
+      return true;
+    }
+
+    return false;
+  }
+
+  void CheckReferences(int* registers, int number_of_references, uint32_t native_pc_offset)
+      SHARED_LOCKS_REQUIRED(Locks::mutator_lock_) {
+    if (GetMethod()->IsOptimized()) {
+      CheckOptimizedMethod(registers, number_of_references, native_pc_offset);
+    } else {
+      CheckQuickMethod(registers, number_of_references, native_pc_offset);
+    }
+  }
+
+ private:
+  void CheckOptimizedMethod(int* registers, int number_of_references, uint32_t native_pc_offset)
+      SHARED_LOCKS_REQUIRED(Locks::mutator_lock_) {
+    mirror::ArtMethod* m = GetMethod();
+    CodeInfo code_info = m->GetOptimizedCodeInfo();
+    StackMap stack_map = code_info.GetStackMapForNativePcOffset(native_pc_offset);
+    DexRegisterMap dex_register_map = code_info.GetDexRegisterMapOf(stack_map, m->GetCodeItem()->registers_size_);
+    MemoryRegion stack_mask = stack_map.GetStackMask();
+    uint32_t register_mask = stack_map.GetRegisterMask();
+    for (int i = 0; i < number_of_references; ++i) {
+      int reg = registers[i];
+      CHECK(reg < m->GetCodeItem()->registers_size_);
+      DexRegisterMap::LocationKind location = dex_register_map.GetLocationKind(reg);
+      switch (location) {
+        case DexRegisterMap::kNone:
+          // Not set, should not be a reference.
+          CHECK(false);
+          break;
+        case DexRegisterMap::kInStack:
+          CHECK(stack_mask.LoadBit(dex_register_map.GetValue(reg) >> 2));
+          break;
+        case DexRegisterMap::kInRegister:
+          CHECK_NE(register_mask & dex_register_map.GetValue(reg), 0u);
+          break;
+        case DexRegisterMap::kConstant:
+          CHECK_EQ(dex_register_map.GetValue(0), 0);
+          break;
+      }
+    }
+  }
+
+  void CheckQuickMethod(int* registers, int number_of_references, uint32_t native_pc_offset)
+      SHARED_LOCKS_REQUIRED(Locks::mutator_lock_) {
+    mirror::ArtMethod* m = GetMethod();
+    NativePcOffsetToReferenceMap map(m->GetNativeGcMap());
+    const uint8_t* ref_bitmap = map.FindBitMap(native_pc_offset);
+    CHECK(ref_bitmap);
+    for (int i = 0; i < number_of_references; ++i) {
+      int reg = registers[i];
+      CHECK(reg < m->GetCodeItem()->registers_size_);
+      CHECK((*((ref_bitmap) + reg / 8) >> (reg % 8) ) & 0x01)
+          << "Error: Reg @" << i << " is not in GC map";
+    }
+  }
+};
+
+}  // namespace art
+
+#endif  // ART_RUNTIME_REFERENCE_MAP_VISITOR_H_
diff --git a/runtime/mirror/art_method-inl.h b/runtime/mirror/art_method-inl.h
index ae170702b5..8447616cf5 100644
--- a/runtime/mirror/art_method-inl.h
+++ b/runtime/mirror/art_method-inl.h
@@ -285,14 +285,17 @@ inline const uint8_t* ArtMethod::GetVmapTable(const void* code_pointer) {
 }
 
 inline StackMap ArtMethod::GetStackMap(uint32_t native_pc_offset) {
+  return GetOptimizedCodeInfo().GetStackMapForNativePcOffset(native_pc_offset);
+}
+
+inline CodeInfo ArtMethod::GetOptimizedCodeInfo() {
   DCHECK(IsOptimized());
   const void* code_pointer = GetQuickOatCodePointer();
   DCHECK(code_pointer != nullptr);
   uint32_t offset =
       reinterpret_cast<const OatQuickMethodHeader*>(code_pointer)[-1].vmap_table_offset_;
   const void* data = reinterpret_cast<const void*>(reinterpret_cast<const uint8_t*>(code_pointer) - offset);
-  CodeInfo code_info(data);
-  return code_info.GetStackMapForNativePcOffset(native_pc_offset);
+  return CodeInfo(data);
 }
 
 inline void ArtMethod::SetOatNativeGcMapOffset(uint32_t gc_map_offset) {
diff --git a/runtime/mirror/art_method.h b/runtime/mirror/art_method.h
index d37aa57688..de6ec05442 100644
--- a/runtime/mirror/art_method.h
+++ b/runtime/mirror/art_method.h
@@ -155,7 +155,9 @@ class MANAGED ArtMethod FINAL : public Object {
     // Temporary solution for detecting if a method has been optimized: the compiler
     // does not create a GC map. Instead, the vmap table contains the stack map
     // (as in stack_map.h).
-    return (GetEntryPointFromQuickCompiledCode() != nullptr) && (GetNativeGcMap() == nullptr);
+    return (GetEntryPointFromQuickCompiledCode() != nullptr)
+        && (GetQuickOatCodePointer() != nullptr)
+        && (GetNativeGcMap() == nullptr);
   }
 
   bool IsPortableCompiled() SHARED_LOCKS_REQUIRED(Locks::mutator_lock_) {
@@ -349,6 +351,7 @@ class MANAGED ArtMethod FINAL : public Object {
       SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
   StackMap GetStackMap(uint32_t native_pc_offset) SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
+  CodeInfo GetOptimizedCodeInfo() SHARED_LOCKS_REQUIRED(Locks::mutator_lock_);
 
   const uint8_t* GetNativeGcMap() SHARED_LOCKS_REQUIRED(Locks::mutator_lock_) {
     return GetFieldPtr<uint8_t*>(OFFSET_OF_OBJECT_MEMBER(ArtMethod, gc_map_));
diff --git a/test/004-ReferenceMap/stack_walk_refmap_jni.cc b/test/004-ReferenceMap/stack_walk_refmap_jni.cc
index 79295549c9..e914bd9d09 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(...)     \
-  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)
+#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);
 
-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_) {
-    mirror::ArtMethod* m = GetMethod();
-    if (!m || m->IsNative() || m->IsRuntimeMethod() || IsShadowFrame()) {
+    if (CheckReferenceMapVisitor::VisitFrame()) {
       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;
+    mirror::ArtMethod* m = GetMethod();
     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 30a0d5906a..c40de7e247 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 2c6d1c2b60..07c407b565 100644
--- a/test/401-optimizing-compiler/src/Main.java
+++ b/test/401-optimizing-compiler/src/Main.java
@@ -97,6 +97,11 @@ public class Main {
     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 @@ public class Main {
     m.o = new Main();
   }
 
+  public static int $opt$InvokeVirtualMethod() {
+    return new Main().virtualMethod();
+  }
+
+  public int virtualMethod() {
+    return 42;
+  }
+
   Object o;
 }