Merge tag 'android-7.1.1_r4' into cm-14.1cm-14.1_prerebase

Android 7.1.1 release 4

Change-Id: I14b36666b5a478024f3a9ffab90fd675a6157461

5 files changed, 545 insertions, 271 deletions

diff --git a/compiler/common_compiler_test.cc b/compiler/common_compiler_test.cc
index bf29e1c31d..bcd8940b57 100644
--- a/compiler/common_compiler_test.cc
+++ b/compiler/common_compiler_test.cc
@@ -225,6 +225,7 @@ void CommonCompilerTest::TearDown() {
   method_inliner_map_.reset();
   verification_results_.reset();
   compiler_options_.reset();
+  image_reservation_.reset();
 
   CommonRuntimeTest::TearDown();
 }
diff --git a/compiler/elf_writer_test.cc b/compiler/elf_writer_test.cc
index 449f514184..a768e0f643 100644
--- a/compiler/elf_writer_test.cc
+++ b/compiler/elf_writer_test.cc
@@ -94,7 +94,7 @@ TEST_F(ElfWriterTest, dlsym) {
                                               /*low_4gb*/false,
                                               &error_msg));
     CHECK(ef.get() != nullptr) << error_msg;
-    CHECK(ef->Load(false, /*low_4gb*/false, &error_msg)) << error_msg;
+    CHECK(ef->Load(file.get(), false, /*low_4gb*/false, &error_msg)) << error_msg;
     EXPECT_EQ(dl_oatdata, ef->FindDynamicSymbolAddress("oatdata"));
     EXPECT_EQ(dl_oatexec, ef->FindDynamicSymbolAddress("oatexec"));
     EXPECT_EQ(dl_oatlastword, ef->FindDynamicSymbolAddress("oatlastword"));
