diff --git a/runtime/oat_file_manager.cc b/runtime/oat_file_manager.cc
index 73b065f..3371a39 100644
--- a/runtime/oat_file_manager.cc
+++ b/runtime/oat_file_manager.cc
@@ -33,9 +33,10 @@
 static constexpr bool kDuplicateClassesCheck = false;
 
 const OatFile* OatFileManager::RegisterOatFile(std::unique_ptr<const OatFile> oat_file) {
-  ReaderMutexLock mu(Thread::Current(), *Locks::oat_file_manager_lock_);
+  WriterMutexLock mu(Thread::Current(), *Locks::oat_file_manager_lock_);
   DCHECK(oat_file != nullptr);
   if (kIsDebugBuild) {
+    CHECK(oat_files_.find(oat_file) == oat_files_.end());
     for (const std::unique_ptr<const OatFile>& existing : oat_files_) {
       CHECK_NE(oat_file.get(), existing.get()) << oat_file->GetLocation();
       // Check that we don't have an oat file with the same address. Copies of the same oat file
@@ -44,13 +45,29 @@
     }
   }
   have_non_pic_oat_file_ = have_non_pic_oat_file_ || !oat_file->IsPic();
-  oat_files_.push_back(std::move(oat_file));
-  return oat_files_.back().get();
+  const OatFile* ret = oat_file.get();
+  oat_files_.insert(std::move(oat_file));
+  return ret;
+}
+
+void OatFileManager::UnRegisterAndDeleteOatFile(const OatFile* oat_file) {
+  WriterMutexLock mu(Thread::Current(), *Locks::oat_file_manager_lock_);
+  DCHECK(oat_file != nullptr);
+  std::unique_ptr<const OatFile> compare(oat_file);
+  auto it = oat_files_.find(compare);
+  CHECK(it != oat_files_.end());
+  oat_files_.erase(it);
+  compare.release();
 }
 
 const OatFile* OatFileManager::FindOpenedOatFileFromOatLocation(const std::string& oat_location)
     const {
   ReaderMutexLock mu(Thread::Current(), *Locks::oat_file_manager_lock_);
+  return FindOpenedOatFileFromOatLocationLocked(oat_location);
+}
+
+const OatFile* OatFileManager::FindOpenedOatFileFromOatLocationLocked(
+    const std::string& oat_location) const {
   for (const std::unique_ptr<const OatFile>& oat_file : oat_files_) {
     if (oat_file->GetLocation() == oat_location) {
       return oat_file.get();
@@ -81,6 +98,9 @@
 }
 
 OatFileManager::~OatFileManager() {
+  // Explicitly clear oat_files_ since the OatFile destructor calls back into OatFileManager for
+  // UnRegisterOatFileLocation.
+  oat_files_.clear();
 }
 
 const OatFile* OatFileManager::RegisterImageOatFile(gc::space::ImageSpace* space) {
@@ -95,17 +115,9 @@
        current_class_index_(current_class_index),
        from_loaded_oat_(from_loaded_oat) {}
 
-  DexFileAndClassPair(DexFileAndClassPair&& rhs) {
-    *this = std::move(rhs);
-  }
+  DexFileAndClassPair(DexFileAndClassPair&& rhs) = default;
 
-  DexFileAndClassPair& operator=(DexFileAndClassPair&& rhs) {
-    cached_descriptor_ = rhs.cached_descriptor_;
-    dex_file_ = std::move(rhs.dex_file_);
-    current_class_index_ = rhs.current_class_index_;
-    from_loaded_oat_ = rhs.from_loaded_oat_;
-    return *this;
-  }
+  DexFileAndClassPair& operator=(DexFileAndClassPair&& rhs) = default;
 
   const char* GetCachedDescriptor() const {
     return cached_descriptor_;
@@ -127,6 +139,7 @@
 
   void Next() {
     ++current_class_index_;
+    cached_descriptor_ = nullptr;
   }
 
   size_t GetCurrentClassIndex() const {
@@ -253,6 +266,7 @@
 std::vector<std::unique_ptr<const DexFile>> OatFileManager::OpenDexFilesFromOat(
     const char* dex_location,
     const char* oat_location,
+    const OatFile** out_oat_file,
     std::vector<std::string>* error_msgs) {
   CHECK(dex_location != nullptr);
   CHECK(error_msgs != nullptr);
@@ -311,6 +325,7 @@
     if (accept_oat_file) {
       VLOG(class_linker) << "Registering " << oat_file->GetLocation();
       source_oat_file = RegisterOatFile(std::move(oat_file));
+      *out_oat_file = source_oat_file;
     }
   }
 
@@ -344,4 +359,26 @@
   return dex_files;
 }
 
+bool OatFileManager::RegisterOatFileLocation(const std::string& oat_location) {
+  WriterMutexLock mu(Thread::Current(), *Locks::oat_file_count_lock_);
+  auto it = oat_file_count_.find(oat_location);
+  if (it != oat_file_count_.end()) {
+    ++it->second;
+    return false;
+  }
+  oat_file_count_.insert(std::pair<std::string, size_t>(oat_location, 1u));
+  return true;
+}
+
+void OatFileManager::UnRegisterOatFileLocation(const std::string& oat_location) {
+  WriterMutexLock mu(Thread::Current(), *Locks::oat_file_count_lock_);
+  auto it = oat_file_count_.find(oat_location);
+  if (it != oat_file_count_.end()) {
+    --it->second;
+    if (it->second == 0) {
+      oat_file_count_.erase(it);
+    }
+  }
+}
+
 }  // namespace art