Cleanup proxy stub compiler code.

Change-Id: I2a5815cc2a38f5cea9e9c36e5e6db92c1a861eaa

16 files changed, 149 insertions, 74 deletions

diff --git a/src/class_linker.cc b/src/class_linker.cc
index 3280410ce6..4133a861bf 100644
--- a/src/class_linker.cc
+++ b/src/class_linker.cc
@@ -1346,9 +1346,6 @@ void ClassLinker::LinkOatCodeFor(Method* method) {
     }
     if (method->GetInvokeStub() == NULL) {
       method->SetInvokeStub(oat_method.GetInvokeStub());
-#if defined(ART_USE_LLVM_COMPILER)
-      method->SetProxyStub(oat_method.GetProxyStub());
-#endif
     }
   }
 }
@@ -2309,7 +2306,8 @@ Method* ClassLinker::CreateProxyMethod(SirtRef<Class>& klass, SirtRef<Method>& p
 #if !defined(ART_USE_LLVM_COMPILER)
   method->SetCode(reinterpret_cast<void*>(art_proxy_invoke_handler));
 #else
-  method->SetCode(prototype->GetProxyStub());
+  OatFile::OatMethod oat_method = GetOatMethodFor(prototype.get());
+  method->SetCode(oat_method.GetProxyStub());
 #endif
 
   return method;
diff --git a/src/common_test.h b/src/common_test.h
index b03d4d6092..6b284fd621 100644
--- a/src/common_test.h
+++ b/src/common_test.h
@@ -211,6 +211,7 @@ class CommonTest : public testing::Test {
                                 static_cast<uint16_t>(-1u),
                                 static_cast<uint16_t>(-1u),
                                 static_cast<uint16_t>(-1u),
+                                static_cast<uint16_t>(-1u),
                                 static_cast<uint16_t>(-1u)
 #endif
                                 );
diff --git a/src/compiled_method.cc b/src/compiled_method.cc
index 07f331cd57..4b3cdeea60 100644
--- a/src/compiled_method.cc
+++ b/src/compiled_method.cc
@@ -180,13 +180,13 @@ const void* CompiledMethod::CodePointer(const void* code_pointer,
 }
 
 #if defined(ART_USE_LLVM_COMPILER)
-CompiledInvokeStub::CompiledInvokeStub(uint16_t elf_idx)
-    : elf_idx_(elf_idx) {
+CompiledInvokeStub::CompiledInvokeStub(uint16_t elf_idx, uint16_t elf_func_idx)
+    : elf_idx_(elf_idx), elf_func_idx_(elf_func_idx) {
 }
 #endif
 
 CompiledInvokeStub::CompiledInvokeStub(std::vector<uint8_t>& code)
-    : elf_idx_(-1) {
+    : elf_idx_(-1), elf_func_idx_(-1) {
   CHECK_NE(code.size(), 0U);
   code_ = code;
 }
diff --git a/src/compiled_method.h b/src/compiled_method.h
index f436214953..3987a5b917 100644
--- a/src/compiled_method.h
+++ b/src/compiled_method.h
@@ -121,33 +121,20 @@ class CompiledInvokeStub {
  public:
   explicit CompiledInvokeStub(std::vector<uint8_t>& code);
 #if defined(ART_USE_LLVM_COMPILER)
-  explicit CompiledInvokeStub(uint16_t elf_idx);
+  explicit CompiledInvokeStub(uint16_t elf_idx, uint16_t elf_func_idx);
 #endif
   ~CompiledInvokeStub();
 
   const std::vector<uint8_t>& GetCode() const;
 
-  uint16_t GetStubElfIndex() const {
+  uint16_t GetElfIndex() const {
     DCHECK(IsExecutableInElf());
     return elf_idx_;
   }
 
-  uint16_t GetInvokeStubElfFuncIndex() const {
-    DCHECK(IsExecutableInElf());
-    return invoke_stub_elf_func_idx_;
-  }
-
-  uint16_t GetProxyStubElfFuncIndex() const {
+  uint16_t GetElfFuncIndex() const {
     DCHECK(IsExecutableInElf());
-    return proxy_stub_elf_func_idx_;
-  }
-
-  void SetInvokeStub(uint16_t invoke_stub_elf_func_idx) {
-    invoke_stub_elf_func_idx_ = invoke_stub_elf_func_idx;
-  }
-
-  void SetProxyStub(uint16_t proxy_stub_elf_func_idx) {
-    proxy_stub_elf_func_idx_ = proxy_stub_elf_func_idx;
+    return elf_func_idx_;
   }
 
   bool IsExecutableInElf() const {
@@ -157,8 +144,7 @@ class CompiledInvokeStub {
  private:
   std::vector<uint8_t> code_;
   uint16_t elf_idx_;
-  uint16_t invoke_stub_elf_func_idx_;
-  uint16_t proxy_stub_elf_func_idx_;
+  uint16_t elf_func_idx_;
 };
 
 }  // namespace art
diff --git a/src/compiler.cc b/src/compiler.cc
index dadfc40513..5389127300 100644
--- a/src/compiler.cc
+++ b/src/compiler.cc
@@ -307,6 +307,9 @@ Compiler::Compiler(InstructionSet instruction_set, bool image, size_t thread_cou
       compiled_classes_lock_("compiled classes lock"),
       compiled_methods_lock_("compiled method lock"),
       compiled_invoke_stubs_lock_("compiled invoke stubs lock"),
+#if defined(ART_USE_LLVM_COMPILER)
+      compiled_proxy_stubs_lock_("compiled proxy stubs lock"),
+#endif
       image_(image),
       thread_count_(thread_count),
       support_debugging_(support_debugging),
@@ -350,6 +353,8 @@ Compiler::Compiler(InstructionSet instruction_set, bool image, size_t thread_cou
   create_invoke_stub_ = FindFunction<CreateInvokeStubFn>(compiler_so_name, compiler_library_, "ArtCreateInvokeStub");
 
 #if defined(ART_USE_LLVM_COMPILER)
+  create_proxy_stub_ = FindFunction<CreateProxyStubFn>(
+      compiler_so_name, compiler_library_, "ArtCreateProxyStub");
   compiler_enable_auto_elf_loading_ = FindFunction<CompilerEnableAutoElfLoadingFn>(
       compiler_so_name, compiler_library_, "compilerLLVMEnableAutoElfLoading");
   compiler_get_method_code_addr_ = FindFunction<CompilerGetMethodCodeAddrFn>(
@@ -378,6 +383,10 @@ Compiler::~Compiler() {
     STLDeleteValues(&compiled_invoke_stubs_);
   }
   {
+    MutexLock mu(compiled_proxy_stubs_lock_);
+    STLDeleteValues(&compiled_proxy_stubs_);
+  }
+  {
     MutexLock mu(compiled_methods_lock_);
     STLDeleteElements(&code_to_patch_);
   }
@@ -1505,6 +1514,18 @@ void Compiler::CompileMethod(const DexFile::CodeItem* code_item, uint32_t access
     CHECK(compiled_invoke_stub != NULL);
     InsertInvokeStub(key, compiled_invoke_stub);
   }
+
+#if defined(ART_USE_LLVM_COMPILER)
+  if (!is_static) {
+    const CompiledInvokeStub* compiled_proxy_stub = FindProxyStub(shorty);
+    if (compiled_proxy_stub == NULL) {
+      compiled_proxy_stub = (*create_proxy_stub_)(*this, shorty, shorty_len);
+      CHECK(compiled_proxy_stub != NULL);
+      InsertProxyStub(shorty, compiled_proxy_stub);
+    }
+  }
+#endif
+
   CHECK(!Thread::Current()->IsExceptionPending()) << PrettyMethod(method_idx, dex_file);
 }
 
@@ -1536,6 +1557,31 @@ void Compiler::InsertInvokeStub(const std::string& key,
   }
 }
 
+#if defined(ART_USE_LLVM_COMPILER)
+const CompiledInvokeStub* Compiler::FindProxyStub(const char* shorty) const {
+  MutexLock mu(compiled_proxy_stubs_lock_);
+  ProxyStubTable::const_iterator it = compiled_proxy_stubs_.find(shorty);
+  if (it == compiled_proxy_stubs_.end()) {
+    return NULL;
+  } else {
+    DCHECK(it->second != NULL);
+    return it->second;
+  }
+}
+
+void Compiler::InsertProxyStub(const char* shorty,
+                               const CompiledInvokeStub* compiled_proxy_stub) {
+  MutexLock mu(compiled_proxy_stubs_lock_);
+  InvokeStubTable::iterator it = compiled_proxy_stubs_.find(shorty);
+  if (it != compiled_proxy_stubs_.end()) {
+    // Someone else won the race.
+    delete compiled_proxy_stub;
+  } else {
+    compiled_proxy_stubs_.Put(shorty, compiled_proxy_stub);
+  }
+}
+#endif
+
 CompiledClass* Compiler::GetCompiledClass(ClassReference ref) const {
   MutexLock mu(compiled_classes_lock_);
   ClassTable::const_iterator it = compiled_classes_.find(ref);
diff --git a/src/compiler.h b/src/compiler.h
index c0732aab94..5ae46d5940 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -94,6 +94,10 @@ class Compiler {
   const CompiledInvokeStub* FindInvokeStub(bool is_static, const char* shorty) const;
   const CompiledInvokeStub* FindInvokeStub(const std::string& key) const;
 
+#if defined(ART_USE_LLVM_COMPILER)
+  const CompiledInvokeStub* FindProxyStub(const char* shorty) const;
+#endif
+
   // Callbacks from OAT/ART compiler to see what runtime checks must be generated
 
   bool CanAssumeTypeIsPresentInDexCache(const DexCache* dex_cache, uint32_t type_idx);
@@ -281,6 +285,10 @@ class Compiler {
 
   void InsertInvokeStub(const std::string& key, const CompiledInvokeStub* compiled_invoke_stub);
 
+#if defined(ART_USE_LLVM_COMPILER)
+  void InsertProxyStub(const char* shorty, const CompiledInvokeStub* compiled_proxy_stub);
+#endif
+
   std::vector<const PatchInformation*> code_to_patch_;
   std::vector<const PatchInformation*> methods_to_patch_;
 
@@ -301,6 +309,13 @@ class Compiler {
   mutable Mutex compiled_invoke_stubs_lock_;
   InvokeStubTable compiled_invoke_stubs_;
 
+#if defined(ART_USE_LLVM_COMPILER)
+  typedef SafeMap<std::string, const CompiledInvokeStub*> ProxyStubTable;
+  // Proxy stubs created for proxy invocation delegation
+  mutable Mutex compiled_proxy_stubs_lock_;
+  ProxyStubTable compiled_proxy_stubs_;
+#endif
+
   bool image_;
   size_t thread_count_;
   bool support_debugging_;
@@ -338,6 +353,10 @@ class Compiler {
   CreateInvokeStubFn create_invoke_stub_;
 
 #if defined(ART_USE_LLVM_COMPILER)
+  typedef CompiledInvokeStub* (*CreateProxyStubFn)
+      (Compiler& compiler, const char* shorty, uint32_t shorty_len);
+  CreateProxyStubFn create_proxy_stub_;
+
   typedef void (*CompilerEnableAutoElfLoadingFn)(Compiler& compiler);
   CompilerEnableAutoElfLoadingFn compiler_enable_auto_elf_loading_;
 
diff --git a/src/compiler_llvm/compiler_llvm.cc b/src/compiler_llvm/compiler_llvm.cc
index 95e9803f68..6c8188ef6f 100644
--- a/src/compiler_llvm/compiler_llvm.cc
+++ b/src/compiler_llvm/compiler_llvm.cc
@@ -225,8 +225,8 @@ const void* CompilerLLVM::GetMethodCodeAddr(const CompiledMethod* cm) const {
 
 const Method::InvokeStub* CompilerLLVM::
 GetMethodInvokeStubAddr(const CompiledInvokeStub* cm) const {
-  return elf_loader_->GetMethodInvokeStubAddr(cm->GetStubElfIndex(),
-                                              cm->GetInvokeStubElfFuncIndex());
+  return elf_loader_->GetMethodInvokeStubAddr(cm->GetElfIndex(),
+                                              cm->GetElfFuncIndex());
 }
 
 
@@ -282,12 +282,21 @@ CompiledInvokeStub* CompilerLLVM::CreateInvokeStub(bool is_static,
   UniquePtr<StubCompiler> stub_compiler(
     new StubCompiler(curr_cunit_, *compiler_));
 
-  CompiledInvokeStub* compiled_stub = new CompiledInvokeStub(curr_cunit_->GetElfIndex());
+  return stub_compiler->CreateInvokeStub(is_static, shorty);
+}
+
+
+CompiledInvokeStub* CompilerLLVM::CreateProxyStub(char const *shorty) {
+  MutexLock GUARD(compiler_lock_);
 
-  compiled_stub->SetInvokeStub(stub_compiler->CreateInvokeStub(is_static, shorty));
-  compiled_stub->SetProxyStub(stub_compiler->CreateProxyStub(is_static, shorty));
+  EnsureCompilationUnit();
+
+  MutexLock GUARD_CUNIT(curr_cunit_->cunit_lock_);
 
-  return compiled_stub;
+  UniquePtr<StubCompiler> stub_compiler(
+    new StubCompiler(curr_cunit_, *compiler_));
+
+  return stub_compiler->CreateProxyStub(shorty);
 }
 
 } // namespace compiler_llvm
@@ -349,14 +358,25 @@ extern "C" art::CompiledMethod* ArtJniCompileMethod(art::Compiler& compiler,
   return result;
 }
 
-extern "C" art::CompiledInvokeStub* ArtCreateInvokeStub(art::Compiler& compiler, bool is_static,
-                                                        const char* shorty, uint32_t shorty_len) {
+extern "C" art::CompiledInvokeStub* ArtCreateInvokeStub(art::Compiler& compiler,
+                                                        bool is_static,
+                                                        const char* shorty,
+                                                        uint32_t shorty_len) {
   art::compiler_llvm::CompilerLLVM* compiler_llvm = ContextOf(compiler);
   art::CompiledInvokeStub* result = compiler_llvm->CreateInvokeStub(is_static, shorty);
   compiler_llvm->MaterializeIfThresholdReached();
   return result;
 }
 
+extern "C" art::CompiledInvokeStub* ArtCreateProxyStub(art::Compiler& compiler,
+                                                       const char* shorty,
+                                                       uint32_t shorty_len) {
+  art::compiler_llvm::CompilerLLVM* compiler_llvm = ContextOf(compiler);
+  art::CompiledInvokeStub* result = compiler_llvm->CreateProxyStub(shorty);
+  compiler_llvm->MaterializeIfThresholdReached();
+  return result;
+}
+
 extern "C" void compilerLLVMSetBitcodeFileName(art::Compiler& compiler,
                                                std::string const& filename) {
   ContextOf(compiler)->SetBitcodeFileName(filename);
diff --git a/src/compiler_llvm/compiler_llvm.h b/src/compiler_llvm/compiler_llvm.h
index a6cf966416..7cc4a0e42d 100644
--- a/src/compiler_llvm/compiler_llvm.h
+++ b/src/compiler_llvm/compiler_llvm.h
@@ -104,7 +104,9 @@ class CompilerLLVM {
 
   CompiledMethod* CompileNativeMethod(OatCompilationUnit* oat_compilation_unit);
 
-  CompiledInvokeStub* CreateInvokeStub(bool is_static, char const *shorty);
+  CompiledInvokeStub* CreateInvokeStub(bool is_static, const char *shorty);
+
+  CompiledInvokeStub* CreateProxyStub(const char *shorty);
 
  private:
   void EnsureCompilationUnit();
diff --git a/src/compiler_llvm/stub_compiler.cc b/src/compiler_llvm/stub_compiler.cc
index a94a5ad333..84447b2cc8 100644
--- a/src/compiler_llvm/stub_compiler.cc
+++ b/src/compiler_llvm/stub_compiler.cc
@@ -46,11 +46,11 @@ StubCompiler::StubCompiler(CompilationUnit* cunit, Compiler& compiler)
 }
 
 
-uint16_t StubCompiler::CreateInvokeStub(bool is_static,
-                                        char const* shorty) {
+CompiledInvokeStub* StubCompiler::CreateInvokeStub(bool is_static,
+                                                   char const* shorty) {
   uint16_t elf_func_idx = cunit_->AcquireUniqueElfFuncIndex();
 
-  CHECK_NE(shorty, static_cast<char const*>(NULL));
+  CHECK(shorty != NULL);
   size_t shorty_size = strlen(shorty);
 
   // Function name
@@ -189,17 +189,12 @@ uint16_t StubCompiler::CreateInvokeStub(bool is_static,
   // store ret_addr, and ret_void.  Beside, we guess that we have to use
   // 50 bytes to represent one LLVM instruction.
 
-  return elf_func_idx;
+  return new CompiledInvokeStub(cunit_->GetElfIndex(), elf_func_idx);
 }
 
-uint16_t StubCompiler::CreateProxyStub(bool is_static,
-                                       char const* shorty) {
-  // Static method dosen't need proxy stub.
-  if (is_static) {
-    return 0;
-  }
 
-  CHECK_NE(shorty, static_cast<char const*>(NULL));
+CompiledInvokeStub* StubCompiler::CreateProxyStub(char const* shorty) {
+  CHECK(shorty != NULL);
   size_t shorty_size = strlen(shorty);
 
   uint16_t elf_func_idx = cunit_->AcquireUniqueElfFuncIndex();
@@ -209,12 +204,15 @@ uint16_t StubCompiler::CreateProxyStub(bool is_static,
 
   // Accurate function type
   llvm::Type* accurate_ret_type = irb_.getJType(shorty[0], kAccurate);
+
   std::vector<llvm::Type*> accurate_arg_types;
   accurate_arg_types.push_back(irb_.getJObjectTy()); // method
   accurate_arg_types.push_back(irb_.getJObjectTy()); // this
+
   for (size_t i = 1; i < shorty_size; ++i) {
     accurate_arg_types.push_back(irb_.getJType(shorty[i], kAccurate));
   }
+
   llvm::FunctionType* accurate_func_type =
     llvm::FunctionType::get(accurate_ret_type, accurate_arg_types, false);
 
@@ -233,13 +231,16 @@ uint16_t StubCompiler::CreateProxyStub(bool is_static,
 
   // Load actual arguments
   llvm::Function::arg_iterator arg_iter = func->arg_begin();
+
   std::vector<llvm::Value*> args;
   args.push_back(arg_iter++); // method
   args.push_back(arg_iter++); // this
   args.push_back(irb_.Runtime().EmitGetCurrentThread()); // thread
+
   for (size_t i = 1; i < shorty_size; ++i) {
     args.push_back(arg_iter++);
   }
+
   if (shorty[0] != 'V') {
     args.push_back(jvalue_temp);
   }
@@ -247,6 +248,7 @@ uint16_t StubCompiler::CreateProxyStub(bool is_static,
   // Call ProxyInvokeHandler
   // TODO: Partial inline ProxyInvokeHandler, don't use VarArg.
   irb_.CreateCall(irb_.GetRuntime(ProxyInvokeHandler), args);
+
   if (shorty[0] != 'V') {
     llvm::Value* result_addr =
         irb_.CreateBitCast(jvalue_temp, accurate_ret_type->getPointerTo());
@@ -262,7 +264,7 @@ uint16_t StubCompiler::CreateProxyStub(bool is_static,
   // Add the memory usage approximation of the compilation unit
   cunit_->AddMemUsageApproximation((shorty_size + 2) * 50);
 
-  return elf_func_idx;
+  return new CompiledInvokeStub(cunit_->GetElfIndex(), elf_func_idx);
 }
 
 
diff --git a/src/compiler_llvm/stub_compiler.h b/src/compiler_llvm/stub_compiler.h
index 4b49ec49a9..0ed626169b 100644
--- a/src/compiler_llvm/stub_compiler.h
+++ b/src/compiler_llvm/stub_compiler.h
@@ -41,8 +41,8 @@ class StubCompiler {
  public:
   StubCompiler(CompilationUnit* cunit, Compiler& compiler);
 
-  uint16_t CreateInvokeStub(bool is_static, char const* shorty);
-  uint16_t CreateProxyStub(bool is_static, char const* shorty);
+  CompiledInvokeStub* CreateInvokeStub(bool is_static, char const* shorty);
+  CompiledInvokeStub* CreateProxyStub(char const* shorty);
 
  private:
   CompilationUnit* cunit_;
diff --git a/src/oat.cc b/src/oat.cc
index d5aaf6489b..ce7c1333a4 100644
--- a/src/oat.cc
+++ b/src/oat.cc
@@ -177,6 +177,7 @@ OatMethodOffsets::OatMethodOffsets(uint32_t code_offset,
                                    uint16_t code_elf_func_idx,
                                    uint16_t invoke_stub_elf_idx,
                                    uint16_t invoke_stub_elf_func_idx,
+                                   uint16_t proxy_stub_elf_idx,
                                    uint16_t proxy_stub_elf_func_idx
 #endif
                                    )
@@ -193,6 +194,7 @@ OatMethodOffsets::OatMethodOffsets(uint32_t code_offset,
     code_elf_func_idx_(code_elf_func_idx),
     invoke_stub_elf_idx_(invoke_stub_elf_idx),
     invoke_stub_elf_func_idx_(invoke_stub_elf_func_idx),
+    proxy_stub_elf_idx_(proxy_stub_elf_idx),
     proxy_stub_elf_func_idx_(proxy_stub_elf_func_idx)
 #endif
 {}
diff --git a/src/oat.h b/src/oat.h
index c2617cf390..d1e177ff3c 100644
--- a/src/oat.h
+++ b/src/oat.h
@@ -87,6 +87,7 @@ class PACKED OatMethodOffsets {
                    uint16_t code_elf_func_idx,
                    uint16_t invoke_stub_elf_idx,
                    uint16_t invoke_stub_elf_func_idx,
+                   uint16_t proxy_stub_elf_idx,
                    uint16_t proxy_stub_elf_func_idx
 #endif
                    );
@@ -106,6 +107,7 @@ class PACKED OatMethodOffsets {
   uint16_t code_elf_func_idx_;
   uint16_t invoke_stub_elf_idx_;
   uint16_t invoke_stub_elf_func_idx_;
+  uint16_t proxy_stub_elf_idx_;
   uint16_t proxy_stub_elf_func_idx_;
 #endif
 };
diff --git a/src/oat_file.cc b/src/oat_file.cc
index 32e3b69ebc..15ad11f5ff 100644
--- a/src/oat_file.cc
+++ b/src/oat_file.cc
@@ -315,6 +315,7 @@ const OatFile::OatMethod OatFile::OatClass::GetOatMethod(uint32_t method_index)
       oat_method_offsets.code_elf_func_idx_,
       oat_method_offsets.invoke_stub_elf_idx_,
       oat_method_offsets.invoke_stub_elf_func_idx_,
+      oat_method_offsets.proxy_stub_elf_idx_,
       oat_method_offsets.proxy_stub_elf_func_idx_
 #endif
       );
@@ -335,6 +336,7 @@ OatFile::OatMethod::OatMethod(const byte* base,
                               const uint16_t code_elf_func_idx,
                               const uint16_t invoke_stub_elf_idx,
                               const uint16_t invoke_stub_elf_func_idx,
+                              const uint16_t proxy_stub_elf_idx,
                               const uint16_t proxy_stub_elf_func_idx
 #endif
                               )
@@ -353,6 +355,7 @@ OatFile::OatMethod::OatMethod(const byte* base,
     code_elf_func_idx_(code_elf_func_idx),
     invoke_stub_elf_idx_(invoke_stub_elf_idx),
     invoke_stub_elf_func_idx_(invoke_stub_elf_func_idx),
+    proxy_stub_elf_idx_(proxy_stub_elf_idx),
     proxy_stub_elf_func_idx_(proxy_stub_elf_func_idx)
 #endif
 {
@@ -452,8 +455,9 @@ const void* OatFile::OatMethod::GetProxyStub() const {
   CHECK(elf_loader_ != NULL);
   const void* stub = NULL;
   if (proxy_stub_elf_func_idx_ != static_cast<uint16_t>(-1)) {
-    stub = elf_loader_->GetMethodCodeAddr(invoke_stub_elf_idx_,
+    stub = elf_loader_->GetMethodCodeAddr(proxy_stub_elf_idx_,
                                           proxy_stub_elf_func_idx_);
+    CHECK(stub != NULL);
   }
   return stub;
 }
@@ -469,9 +473,6 @@ void OatFile::OatMethod::LinkMethodPointers(Method* method) const {
   method->SetVmapTable(GetVmapTable());
   method->SetGcMap(GetGcMap());  // Note, used by native methods in work around JNI mode.
   method->SetInvokeStub(GetInvokeStub());
-#if defined(ART_USE_LLVM_COMPILER)
-  method->SetProxyStub(GetProxyStub());
-#endif
 }
 
 void OatFile::OatMethod::LinkMethodOffsets(Method* method) const {
diff --git a/src/oat_file.h b/src/oat_file.h
index e958c4d1c2..1ad6e5d4c9 100644
--- a/src/oat_file.h
+++ b/src/oat_file.h
@@ -172,6 +172,7 @@ class OatFile {
               const uint16_t code_elf_func_idx,
               const uint16_t invoke_stub_elf_idx,
               const uint16_t invoke_stub_elf_func_idx,
+              const uint16_t proxy_stub_elf_idx,
               const uint16_t proxy_stub_elf_func_idx
 #endif
               );
@@ -203,6 +204,7 @@ class OatFile {
     uint16_t code_elf_func_idx_;
     uint16_t invoke_stub_elf_idx_;
     uint16_t invoke_stub_elf_func_idx_;
+    uint16_t proxy_stub_elf_idx_;
     uint16_t proxy_stub_elf_func_idx_;
 #endif
 
diff --git a/src/oat_writer.cc b/src/oat_writer.cc
index d90509067e..d79e61b1ff 100644
--- a/src/oat_writer.cc
+++ b/src/oat_writer.cc
@@ -273,6 +273,7 @@ size_t OatWriter::InitOatCodeMethod(size_t offset, size_t oat_class_index,
 #if defined(ART_USE_LLVM_COMPILER)
   uint16_t invoke_stub_elf_idx = static_cast<uint16_t>(-1u);
   uint16_t invoke_stub_elf_func_idx = static_cast<uint16_t>(-1u);
+  uint16_t proxy_stub_elf_idx = static_cast<uint16_t>(-1u);
   uint16_t proxy_stub_elf_func_idx = static_cast<uint16_t>(-1u);
 #endif
 
@@ -367,11 +368,8 @@ size_t OatWriter::InitOatCodeMethod(size_t offset, size_t oat_class_index,
   if (compiled_invoke_stub != NULL) {
     if (compiled_invoke_stub->IsExecutableInElf()) {
 #if defined(ART_USE_LLVM_COMPILER)
-      invoke_stub_elf_idx = compiled_invoke_stub->GetStubElfIndex();
-      invoke_stub_elf_func_idx = compiled_invoke_stub->GetInvokeStubElfFuncIndex();
-      if (!is_static) {
-        proxy_stub_elf_func_idx = compiled_invoke_stub->GetProxyStubElfFuncIndex();
-      }
+      invoke_stub_elf_idx = compiled_invoke_stub->GetElfIndex();
+      invoke_stub_elf_func_idx = compiled_invoke_stub->GetElfFuncIndex();
 #endif
     } else {
       offset = CompiledMethod::AlignCode(offset, compiler_->GetInstructionSet());
@@ -394,6 +392,17 @@ size_t OatWriter::InitOatCodeMethod(size_t offset, size_t oat_class_index,
     }
   }
 
+#if defined(ART_USE_LLVM_COMPILER)
+  if (!is_static) {
+    const CompiledInvokeStub* compiled_proxy_stub = compiler_->FindProxyStub(shorty);
+    if (compiled_proxy_stub != NULL) {
+      DCHECK(compiled_proxy_stub->IsExecutableInElf());
+      proxy_stub_elf_idx = compiled_proxy_stub->GetElfIndex();
+      proxy_stub_elf_func_idx = compiled_proxy_stub->GetElfFuncIndex();
+    }
+  }
+#endif
+
   oat_classes_[oat_class_index]->method_offsets_[class_def_method_index]
       = OatMethodOffsets(code_offset,
                          frame_size_in_bytes,
@@ -408,6 +417,7 @@ size_t OatWriter::InitOatCodeMethod(size_t offset, size_t oat_class_index,
                          code_elf_func_idx,
                          invoke_stub_elf_idx,
                          invoke_stub_elf_func_idx,
+                         proxy_stub_elf_idx,
                          proxy_stub_elf_func_idx
 #endif
                          );
diff --git a/src/object.h b/src/object.h
index b6455db99b..cbb3aa91b5 100644
--- a/src/object.h
+++ b/src/object.h
@@ -861,22 +861,6 @@ class MANAGED Method : public Object {
     return OFFSET_OF_OBJECT_MEMBER(Method, invoke_stub_);
   }
 
-#if defined(ART_USE_LLVM_COMPILER)
-  // Proxy stub entry point
-  const void* GetProxyStub() const {
-    // NOTE: LLVM doesn't use vmap table, so we reuse it for proxy stub.
-    const void* result = GetFieldPtr<const void*>(
-        OFFSET_OF_OBJECT_MEMBER(Method, vmap_table_), false);
-    // TODO: DCHECK(result != NULL);  should be ahead of time compiled
-    return result;
-  }
-
-  void SetProxyStub(const void* proxy_stub) {
-    SetFieldPtr<const void*>(OFFSET_OF_OBJECT_MEMBER(Method, vmap_table_),
-                             proxy_stub, false);
-  }
-#endif
-
   static MemberOffset GetMethodIndexOffset() {
     return OFFSET_OF_OBJECT_MEMBER(Method, method_index_);
   }