diff --git a/compiler/image_test.cc b/compiler/image_test.cc
index 91579e9daf..a68ab7cc9d 100644
--- a/compiler/image_test.cc
+++ b/compiler/image_test.cc
@@ -39,156 +39,303 @@
 
 namespace art {
 
+static const uintptr_t kRequestedImageBase = ART_BASE_ADDRESS;
+
+struct CompilationHelper {
+  std::vector<std::string> dex_file_locations;
+  std::vector<ScratchFile> image_locations;
+  std::vector<std::unique_ptr<const DexFile>> extra_dex_files;
+  std::vector<ScratchFile> image_files;
+  std::vector<ScratchFile> oat_files;
+  std::string image_dir;
+
+  void Compile(CompilerDriver* driver,
+               ImageHeader::StorageMode storage_mode);
+
+  std::vector<size_t> GetImageObjectSectionSizes();
+
+  ~CompilationHelper();
+};
+
 class ImageTest : public CommonCompilerTest {
  protected:
   virtual void SetUp() {
     ReserveImageSpace();
     CommonCompilerTest::SetUp();
   }
+
   void TestWriteRead(ImageHeader::StorageMode storage_mode);
+
+  void Compile(ImageHeader::StorageMode storage_mode,
+               CompilationHelper& out_helper,
+               const std::string& extra_dex = "",
+               const std::string& image_class = "");
+
+  std::unordered_set<std::string>* GetImageClasses() OVERRIDE {
+    return new std::unordered_set<std::string>(image_classes_);
+  }
+
+ private:
+  std::unordered_set<std::string> image_classes_;
 };
 
-void ImageTest::TestWriteRead(ImageHeader::StorageMode storage_mode) {
-  CreateCompilerDriver(Compiler::kOptimizing, kRuntimeISA, kIsTargetBuild ? 2U : 16U);
+CompilationHelper::~CompilationHelper() {
+  for (ScratchFile& image_file : image_files) {
+    image_file.Unlink();
+  }
+  for (ScratchFile& oat_file : oat_files) {
+    oat_file.Unlink();
+  }
+  const int rmdir_result = rmdir(image_dir.c_str());
+  CHECK_EQ(0, rmdir_result);
+}
 
-  // Set inline filter values.
-  compiler_options_->SetInlineDepthLimit(CompilerOptions::kDefaultInlineDepthLimit);
-  compiler_options_->SetInlineMaxCodeUnits(CompilerOptions::kDefaultInlineMaxCodeUnits);
+std::vector<size_t> CompilationHelper::GetImageObjectSectionSizes() {
+  std::vector<size_t> ret;
+  for (ScratchFile& image_file : image_files) {
+    std::unique_ptr<File> file(OS::OpenFileForReading(image_file.GetFilename().c_str()));
+    CHECK(file.get() != nullptr);
+    ImageHeader image_header;
+    CHECK_EQ(file->ReadFully(&image_header, sizeof(image_header)), true);
+    CHECK(image_header.IsValid());
+    ret.push_back(image_header.GetImageSize());
+  }
+  return ret;
+}
 
+void CompilationHelper::Compile(CompilerDriver* driver,
+                                ImageHeader::StorageMode storage_mode) {
   ClassLinker* class_linker = Runtime::Current()->GetClassLinker();
+  std::vector<const DexFile*> class_path = class_linker->GetBootClassPath();
+
+  for (const std::unique_ptr<const DexFile>& dex_file : extra_dex_files) {
+    {
+      ScopedObjectAccess soa(Thread::Current());
+      // Inject in boot class path so that the compiler driver can see it.
+      class_linker->AppendToBootClassPath(soa.Self(), *dex_file.get());
+    }
+    class_path.push_back(dex_file.get());
+  }
+
   // Enable write for dex2dex.
-  for (const DexFile* dex_file : class_linker->GetBootClassPath()) {
-    dex_file->EnableWrite();
+  for (const DexFile* dex_file : class_path) {
+    dex_file_locations.push_back(dex_file->GetLocation());
+    if (dex_file->IsReadOnly()) {
+      dex_file->EnableWrite();
+    }
   }
-  // Create a generic location tmp file, to be the base of the .art and .oat temporary files.
-  ScratchFile location;
-  ScratchFile image_location(location, ".art");
-
-  std::string image_filename(GetSystemImageFilename(image_location.GetFilename().c_str(),
-                                                    kRuntimeISA));
-  size_t pos = image_filename.rfind('/');
-  CHECK_NE(pos, std::string::npos) << image_filename;
-  std::string image_dir(image_filename, 0, pos);
-  int mkdir_result = mkdir(image_dir.c_str(), 0700);
-  CHECK_EQ(0, mkdir_result) << image_dir;
-  ScratchFile image_file(OS::CreateEmptyFile(image_filename.c_str()));
-
-  std::string oat_filename(image_filename, 0, image_filename.size() - 3);
-  oat_filename += "oat";
-  ScratchFile oat_file(OS::CreateEmptyFile(oat_filename.c_str()));
-
-  const uintptr_t requested_image_base = ART_BASE_ADDRESS;
+
+  {
+    // Create a generic tmp file, to be the base of the .art and .oat temporary files.
+    ScratchFile location;
+    for (int i = 0; i < static_cast<int>(class_path.size()); ++i) {
+      std::string cur_location(StringPrintf("%s-%d.art", location.GetFilename().c_str(), i));
+      image_locations.push_back(ScratchFile(cur_location));
+    }
+  }
+  std::vector<std::string> image_filenames;
+  for (ScratchFile& file : image_locations) {
+    std::string image_filename(GetSystemImageFilename(file.GetFilename().c_str(), kRuntimeISA));
+    image_filenames.push_back(image_filename);
+    size_t pos = image_filename.rfind('/');
+    CHECK_NE(pos, std::string::npos) << image_filename;
+    if (image_dir.empty()) {
+      image_dir = image_filename.substr(0, pos);
+      int mkdir_result = mkdir(image_dir.c_str(), 0700);
+      CHECK_EQ(0, mkdir_result) << image_dir;
+    }
+    image_files.push_back(ScratchFile(OS::CreateEmptyFile(image_filename.c_str())));
+  }
+
+  std::vector<std::string> oat_filenames;
+  for (const std::string& image_filename : image_filenames) {
+    std::string oat_filename(image_filename.substr(0, image_filename.size() - strlen("art")) + "oat");
+    oat_files.push_back(ScratchFile(OS::CreateEmptyFile(oat_filename.c_str())));
+    oat_filenames.push_back(oat_filename);
+  }
+
   std::unordered_map<const DexFile*, size_t> dex_file_to_oat_index_map;
-  std::vector<const char*> oat_filename_vector(1, oat_filename.c_str());
-  for (const DexFile* dex_file : class_linker->GetBootClassPath()) {
-    dex_file_to_oat_index_map.emplace(dex_file, 0);
+  std::vector<const char*> oat_filename_vector;
+  for (const std::string& file : oat_filenames) {
+    oat_filename_vector.push_back(file.c_str());
+  }
+  std::vector<const char*> image_filename_vector;
+  for (const std::string& file : image_filenames) {
+    image_filename_vector.push_back(file.c_str());
   }
-  std::unique_ptr<ImageWriter> writer(new ImageWriter(*compiler_driver_,
-                                                      requested_image_base,
+  size_t image_idx = 0;
+  for (const DexFile* dex_file : class_path) {
+    dex_file_to_oat_index_map.emplace(dex_file, image_idx);
+    ++image_idx;
+  }
+  // TODO: compile_pic should be a test argument.
+  std::unique_ptr<ImageWriter> writer(new ImageWriter(*driver,
+                                                      kRequestedImageBase,
                                                       /*compile_pic*/false,
                                                       /*compile_app_image*/false,
                                                       storage_mode,
                                                       oat_filename_vector,
                                                       dex_file_to_oat_index_map));
-  // TODO: compile_pic should be a test argument.
   {
     {
       jobject class_loader = nullptr;
       TimingLogger timings("ImageTest::WriteRead", false, false);
       TimingLogger::ScopedTiming t("CompileAll", &timings);
-      compiler_driver_->SetDexFilesForOatFile(class_linker->GetBootClassPath());
-      compiler_driver_->CompileAll(class_loader, class_linker->GetBootClassPath(), &timings);
+      driver->SetDexFilesForOatFile(class_path);
+      driver->CompileAll(class_loader, class_path, &timings);
 
       t.NewTiming("WriteElf");
       SafeMap<std::string, std::string> key_value_store;
-      const std::vector<const DexFile*>& dex_files = class_linker->GetBootClassPath();
-      std::unique_ptr<ElfWriter> elf_writer = CreateElfWriterQuick(
-          compiler_driver_->GetInstructionSet(),
-          compiler_driver_->GetInstructionSetFeatures(),
-          &compiler_driver_->GetCompilerOptions(),
-          oat_file.GetFile());
-      elf_writer->Start();
-      OatWriter oat_writer(/*compiling_boot_image*/true, &timings);
-      OutputStream* rodata = elf_writer->StartRoData();
-      for (const DexFile* dex_file : dex_files) {
+      std::vector<const char*> dex_filename_vector;
+      for (size_t i = 0; i < class_path.size(); ++i) {
+        dex_filename_vector.push_back("");
+      }
+      key_value_store.Put(OatHeader::kBootClassPathKey,
+                          gc::space::ImageSpace::GetMultiImageBootClassPath(
+                              dex_filename_vector,
+                              oat_filename_vector,
+                              image_filename_vector));
+
+      std::vector<std::unique_ptr<ElfWriter>> elf_writers;
+      std::vector<std::unique_ptr<OatWriter>> oat_writers;
+      for (ScratchFile& oat_file : oat_files) {
+        elf_writers.emplace_back(CreateElfWriterQuick(driver->GetInstructionSet(),
+                                                      driver->GetInstructionSetFeatures(),
+                                                      &driver->GetCompilerOptions(),
+                                                      oat_file.GetFile()));
+        elf_writers.back()->Start();
+        oat_writers.emplace_back(new OatWriter(/*compiling_boot_image*/true, &timings));
+      }
+
+      std::vector<OutputStream*> rodata;
+      std::vector<std::unique_ptr<MemMap>> opened_dex_files_map;
+      std::vector<std::unique_ptr<const DexFile>> opened_dex_files;
+      // Now that we have finalized key_value_store_, start writing the oat file.
+      for (size_t i = 0, size = oat_writers.size(); i != size; ++i) {
+        const DexFile* dex_file = class_path[i];
+        rodata.push_back(elf_writers[i]->StartRoData());
         ArrayRef<const uint8_t> raw_dex_file(
             reinterpret_cast<const uint8_t*>(&dex_file->GetHeader()),
             dex_file->GetHeader().file_size_);
-        oat_writer.AddRawDexFileSource(raw_dex_file,
-                                       dex_file->GetLocation().c_str(),
-                                       dex_file->GetLocationChecksum());
+        oat_writers[i]->AddRawDexFileSource(raw_dex_file,
+                                            dex_file->GetLocation().c_str(),
+                                            dex_file->GetLocationChecksum());
+
+        std::unique_ptr<MemMap> cur_opened_dex_files_map;
+        std::vector<std::unique_ptr<const DexFile>> cur_opened_dex_files;
+        bool dex_files_ok = oat_writers[i]->WriteAndOpenDexFiles(
+            rodata.back(),
+            oat_files[i].GetFile(),
+            driver->GetInstructionSet(),
+            driver->GetInstructionSetFeatures(),
+            &key_value_store,
+            /* verify */ false,           // Dex files may be dex-to-dex-ed, don't verify.
+            &cur_opened_dex_files_map,
+            &cur_opened_dex_files);
+        ASSERT_TRUE(dex_files_ok);
+
+        if (cur_opened_dex_files_map != nullptr) {
+          opened_dex_files_map.push_back(std::move(cur_opened_dex_files_map));
+          for (std::unique_ptr<const DexFile>& cur_dex_file : cur_opened_dex_files) {
+            // dex_file_oat_index_map_.emplace(dex_file.get(), i);
+            opened_dex_files.push_back(std::move(cur_dex_file));
+          }
+        } else {
+          ASSERT_TRUE(cur_opened_dex_files.empty());
+        }
       }
-      std::unique_ptr<MemMap> opened_dex_files_map;
-      std::vector<std::unique_ptr<const DexFile>> opened_dex_files;
-      bool dex_files_ok = oat_writer.WriteAndOpenDexFiles(
-          rodata,
-          oat_file.GetFile(),
-          compiler_driver_->GetInstructionSet(),
-          compiler_driver_->GetInstructionSetFeatures(),
-          &key_value_store,
-          /* verify */ false,           // Dex files may be dex-to-dex-ed, don't verify.
-          &opened_dex_files_map,
-          &opened_dex_files);
-      ASSERT_TRUE(dex_files_ok);
 
       bool image_space_ok = writer->PrepareImageAddressSpace();
       ASSERT_TRUE(image_space_ok);
 
-      linker::MultiOatRelativePatcher patcher(compiler_driver_->GetInstructionSet(),
-                                              instruction_set_features_.get());
-      oat_writer.PrepareLayout(compiler_driver_.get(), writer.get(), dex_files, &patcher);
-      size_t rodata_size = oat_writer.GetOatHeader().GetExecutableOffset();
-      size_t text_size = oat_writer.GetSize() - rodata_size;
-      elf_writer->SetLoadedSectionSizes(rodata_size, text_size, oat_writer.GetBssSize());
-
-      writer->UpdateOatFileLayout(/* oat_index */ 0u,
-                                  elf_writer->GetLoadedSize(),
-                                  oat_writer.GetOatDataOffset(),
-                                  oat_writer.GetSize());
-
-      bool rodata_ok = oat_writer.WriteRodata(rodata);
-      ASSERT_TRUE(rodata_ok);
-      elf_writer->EndRoData(rodata);
-
-      OutputStream* text = elf_writer->StartText();
-      bool text_ok = oat_writer.WriteCode(text);
-      ASSERT_TRUE(text_ok);
-      elf_writer->EndText(text);
-
-      bool header_ok = oat_writer.WriteHeader(elf_writer->GetStream(), 0u, 0u, 0u);
-      ASSERT_TRUE(header_ok);
-
-      writer->UpdateOatFileHeader(/* oat_index */ 0u, oat_writer.GetOatHeader());
-
-      elf_writer->WriteDynamicSection();
-      elf_writer->WriteDebugInfo(oat_writer.GetMethodDebugInfo());
-      elf_writer->WritePatchLocations(oat_writer.GetAbsolutePatchLocations());
-
-      bool success = elf_writer->End();
-      ASSERT_TRUE(success);
+      for (size_t i = 0, size = oat_files.size(); i != size; ++i) {
+        linker::MultiOatRelativePatcher patcher(driver->GetInstructionSet(),
+                                                driver->GetInstructionSetFeatures());
+        OatWriter* const oat_writer = oat_writers[i].get();
+        ElfWriter* const elf_writer = elf_writers[i].get();
+        std::vector<const DexFile*> cur_dex_files(1u, class_path[i]);
+        oat_writer->PrepareLayout(driver, writer.get(), cur_dex_files, &patcher);
+        size_t rodata_size = oat_writer->GetOatHeader().GetExecutableOffset();
+        size_t text_size = oat_writer->GetSize() - rodata_size;
+        elf_writer->SetLoadedSectionSizes(rodata_size, text_size, oat_writer->GetBssSize());
+
+        writer->UpdateOatFileLayout(i,
+                                    elf_writer->GetLoadedSize(),
+                                    oat_writer->GetOatDataOffset(),
+                                    oat_writer->GetSize());
+
+        bool rodata_ok = oat_writer->WriteRodata(rodata[i]);
+        ASSERT_TRUE(rodata_ok);
+        elf_writer->EndRoData(rodata[i]);
+
+        OutputStream* text = elf_writer->StartText();
+        bool text_ok = oat_writer->WriteCode(text);
+        ASSERT_TRUE(text_ok);
+        elf_writer->EndText(text);
+
+        bool header_ok = oat_writer->WriteHeader(elf_writer->GetStream(), 0u, 0u, 0u);
+        ASSERT_TRUE(header_ok);
+
+        writer->UpdateOatFileHeader(i, oat_writer->GetOatHeader());
+
+        elf_writer->WriteDynamicSection();
+        elf_writer->WriteDebugInfo(oat_writer->GetMethodDebugInfo());
+        elf_writer->WritePatchLocations(oat_writer->GetAbsolutePatchLocations());
+
+        bool success = elf_writer->End();
+        ASSERT_TRUE(success);
+      }
     }
-  }
-  // Workound bug that mcld::Linker::emit closes oat_file by reopening as dup_oat.
-  std::unique_ptr<File> dup_oat(OS::OpenFileReadWrite(oat_file.GetFilename().c_str()));
-  ASSERT_TRUE(dup_oat.get() != nullptr);
 
-  {
-    std::vector<const char*> dup_oat_filename(1, dup_oat->GetPath().c_str());
-    std::vector<const char*> dup_image_filename(1, image_file.GetFilename().c_str());
     bool success_image = writer->Write(kInvalidFd,
-                                       dup_image_filename,
-                                       dup_oat_filename);
+                                       image_filename_vector,
+                                       oat_filename_vector);
     ASSERT_TRUE(success_image);
-    bool success_fixup = ElfWriter::Fixup(dup_oat.get(),
-                                          writer->GetOatDataBegin(0));
-    ASSERT_TRUE(success_fixup);
 
-    ASSERT_EQ(dup_oat->FlushCloseOrErase(), 0) << "Could not flush and close oat file "
-                                               << oat_file.GetFilename();
+    for (size_t i = 0, size = oat_filenames.size(); i != size; ++i) {
+      const char* oat_filename = oat_filenames[i].c_str();
+      std::unique_ptr<File> oat_file(OS::OpenFileReadWrite(oat_filename));
+      ASSERT_TRUE(oat_file != nullptr);
+      bool success_fixup = ElfWriter::Fixup(oat_file.get(),
+                                            writer->GetOatDataBegin(i));
+      ASSERT_TRUE(success_fixup);
+      ASSERT_EQ(oat_file->FlushCloseOrErase(), 0) << "Could not flush and close oat file "
+                                                  << oat_filename;
+    }
   }
+}
 
-  uint64_t image_file_size;
-  {
+void ImageTest::Compile(ImageHeader::StorageMode storage_mode,
+                        CompilationHelper& helper,
+                        const std::string& extra_dex,
+                        const std::string& image_class) {
+  if (!image_class.empty()) {
+    image_classes_.insert(image_class);
+  }
+  CreateCompilerDriver(Compiler::kOptimizing, kRuntimeISA, kIsTargetBuild ? 2U : 16U);
+  // Set inline filter values.
+  compiler_options_->SetInlineDepthLimit(CompilerOptions::kDefaultInlineDepthLimit);
+  compiler_options_->SetInlineMaxCodeUnits(CompilerOptions::kDefaultInlineMaxCodeUnits);
+  image_classes_.clear();
+  if (!extra_dex.empty()) {
+    helper.extra_dex_files = OpenTestDexFiles(extra_dex.c_str());
+  }
+  helper.Compile(compiler_driver_.get(), storage_mode);
+  if (!image_class.empty()) {
+    // Make sure the class got initialized.
+    ScopedObjectAccess soa(Thread::Current());
+    ClassLinker* const class_linker = Runtime::Current()->GetClassLinker();
+    mirror::Class* klass = class_linker->FindSystemClass(Thread::Current(), image_class.c_str());
+    EXPECT_TRUE(klass != nullptr);
+    EXPECT_TRUE(klass->IsInitialized());
+  }
+}
+
+void ImageTest::TestWriteRead(ImageHeader::StorageMode storage_mode) {
+  CompilationHelper helper;
+  Compile(storage_mode, /*out*/ helper);
+  std::vector<uint64_t> image_file_sizes;
+  for (ScratchFile& image_file : helper.image_files) {
     std::unique_ptr<File> file(OS::OpenFileForReading(image_file.GetFilename().c_str()));
     ASSERT_TRUE(file.get() != nullptr);
     ImageHeader image_header;
@@ -205,7 +352,7 @@ void ImageTest::TestWriteRead(ImageHeader::StorageMode storage_mode) {
     ASSERT_TRUE(space != nullptr);
     ASSERT_TRUE(space->IsMallocSpace());
 
-    image_file_size = file->GetLength();
+    image_file_sizes.push_back(file->GetLength());
   }
 
   ASSERT_TRUE(compiler_driver_->GetImageClasses() != nullptr);
@@ -219,17 +366,16 @@ void ImageTest::TestWriteRead(ImageHeader::StorageMode storage_mode) {
   // Remove the reservation of the memory for use to load the image.
   // Need to do this before we reset the runtime.
   UnreserveImageSpace();
-  writer.reset(nullptr);
 
+  helper.extra_dex_files.clear();
   runtime_.reset();
   java_lang_dex_file_ = nullptr;
 
   MemMap::Init();
-  std::unique_ptr<const DexFile> dex(LoadExpectSingleDexFile(GetLibCoreDexFileNames()[0].c_str()));
 
   RuntimeOptions options;
   std::string image("-Ximage:");
-  image.append(image_location.GetFilename());
+  image.append(helper.image_locations[0].GetFilename());
   options.push_back(std::make_pair(image.c_str(), static_cast<void*>(nullptr)));
   // By default the compiler this creates will not include patch information.
   options.push_back(std::make_pair("-Xnorelocate", nullptr));
@@ -251,40 +397,45 @@ void ImageTest::TestWriteRead(ImageHeader::StorageMode storage_mode) {
   ASSERT_TRUE(heap->GetNonMovingSpace()->IsMallocSpace());
 
   // We loaded the runtime with an explicit image, so it must exist.
-  gc::space::ImageSpace* image_space = heap->GetBootImageSpaces()[0];
-  ASSERT_TRUE(image_space != nullptr);
-  if (storage_mode == ImageHeader::kStorageModeUncompressed) {
-    // Uncompressed, image should be smaller than file.
-    ASSERT_LE(image_space->Size(), image_file_size);
-  } else {
-    // Compressed, file should be smaller than image.
-    ASSERT_LE(image_file_size, image_space->Size());
-  }
+  ASSERT_EQ(heap->GetBootImageSpaces().size(), image_file_sizes.size());
+  for (size_t i = 0; i < helper.dex_file_locations.size(); ++i) {
+    std::unique_ptr<const DexFile> dex(
+        LoadExpectSingleDexFile(helper.dex_file_locations[i].c_str()));
+    ASSERT_TRUE(dex != nullptr);
+    uint64_t image_file_size = image_file_sizes[i];
+    gc::space::ImageSpace* image_space = heap->GetBootImageSpaces()[i];
+    ASSERT_TRUE(image_space != nullptr);
+    if (storage_mode == ImageHeader::kStorageModeUncompressed) {
+      // Uncompressed, image should be smaller than file.
+      ASSERT_LE(image_space->GetImageHeader().GetImageSize(), image_file_size);
+    } else if (image_file_size > 16 * KB) {
+      // Compressed, file should be smaller than image. Not really valid for small images.
+      ASSERT_LE(image_file_size, image_space->GetImageHeader().GetImageSize());
+    }
 
-  image_space->VerifyImageAllocations();
-  uint8_t* image_begin = image_space->Begin();
-  uint8_t* image_end = image_space->End();
-  CHECK_EQ(requested_image_base, reinterpret_cast<uintptr_t>(image_begin));
-  for (size_t i = 0; i < dex->NumClassDefs(); ++i) {
-    const DexFile::ClassDef& class_def = dex->GetClassDef(i);
-    const char* descriptor = dex->GetClassDescriptor(class_def);
-    mirror::Class* klass = class_linker_->FindSystemClass(soa.Self(), descriptor);
-    EXPECT_TRUE(klass != nullptr) << descriptor;
-    if (image_classes.find(descriptor) != image_classes.end()) {
-      // Image classes should be located inside the image.
-      EXPECT_LT(image_begin, reinterpret_cast<uint8_t*>(klass)) << descriptor;
-      EXPECT_LT(reinterpret_cast<uint8_t*>(klass), image_end) << descriptor;
-    } else {
-      EXPECT_TRUE(reinterpret_cast<uint8_t*>(klass) >= image_end ||
-                  reinterpret_cast<uint8_t*>(klass) < image_begin) << descriptor;
+    image_space->VerifyImageAllocations();
+    uint8_t* image_begin = image_space->Begin();
+    uint8_t* image_end = image_space->End();
+    if (i == 0) {
+      // This check is only valid for image 0.
+      CHECK_EQ(kRequestedImageBase, reinterpret_cast<uintptr_t>(image_begin));
+    }
+    for (size_t j = 0; j < dex->NumClassDefs(); ++j) {
+      const DexFile::ClassDef& class_def = dex->GetClassDef(j);
+      const char* descriptor = dex->GetClassDescriptor(class_def);
+      mirror::Class* klass = class_linker_->FindSystemClass(soa.Self(), descriptor);
+      EXPECT_TRUE(klass != nullptr) << descriptor;
+      if (image_classes.find(descriptor) == image_classes.end()) {
+        EXPECT_TRUE(reinterpret_cast<uint8_t*>(klass) >= image_end ||
+                    reinterpret_cast<uint8_t*>(klass) < image_begin) << descriptor;
+      } else {
+        // Image classes should be located inside the image.
+        EXPECT_LT(image_begin, reinterpret_cast<uint8_t*>(klass)) << descriptor;
+        EXPECT_LT(reinterpret_cast<uint8_t*>(klass), image_end) << descriptor;
+      }
+      EXPECT_TRUE(Monitor::IsValidLockWord(klass->GetLockWord(false)));
     }
-    EXPECT_TRUE(Monitor::IsValidLockWord(klass->GetLockWord(false)));
   }
-
-  image_file.Unlink();
-  oat_file.Unlink();
-  int rmdir_result = rmdir(image_dir.c_str());
-  CHECK_EQ(0, rmdir_result);
 }
 
 TEST_F(ImageTest, WriteReadUncompressed) {
@@ -299,6 +450,34 @@ TEST_F(ImageTest, WriteReadLZ4HC) {
   TestWriteRead(ImageHeader::kStorageModeLZ4HC);
 }
 
+TEST_F(ImageTest, TestImageLayout) {
+  std::vector<size_t> image_sizes;
+  std::vector<size_t> image_sizes_extra;
+  // Compile multi-image with ImageLayoutA being the last image.
+  {
+    CompilationHelper helper;
+    Compile(ImageHeader::kStorageModeUncompressed, helper, "ImageLayoutA", "LMyClass;");
+    image_sizes = helper.GetImageObjectSectionSizes();
+  }
+  TearDown();
+  runtime_.reset();
+  SetUp();
+  // Compile multi-image with ImageLayoutB being the last image.
+  {
+    CompilationHelper helper;
+    Compile(ImageHeader::kStorageModeUncompressed, helper, "ImageLayoutB", "LMyClass;");
+    image_sizes_extra = helper.GetImageObjectSectionSizes();
+  }
+  // Make sure that the new stuff in the clinit in ImageLayoutB is in the last image and not in the
+  // first two images.
+  ASSERT_EQ(image_sizes.size(), image_sizes.size());
+  // Sizes of the images should be the same. These sizes are for the whole image unrounded.
+  for (size_t i = 0; i < image_sizes.size() - 1; ++i) {
+    EXPECT_EQ(image_sizes[i], image_sizes_extra[i]);
+  }
+  // Last image should be larger since it has a hash map and a string.
+  EXPECT_LT(image_sizes.back(), image_sizes_extra.back());
+}
 
 TEST_F(ImageTest, ImageHeaderIsValid) {
     uint32_t image_begin = ART_BASE_ADDRESS;
diff --git a/compiler/image_writer.cc b/compiler/image_writer.cc
index 063eb11718..61cf00942e 100644
--- a/compiler/image_writer.cc
+++ b/compiler/image_writer.cc
@@ -387,7 +387,6 @@ void ImageWriter::SetImageBinSlot(mirror::Object* object, BinSlot bin_slot) {
   DCHECK(!IsImageBinSlotAssigned(object));
 
   // Before we stomp over the lock word, save the hash code for later.
-  Monitor::Deflate(Thread::Current(), object);;
   LockWord lw(object->GetLockWord(false));
   switch (lw.GetState()) {
     case LockWord::kFatLocked: {
@@ -488,7 +487,7 @@ void ImageWriter::AddMethodPointerArray(mirror::PointerArray* arr) {
   pointer_arrays_.emplace(arr, kBinArtMethodClean);
 }
 
-void ImageWriter::AssignImageBinSlot(mirror::Object* object) {
+void ImageWriter::AssignImageBinSlot(mirror::Object* object, size_t oat_index) {
   DCHECK(object != nullptr);
   size_t object_size = object->SizeOf();
 
@@ -591,7 +590,10 @@ void ImageWriter::AssignImageBinSlot(mirror::Object* object) {
     // else bin = kBinRegular
   }
 
-  size_t oat_index = GetOatIndex(object);
+  // Assign the oat index too.
+  DCHECK(oat_index_map_.find(object) == oat_index_map_.end());
+  oat_index_map_.emplace(object, oat_index);
+
   ImageInfo& image_info = GetImageInfo(oat_index);
 
   size_t offset_delta = RoundUp(object_size, kObjectAlignment);  // 64-bit alignment
@@ -972,39 +974,6 @@ mirror::String* ImageWriter::FindInternedString(mirror::String* string) {
   return nullptr;
 }
 
-void ImageWriter::CalculateObjectBinSlots(Object* obj) {
-  DCHECK(obj != nullptr);
-  // if it is a string, we want to intern it if its not interned.
-  if (obj->GetClass()->IsStringClass()) {
-    size_t oat_index = GetOatIndex(obj);
-    ImageInfo& image_info = GetImageInfo(oat_index);
-
-    // we must be an interned string that was forward referenced and already assigned
-    if (IsImageBinSlotAssigned(obj)) {
-      DCHECK_EQ(obj, FindInternedString(obj->AsString()));
-      return;
-    }
-    // Need to check if the string is already interned in another image info so that we don't have
-    // the intern tables of two different images contain the same string.
-    mirror::String* interned = FindInternedString(obj->AsString());
-    if (interned == nullptr) {
-      // Not in another image space, insert to our table.
-      interned = image_info.intern_table_->InternStrongImageString(obj->AsString());
-    }
-    if (obj != interned) {
-      if (!IsImageBinSlotAssigned(interned)) {
-        // interned obj is after us, allocate its location early
-        AssignImageBinSlot(interned);
-      }
-      // point those looking for this object to the interned version.
-      SetImageBinSlot(obj, GetImageBinSlot(interned));
-      return;
-    }
-    // else (obj == interned), nothing to do but fall through to the normal case
-  }
-
-  AssignImageBinSlot(obj);
-}
 
 ObjectArray<Object>* ImageWriter::CreateImageRoots(size_t oat_index) const {
   Runtime* runtime = Runtime::Current();
@@ -1090,61 +1059,33 @@ ObjectArray<Object>* ImageWriter::CreateImageRoots(size_t oat_index) const {
   return image_roots.Get();
 }
 
-// Walk instance fields of the given Class. Separate function to allow recursion on the super
-// class.
-void ImageWriter::WalkInstanceFields(mirror::Object* obj, mirror::Class* klass) {
-  // Visit fields of parent classes first.
-  StackHandleScope<1> hs(Thread::Current());
-  Handle<mirror::Class> h_class(hs.NewHandle(klass));
-  mirror::Class* super = h_class->GetSuperClass();
-  if (super != nullptr) {
-    WalkInstanceFields(obj, super);
-  }
-  //
-  size_t num_reference_fields = h_class->NumReferenceInstanceFields();
-  MemberOffset field_offset = h_class->GetFirstReferenceInstanceFieldOffset();
-  for (size_t i = 0; i < num_reference_fields; ++i) {
-    mirror::Object* value = obj->GetFieldObject<mirror::Object>(field_offset);
-    if (value != nullptr) {
-      WalkFieldsInOrder(value);
-    }
-    field_offset = MemberOffset(field_offset.Uint32Value() +
-                                sizeof(mirror::HeapReference<mirror::Object>));
-  }
-}
-
-// For an unvisited object, visit it then all its children found via fields.
-void ImageWriter::WalkFieldsInOrder(mirror::Object* obj) {
-  if (IsInBootImage(obj)) {
-    // Object is in the image, don't need to fix it up.
-    return;
+mirror::Object* ImageWriter::TryAssignBinSlot(WorkStack& work_stack,
+                                              mirror::Object* obj,
+                                              size_t oat_index) {
+  if (obj == nullptr || IsInBootImage(obj)) {
+    // Object is null or already in the image, there is no work to do.
+    return obj;
   }
-  // Use our own visitor routine (instead of GC visitor) to get better locality between
-  // an object and its fields
   if (!IsImageBinSlotAssigned(obj)) {
-    // Walk instance fields of all objects
-    StackHandleScope<2> hs(Thread::Current());
-    Handle<mirror::Object> h_obj(hs.NewHandle(obj));
-    Handle<mirror::Class> klass(hs.NewHandle(obj->GetClass()));
-    // visit the object itself.
-    CalculateObjectBinSlots(h_obj.Get());
-    WalkInstanceFields(h_obj.Get(), klass.Get());
-    // Walk static fields of a Class.
-    if (h_obj->IsClass()) {
-      size_t num_reference_static_fields = klass->NumReferenceStaticFields();
-      MemberOffset field_offset = klass->GetFirstReferenceStaticFieldOffset(target_ptr_size_);
-      for (size_t i = 0; i < num_reference_static_fields; ++i) {
-        mirror::Object* value = h_obj->GetFieldObject<mirror::Object>(field_offset);
-        if (value != nullptr) {
-          WalkFieldsInOrder(value);
-        }
-        field_offset = MemberOffset(field_offset.Uint32Value() +
-                                    sizeof(mirror::HeapReference<mirror::Object>));
+    // We want to intern all strings but also assign offsets for the source string. Since the
+    // pruning phase has already happened, if we intern a string to one in the image we still
+    // end up copying an unreachable string.
+    if (obj->IsString()) {
+      // Need to check if the string is already interned in another image info so that we don't have
+      // the intern tables of two different images contain the same string.
+      mirror::String* interned = FindInternedString(obj->AsString());
+      if (interned == nullptr) {
+        // Not in another image space, insert to our table.
+        interned = GetImageInfo(oat_index).intern_table_->InternStrongImageString(obj->AsString());
+        DCHECK_EQ(interned, obj);
       }
+    } else if (obj->IsDexCache()) {
+      oat_index = GetOatIndexForDexCache(obj->AsDexCache());
+    } else if (obj->IsClass()) {
       // Visit and assign offsets for fields and field arrays.
-      auto* as_klass = h_obj->AsClass();
+      mirror::Class* as_klass = obj->AsClass();
       mirror::DexCache* dex_cache = as_klass->GetDexCache();
-      DCHECK_NE(klass->GetStatus(), mirror::Class::kStatusError);
+      DCHECK_NE(as_klass->GetStatus(), mirror::Class::kStatusError);
       if (compile_app_image_) {
         // Extra sanity, no boot loader classes should be left!
         CHECK(!IsBootClassLoaderClass(as_klass)) << PrettyClass(as_klass);
@@ -1152,14 +1093,14 @@ void ImageWriter::WalkFieldsInOrder(mirror::Object* obj) {
       LengthPrefixedArray<ArtField>* fields[] = {
           as_klass->GetSFieldsPtr(), as_klass->GetIFieldsPtr(),
       };
-      size_t oat_index = GetOatIndexForDexCache(dex_cache);
+      // Overwrite the oat index value since the class' dex cache is more accurate of where it
+      // belongs.
+      oat_index = GetOatIndexForDexCache(dex_cache);
       ImageInfo& image_info = GetImageInfo(oat_index);
       {
-        // Note: This table is only accessed from the image writer, so the lock is technically
-        // unnecessary.
-        WriterMutexLock mu(Thread::Current(), *Locks::classlinker_classes_lock_);
-        // Insert in the class table for this iamge.
-        image_info.class_table_->Insert(as_klass);
+        // Note: This table is only accessed from the image writer, avoid locking to prevent lock
+        // order violations from root visiting.
+        image_info.class_table_->InsertWithoutLocks(as_klass);
       }
       for (LengthPrefixedArray<ArtField>* cur_fields : fields) {
         // Total array length including header.
@@ -1249,26 +1190,26 @@ void ImageWriter::WalkFieldsInOrder(mirror::Object* obj) {
         ImTable* imt = as_klass->GetImt(target_ptr_size_);
         TryAssignImTableOffset(imt, oat_index);
       }
-    } else if (h_obj->IsObjectArray()) {
-      // Walk elements of an object array.
-      int32_t length = h_obj->AsObjectArray<mirror::Object>()->GetLength();
-      for (int32_t i = 0; i < length; i++) {
-        mirror::ObjectArray<mirror::Object>* obj_array = h_obj->AsObjectArray<mirror::Object>();
-        mirror::Object* value = obj_array->Get(i);
-        if (value != nullptr) {
-          WalkFieldsInOrder(value);
-        }
-      }
-    } else if (h_obj->IsClassLoader()) {
+    } else if (obj->IsClassLoader()) {
       // Register the class loader if it has a class table.
       // The fake boot class loader should not get registered and we should end up with only one
       // class loader.
-      mirror::ClassLoader* class_loader = h_obj->AsClassLoader();
+      mirror::ClassLoader* class_loader = obj->AsClassLoader();
       if (class_loader->GetClassTable() != nullptr) {
         class_loaders_.insert(class_loader);
       }
     }
+    AssignImageBinSlot(obj, oat_index);
+    work_stack.emplace(obj, oat_index);
   }
+  if (obj->IsString()) {
+    // Always return the interned string if there exists one.
+    mirror::String* interned = FindInternedString(obj->AsString());
+    if (interned != nullptr) {
+      return interned;
+    }
+  }
+  return obj;
 }
 
 bool ImageWriter::NativeRelocationAssigned(void* ptr) const {
@@ -1325,10 +1266,16 @@ void ImageWriter::AssignMethodOffset(ArtMethod* method,
   offset += ArtMethod::Size(target_ptr_size_);
 }
 
-void ImageWriter::WalkFieldsCallback(mirror::Object* obj, void* arg) {
+void ImageWriter::EnsureBinSlotAssignedCallback(mirror::Object* obj, void* arg) {
   ImageWriter* writer = reinterpret_cast<ImageWriter*>(arg);
   DCHECK(writer != nullptr);
-  writer->WalkFieldsInOrder(obj);
+  if (!Runtime::Current()->GetHeap()->ObjectIsInBootImageSpace(obj)) {
+    CHECK(writer->IsImageBinSlotAssigned(obj)) << PrettyTypeOf(obj) << " " << obj;
+  }
+}
+
+void ImageWriter::DeflateMonitorCallback(mirror::Object* obj, void* arg ATTRIBUTE_UNUSED) {
+  Monitor::Deflate(Thread::Current(), obj);
 }
 
 void ImageWriter::UnbinObjectsIntoOffsetCallback(mirror::Object* obj, void* arg) {
@@ -1352,6 +1299,88 @@ void ImageWriter::UnbinObjectsIntoOffset(mirror::Object* obj) {
   AssignImageOffset(obj, bin_slot);
 }
 
+class ImageWriter::VisitReferencesVisitor {
+ public:
+  VisitReferencesVisitor(ImageWriter* image_writer, WorkStack* work_stack, size_t oat_index)
+      : image_writer_(image_writer), work_stack_(work_stack), oat_index_(oat_index) {}
+
+  // Fix up separately since we also need to fix up method entrypoints.
+  ALWAYS_INLINE void VisitRootIfNonNull(mirror::CompressedReference<mirror::Object>* root) const
+      SHARED_REQUIRES(Locks::mutator_lock_) {
+    if (!root->IsNull()) {
+      VisitRoot(root);
+    }
+  }
+
+  ALWAYS_INLINE void VisitRoot(mirror::CompressedReference<mirror::Object>* root) const
+      SHARED_REQUIRES(Locks::mutator_lock_) {
+    root->Assign(VisitReference(root->AsMirrorPtr()));
+  }
+
+  ALWAYS_INLINE void operator() (mirror::Object* obj,
+                                 MemberOffset offset,
+                                 bool is_static ATTRIBUTE_UNUSED) const
+      SHARED_REQUIRES(Locks::mutator_lock_) {
+    mirror::Object* ref =
+        obj->GetFieldObject<mirror::Object, kVerifyNone, kWithoutReadBarrier>(offset);
+    obj->SetFieldObject</*kTransactionActive*/false>(offset, VisitReference(ref));
+  }
+
+  ALWAYS_INLINE void operator() (mirror::Class* klass ATTRIBUTE_UNUSED,
+                                 mirror::Reference* ref) const
+      SHARED_REQUIRES(Locks::mutator_lock_) {
+    ref->SetReferent</*kTransactionActive*/false>(
+        VisitReference(ref->GetReferent<kWithoutReadBarrier>()));
+  }
+
+ private:
+  mirror::Object* VisitReference(mirror::Object* ref) const SHARED_REQUIRES(Locks::mutator_lock_) {
+    return image_writer_->TryAssignBinSlot(*work_stack_, ref, oat_index_);
+  }
+
+  ImageWriter* const image_writer_;
+  WorkStack* const work_stack_;
+  const size_t oat_index_;
+};
+
+class ImageWriter::GetRootsVisitor : public RootVisitor  {
+ public:
+  explicit GetRootsVisitor(std::vector<mirror::Object*>* roots) : roots_(roots) {}
+
+  void VisitRoots(mirror::Object*** roots,
+                  size_t count,
+                  const RootInfo& info ATTRIBUTE_UNUSED) OVERRIDE
+      SHARED_REQUIRES(Locks::mutator_lock_) {
+    for (size_t i = 0; i < count; ++i) {
+      roots_->push_back(*roots[i]);
+    }
+  }
+
+  void VisitRoots(mirror::CompressedReference<mirror::Object>** roots,
+                  size_t count,
+                  const RootInfo& info ATTRIBUTE_UNUSED) OVERRIDE
+      SHARED_REQUIRES(Locks::mutator_lock_) {
+    for (size_t i = 0; i < count; ++i) {
+      roots_->push_back(roots[i]->AsMirrorPtr());
+    }
+  }
+
+ private:
+  std::vector<mirror::Object*>* const roots_;
+};
+
+void ImageWriter::ProcessWorkStack(WorkStack* work_stack) {
+  while (!work_stack->empty()) {
+    std::pair<mirror::Object*, size_t> pair(work_stack->top());
+    work_stack->pop();
+    VisitReferencesVisitor visitor(this, work_stack, /*oat_index*/ pair.second);
+    // Walk references and assign bin slots for them.
+    pair.first->VisitReferences</*kVisitNativeRoots*/true, kVerifyNone, kWithoutReadBarrier>(
+        visitor,
+        visitor);
+  }
+}
+
 void ImageWriter::CalculateNewObjectOffsets() {
   Thread* const self = Thread::Current();
   StackHandleScopeCollection handles(self);
@@ -1360,8 +1389,8 @@ void ImageWriter::CalculateNewObjectOffsets() {
     image_roots.push_back(handles.NewHandle(CreateImageRoots(i)));
   }
 
-  auto* runtime = Runtime::Current();
-  auto* heap = runtime->GetHeap();
+  Runtime* const runtime = Runtime::Current();
+  gc::Heap* const heap = runtime->GetHeap();
 
   // Leave space for the header, but do not write it yet, we need to
   // know where image_roots is going to end up
@@ -1387,8 +1416,64 @@ void ImageWriter::CalculateNewObjectOffsets() {
     }
   }
 
-  // Clear any pre-existing monitors which may have been in the monitor words, assign bin slots.
-  heap->VisitObjects(WalkFieldsCallback, this);
+  // Deflate monitors before we visit roots since deflating acquires the monitor lock. Acquiring
+  // this lock while holding other locks may cause lock order violations.
+  heap->VisitObjects(DeflateMonitorCallback, this);
+
+  // Work list of <object, oat_index> for objects. Everything on the stack must already be
+  // assigned a bin slot.
+  WorkStack work_stack;
+
+  // Special case interned strings to put them in the image they are likely to be resolved from.
+  for (const DexFile* dex_file : compiler_driver_.GetDexFilesForOatFile()) {
+    auto it = dex_file_oat_index_map_.find(dex_file);
+    DCHECK(it != dex_file_oat_index_map_.end()) << dex_file->GetLocation();
+    const size_t oat_index = it->second;
+    InternTable* const intern_table = runtime->GetInternTable();
+    for (size_t i = 0, count = dex_file->NumStringIds(); i < count; ++i) {
+      uint32_t utf16_length;
+      const char* utf8_data = dex_file->StringDataAndUtf16LengthByIdx(i, &utf16_length);
+      mirror::String* string = intern_table->LookupStrong(self, utf16_length, utf8_data);
+      TryAssignBinSlot(work_stack, string, oat_index);
+    }
+  }
+
+  // Get the GC roots and then visit them separately to avoid lock violations since the root visitor
+  // visits roots while holding various locks.
+  {
+    std::vector<mirror::Object*> roots;
+    GetRootsVisitor root_visitor(&roots);
+    runtime->VisitRoots(&root_visitor);
+    for (mirror::Object* obj : roots) {
+      TryAssignBinSlot(work_stack, obj, GetDefaultOatIndex());
+    }
+  }
+  ProcessWorkStack(&work_stack);
+
+  // For app images, there may be objects that are only held live by the by the boot image. One
+  // example is finalizer references. Forward these objects so that EnsureBinSlotAssignedCallback
+  // does not fail any checks. TODO: We should probably avoid copying these objects.
+  if (compile_app_image_) {
+    for (gc::space::ImageSpace* space : heap->GetBootImageSpaces()) {
+      DCHECK(space->IsImageSpace());
+      gc::accounting::ContinuousSpaceBitmap* live_bitmap = space->GetLiveBitmap();
+      live_bitmap->VisitMarkedRange(reinterpret_cast<uintptr_t>(space->Begin()),
+                                    reinterpret_cast<uintptr_t>(space->Limit()),
+                                    [this, &work_stack](mirror::Object* obj)
+          SHARED_REQUIRES(Locks::mutator_lock_) {
+        VisitReferencesVisitor visitor(this, &work_stack, GetDefaultOatIndex());
+        // Visit all references and try to assign bin slots for them (calls TryAssignBinSlot).
+        obj->VisitReferences</*kVisitNativeRoots*/true, kVerifyNone, kWithoutReadBarrier>(
+            visitor,
+            visitor);
+      });
+    }
+    // Process the work stack in case anything was added by TryAssignBinSlot.
+    ProcessWorkStack(&work_stack);
+  }
+
+  // Verify that all objects have assigned image bin slots.
+  heap->VisitObjects(EnsureBinSlotAssignedCallback, this);
 
   // Calculate size of the dex cache arrays slot and prepare offsets.
   PrepareDexCacheArraySlots();
@@ -2272,25 +2357,21 @@ ImageWriter::Bin ImageWriter::BinTypeForNativeRelocationType(NativeObjectRelocat
 }
 
 size_t ImageWriter::GetOatIndex(mirror::Object* obj) const {
-  if (compile_app_image_) {
+  if (!IsMultiImage()) {
     return GetDefaultOatIndex();
-  } else {
-    mirror::DexCache* dex_cache =
-        obj->IsDexCache() ? obj->AsDexCache()
-                          : obj->IsClass() ? obj->AsClass()->GetDexCache()
-                                           : obj->GetClass()->GetDexCache();
-    return GetOatIndexForDexCache(dex_cache);
   }
+  auto it = oat_index_map_.find(obj);
+  DCHECK(it != oat_index_map_.end());
+  return it->second;
 }
 
 size_t ImageWriter::GetOatIndexForDexFile(const DexFile* dex_file) const {
-  if (compile_app_image_) {
+  if (!IsMultiImage()) {
     return GetDefaultOatIndex();
-  } else {
-    auto it = dex_file_oat_index_map_.find(dex_file);
-    DCHECK(it != dex_file_oat_index_map_.end()) << dex_file->GetLocation();
-    return it->second;
   }
+  auto it = dex_file_oat_index_map_.find(dex_file);
+  DCHECK(it != dex_file_oat_index_map_.end()) << dex_file->GetLocation();
+  return it->second;
 }
 
 size_t ImageWriter::GetOatIndexForDexCache(mirror::DexCache* dex_cache) const {
diff --git a/compiler/image_writer.h b/compiler/image_writer.h
index 1efdc22c0a..37f108f660 100644
--- a/compiler/image_writer.h
+++ b/compiler/image_writer.h
@@ -23,6 +23,7 @@
 #include <cstddef>
 #include <memory>
 #include <set>
+#include <stack>
 #include <string>
 #include <ostream>
 
@@ -143,6 +144,8 @@ class ImageWriter FINAL {
   void UpdateOatFileHeader(size_t oat_index, const OatHeader& oat_header);
 
  private:
+  using WorkStack = std::stack<std::pair<mirror::Object*, size_t>>;
+
   bool AllocMemory();
 
   // Mark the objects defined in this space in the given live bitmap.
@@ -321,7 +324,10 @@ class ImageWriter FINAL {
       SHARED_REQUIRES(Locks::mutator_lock_);
 
   void PrepareDexCacheArraySlots() SHARED_REQUIRES(Locks::mutator_lock_);
-  void AssignImageBinSlot(mirror::Object* object) SHARED_REQUIRES(Locks::mutator_lock_);
+  void AssignImageBinSlot(mirror::Object* object, size_t oat_index)
+      SHARED_REQUIRES(Locks::mutator_lock_);
+  mirror::Object* TryAssignBinSlot(WorkStack& work_stack, mirror::Object* obj, size_t oat_index)
+      SHARED_REQUIRES(Locks::mutator_lock_);
   void SetImageBinSlot(mirror::Object* object, BinSlot bin_slot)
       SHARED_REQUIRES(Locks::mutator_lock_);
   bool IsImageBinSlotAssigned(mirror::Object* object) const
@@ -378,20 +384,18 @@ class ImageWriter FINAL {
   // Lays out where the image objects will be at runtime.
   void CalculateNewObjectOffsets()
       SHARED_REQUIRES(Locks::mutator_lock_);
+  void ProcessWorkStack(WorkStack* work_stack)
+      SHARED_REQUIRES(Locks::mutator_lock_);
   void CreateHeader(size_t oat_index)
       SHARED_REQUIRES(Locks::mutator_lock_);
   mirror::ObjectArray<mirror::Object>* CreateImageRoots(size_t oat_index) const
       SHARED_REQUIRES(Locks::mutator_lock_);
-  void CalculateObjectBinSlots(mirror::Object* obj)
-      SHARED_REQUIRES(Locks::mutator_lock_);
   void UnbinObjectsIntoOffset(mirror::Object* obj)
       SHARED_REQUIRES(Locks::mutator_lock_);
 
-  void WalkInstanceFields(mirror::Object* obj, mirror::Class* klass)
+  static void EnsureBinSlotAssignedCallback(mirror::Object* obj, void* arg)
       SHARED_REQUIRES(Locks::mutator_lock_);
-  void WalkFieldsInOrder(mirror::Object* obj)
-      SHARED_REQUIRES(Locks::mutator_lock_);
-  static void WalkFieldsCallback(mirror::Object* obj, void* arg)
+  static void DeflateMonitorCallback(mirror::Object* obj, void* arg)
       SHARED_REQUIRES(Locks::mutator_lock_);
   static void UnbinObjectsIntoOffsetCallback(mirror::Object* obj, void* arg)
       SHARED_REQUIRES(Locks::mutator_lock_);
@@ -461,6 +465,10 @@ class ImageWriter FINAL {
                                   std::unordered_set<mirror::Class*>* visited)
       SHARED_REQUIRES(Locks::mutator_lock_);
 
+  bool IsMultiImage() const {
+    return image_infos_.size() > 1;
+  }
+
   static Bin BinTypeForNativeRelocationType(NativeObjectRelocationType type);
 
   uintptr_t NativeOffsetInImage(void* obj) SHARED_REQUIRES(Locks::mutator_lock_);
@@ -519,6 +527,9 @@ class ImageWriter FINAL {
   // forwarding addresses as well as copying over hash codes.
   std::unordered_map<mirror::Object*, uint32_t> saved_hashcode_map_;
 
+  // Oat index map for objects.
+  std::unordered_map<mirror::Object*, uint32_t> oat_index_map_;
+
   // Boolean flags.
   const bool compile_pic_;
   const bool compile_app_image_;
@@ -573,8 +584,10 @@ class ImageWriter FINAL {
   friend class FixupClassVisitor;
   friend class FixupRootVisitor;
   friend class FixupVisitor;
+  class GetRootsVisitor;
   friend class NativeLocationVisitor;
   friend class NonImageClassesVisitor;
+  class VisitReferencesVisitor;
   DISALLOW_COPY_AND_ASSIGN(ImageWriter);
 };