diff --git a/compiler/driver/compiler_driver.cc b/compiler/driver/compiler_driver.cc
index b3adf10..fb428b8 100644
--- a/compiler/driver/compiler_driver.cc
+++ b/compiler/driver/compiler_driver.cc
@@ -656,7 +656,7 @@
               optimizer::DexToDexCompiler::CompilationLevel::kDontDexToDexCompile) {
         DCHECK(!Runtime::Current()->UseJitCompilation());
         // TODO: add a command-line option to disable DEX-to-DEX compilation ?
-        driver->GetDexToDexCompiler().MarkForCompilation(self, method_ref, code_item);
+        driver->GetDexToDexCompiler().MarkForCompilation(self, method_ref);
       }
     }
     return compiled_method;
@@ -730,7 +730,7 @@
                      true,
                      dex_cache);
 
-  const size_t num_methods = dex_to_dex_compiler_.NumUniqueCodeItems(self);
+  const size_t num_methods = dex_to_dex_compiler_.NumCodeItemsToQuicken(self);
   if (num_methods != 0) {
     DCHECK_EQ(num_methods, 1u);
     CompileMethodDex2Dex(self,
@@ -2808,7 +2808,7 @@
     Runtime::Current()->ReclaimArenaPoolMemory();
   }
 
-  if (dex_to_dex_compiler_.NumUniqueCodeItems(Thread::Current()) > 0u) {
+  if (dex_to_dex_compiler_.NumCodeItemsToQuicken(Thread::Current()) > 0u) {
     // TODO: Not visit all of the dex files, its probably rare that only one would have quickened
     // methods though.
     for (const DexFile* dex_file : dex_files) {
@@ -2840,6 +2840,12 @@
   DCHECK(GetCompiledMethod(method_ref) != nullptr) << method_ref.PrettyMethod();
 }
 
+CompiledMethod* CompilerDriver::RemoveCompiledMethod(const MethodReference& method_ref) {
+  CompiledMethod* ret = nullptr;
+  CHECK(compiled_methods_.Remove(method_ref, &ret));
+  return ret;
+}
+
 bool CompilerDriver::GetCompiledClass(const ClassReference& ref, ClassStatus* status) const {
   DCHECK(status != nullptr);
   // The table doesn't know if something wasn't inserted. For this case it will return
@@ -3013,6 +3019,7 @@
 void CompilerDriver::SetDexFilesForOatFile(const std::vector<const DexFile*>& dex_files) {
   dex_files_for_oat_file_ = dex_files;
   compiled_classes_.AddDexFiles(dex_files);
+  dex_to_dex_compiler_.SetDexFiles(dex_files);
 }
 
 void CompilerDriver::SetClasspathDexFiles(const std::vector<const DexFile*>& dex_files) {
diff --git a/compiler/driver/compiler_driver.h b/compiler/driver/compiler_driver.h
index e32b5c4..2b524a3 100644
--- a/compiler/driver/compiler_driver.h
+++ b/compiler/driver/compiler_driver.h
@@ -165,6 +165,7 @@
   void AddCompiledMethod(const MethodReference& method_ref,
                          CompiledMethod* const compiled_method,
                          size_t non_relative_linker_patch_count);
+  CompiledMethod* RemoveCompiledMethod(const MethodReference& method_ref);
 
   void SetRequiresConstructorBarrier(Thread* self,
                                      const DexFile* dex_file,