Revert "Refactor oat file writing."

This reverts commit 0d192d2c08116cab8244a6eb229aa784bc6cfefa.

Reason for revert: Boot time regression.
Bug: 142680736

Change-Id: I65c8f687668c1d7cec293b6b8bfe7147181c0a10
diff --git a/dex2oat/dex2oat.cc b/dex2oat/dex2oat.cc
index 2f66f3c..91d6513 100644
--- a/dex2oat/dex2oat.cc
+++ b/dex2oat/dex2oat.cc
@@ -1448,51 +1448,6 @@
       return dex2oat::ReturnCode::kOther;
     }
 
-    {
-      TimingLogger::ScopedTiming t_dex("Writing and opening dex files", timings_);
-      for (size_t i = 0, size = oat_writers_.size(); i != size; ++i) {
-        // Unzip or copy dex files straight to the oat file.
-        std::vector<MemMap> opened_dex_files_map;
-        std::vector<std::unique_ptr<const DexFile>> opened_dex_files;
-        // No need to verify the dex file when we have a vdex file, which means it was already
-        // verified.
-        const bool verify = (input_vdex_file_ == nullptr);
-        if (!oat_writers_[i]->WriteAndOpenDexFiles(
-            vdex_files_[i].get(),
-            verify,
-            update_input_vdex_,
-            copy_dex_files_,
-            &opened_dex_files_map,
-            &opened_dex_files)) {
-          return dex2oat::ReturnCode::kOther;
-        }
-        dex_files_per_oat_file_.push_back(MakeNonOwningPointerVector(opened_dex_files));
-        if (opened_dex_files_map.empty()) {
-          DCHECK(opened_dex_files.empty());
-        } else {
-          for (MemMap& map : opened_dex_files_map) {
-            opened_dex_files_maps_.push_back(std::move(map));
-          }
-          for (std::unique_ptr<const DexFile>& dex_file : opened_dex_files) {
-            dex_file_oat_index_map_.emplace(dex_file.get(), i);
-            opened_dex_files_.push_back(std::move(dex_file));
-          }
-        }
-      }
-    }
-
-    compiler_options_->dex_files_for_oat_file_ = MakeNonOwningPointerVector(opened_dex_files_);
-    const std::vector<const DexFile*>& dex_files = compiler_options_->dex_files_for_oat_file_;
-
-    if (IsBootImage()) {
-      // For boot image, pass opened dex files to the Runtime::Create().
-      // Note: Runtime acquires ownership of these dex files.
-      runtime_options.Set(RuntimeArgumentMap::BootClassPathDexList, &opened_dex_files_);
-    }
-    if (!CreateRuntime(std::move(runtime_options))) {
-      return dex2oat::ReturnCode::kCreateRuntime;
-    }
-
     if (!compilation_reason_.empty()) {
       key_value_store_->Put(OatHeader::kCompilationReasonKey, compilation_reason_);
     }
@@ -1504,6 +1459,12 @@
     }
 
     if (!IsBootImage()) {
+      // When compiling an app, create the runtime early to retrieve
+      // the boot image checksums needed for the oat header.
+      if (!CreateRuntime(std::move(runtime_options))) {
+        return dex2oat::ReturnCode::kCreateRuntime;
+      }
+
       if (CompilerFilter::DependsOnImageChecksum(compiler_options_->GetCompilerFilter())) {
         TimingLogger::ScopedTiming t3("Loading image checksum", timings_);
         Runtime* runtime = Runtime::Current();
@@ -1557,6 +1518,47 @@
       key_value_store_->Put(OatHeader::kClassPathKey, class_path_key);
     }
 
+    // Now that we have finalized key_value_store_, start writing the oat file.
+    {
+      TimingLogger::ScopedTiming t_dex("Writing and opening dex files", timings_);
+      rodata_.reserve(oat_writers_.size());
+      for (size_t i = 0, size = oat_writers_.size(); i != size; ++i) {
+        rodata_.push_back(elf_writers_[i]->StartRoData());
+        // Unzip or copy dex files straight to the oat file.
+        std::vector<MemMap> opened_dex_files_map;
+        std::vector<std::unique_ptr<const DexFile>> opened_dex_files;
+        // No need to verify the dex file when we have a vdex file, which means it was already
+        // verified.
+        const bool verify = (input_vdex_file_ == nullptr);
+        if (!oat_writers_[i]->WriteAndOpenDexFiles(
+            vdex_files_[i].get(),
+            rodata_.back(),
+            (i == 0u) ? key_value_store_.get() : nullptr,
+            verify,
+            update_input_vdex_,
+            copy_dex_files_,
+            &opened_dex_files_map,
+            &opened_dex_files)) {
+          return dex2oat::ReturnCode::kOther;
+        }
+        dex_files_per_oat_file_.push_back(MakeNonOwningPointerVector(opened_dex_files));
+        if (opened_dex_files_map.empty()) {
+          DCHECK(opened_dex_files.empty());
+        } else {
+          for (MemMap& map : opened_dex_files_map) {
+            opened_dex_files_maps_.push_back(std::move(map));
+          }
+          for (std::unique_ptr<const DexFile>& dex_file : opened_dex_files) {
+            dex_file_oat_index_map_.emplace(dex_file.get(), i);
+            opened_dex_files_.push_back(std::move(dex_file));
+          }
+        }
+      }
+    }
+
+    compiler_options_->dex_files_for_oat_file_ = MakeNonOwningPointerVector(opened_dex_files_);
+    const std::vector<const DexFile*>& dex_files = compiler_options_->dex_files_for_oat_file_;
+
     // If we need to downgrade the compiler-filter for size reasons.
     if (!IsBootImage() && IsVeryLarge(dex_files)) {
       // Disable app image to make sure dex2oat unloading is enabled.
@@ -1600,6 +1602,14 @@
       }
     }
     // Note that dex2oat won't close the swap_fd_. The compiler driver's swap space will do that.
+    if (IsBootImage()) {
+      // For boot image, pass opened dex files to the Runtime::Create().
+      // Note: Runtime acquires ownership of these dex files.
+      runtime_options.Set(RuntimeArgumentMap::BootClassPathDexList, &opened_dex_files_);
+      if (!CreateRuntime(std::move(runtime_options))) {
+        return dex2oat::ReturnCode::kOther;
+      }
+    }
 
     // If we're doing the image, override the compiler filter to force full compilation. Must be
     // done ahead of WellKnownClasses::Init that causes verification.  Note: doesn't force
@@ -1924,21 +1934,12 @@
       }
     }
 
-    // Initialize the writers with the compiler driver, image writer and their dex files
-    // and start writing the .rodata sections.
-    {
-      TimingLogger::ScopedTiming t2("dex2oat Starting oat .rodata section(s)", timings_);
-      rodata_.reserve(oat_writers_.size());
-      for (size_t i = 0, size = oat_files_.size(); i != size; ++i) {
-        rodata_.push_back(elf_writers_[i]->StartRoData());
-        if (!oat_writers_[i]->StartRoData(driver_.get(),
-                                          image_writer_.get(),
-                                          dex_files_per_oat_file_[i],
-                                          rodata_.back(),
-                                          (i == 0u) ? key_value_store_.get() : nullptr)) {
-          return false;
-        }
-      }
+    // Initialize the writers with the compiler driver, image writer, and their
+    // dex files. The writers were created without those being there yet.
+    for (size_t i = 0, size = oat_files_.size(); i != size; ++i) {
+      std::unique_ptr<linker::OatWriter>& oat_writer = oat_writers_[i];
+      std::vector<const DexFile*>& dex_files = dex_files_per_oat_file_[i];
+      oat_writer->Initialize(driver_.get(), image_writer_.get(), dex_files);
     }
 
     {
@@ -2007,7 +2008,7 @@
         debug::DebugInfo debug_info = oat_writer->GetDebugInfo();  // Keep the variable alive.
         elf_writer->PrepareDebugInfo(debug_info);  // Processes the data on background thread.
 
-        OutputStream* rodata = rodata_[i];
+        OutputStream*& rodata = rodata_[i];
         DCHECK(rodata != nullptr);
         if (!oat_writer->WriteRodata(rodata)) {
           LOG(ERROR) << "Failed to write .rodata section to the ELF file " << oat_file->GetPath();
diff --git a/dex2oat/linker/image_test.h b/dex2oat/linker/image_test.h
index d1d22a5..4c984d3 100644
--- a/dex2oat/linker/image_test.h
+++ b/dex2oat/linker/image_test.h
@@ -257,6 +257,8 @@
         std::vector<std::unique_ptr<const DexFile>> cur_opened_dex_files;
         bool dex_files_ok = oat_writers[i]->WriteAndOpenDexFiles(
             out_helper.vdex_files[i].GetFile(),
+            rodata.back(),
+            (i == 0u) ? &key_value_store : nullptr,
             /* verify */ false,           // Dex files may be dex-to-dex-ed, don't verify.
             /* update_input_vdex */ false,
             /* copy_dex_files */ CopyOption::kOnlyIfCompressed,
@@ -287,12 +289,7 @@
         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]);
-        bool initialize_ok = oat_writer->StartRoData(driver,
-                                                     writer.get(),
-                                                     cur_dex_files,
-                                                     rodata[i],
-                                                     (i == 0u) ? &key_value_store : nullptr);
-        ASSERT_TRUE(initialize_ok);
+        oat_writer->Initialize(driver, writer.get(), cur_dex_files);
 
         std::unique_ptr<BufferedOutputStream> vdex_out =
             std::make_unique<BufferedOutputStream>(
diff --git a/dex2oat/linker/oat_writer.cc b/dex2oat/linker/oat_writer.cc
index 2c1837e..b478e28 100644
--- a/dex2oat/linker/oat_writer.cc
+++ b/dex2oat/linker/oat_writer.cc
@@ -658,6 +658,8 @@
 
 bool OatWriter::WriteAndOpenDexFiles(
     File* vdex_file,
+    OutputStream* oat_rodata,
+    SafeMap<std::string, std::string>* key_value_store,
     bool verify,
     bool update_input_vdex,
     CopyOption copy_dex_files,
@@ -665,9 +667,23 @@
     /*out*/ std::vector<std::unique_ptr<const DexFile>>* opened_dex_files) {
   CHECK(write_state_ == WriteState::kAddingDexFileSources);
 
+  // Record the ELF rodata section offset, i.e. the beginning of the OAT data.
+  if (!RecordOatDataOffset(oat_rodata)) {
+     return false;
+  }
+
+  // Record whether this is the primary oat file.
+  primary_oat_file_ = (key_value_store != nullptr);
+
+  // Initialize VDEX and OAT headers.
+
   // Reserve space for Vdex header and checksums.
   vdex_size_ = sizeof(VdexFile::VerifierDepsHeader) +
       oat_dex_files_.size() * sizeof(VdexFile::VdexChecksum);
+  oat_size_ = InitOatHeader(dchecked_integral_cast<uint32_t>(oat_dex_files_.size()),
+                            key_value_store);
+
+  ChecksumUpdatingOutputStream checksum_updating_rodata(oat_rodata, this);
 
   std::unique_ptr<BufferedOutputStream> vdex_out =
       std::make_unique<BufferedOutputStream>(std::make_unique<FileOutputStream>(vdex_file));
@@ -679,37 +695,6 @@
     return false;
   }
 
-  *opened_dex_files_map = std::move(dex_files_map);
-  *opened_dex_files = std::move(dex_files);
-  write_state_ = WriteState::kStartRoData;
-  return true;
-}
-
-// Initialize the writer with the given parameters.
-bool OatWriter::StartRoData(const CompilerDriver* compiler_driver,
-                            ImageWriter* image_writer,
-                            const std::vector<const DexFile*>& dex_files,
-                            OutputStream* oat_rodata,
-                            SafeMap<std::string, std::string>* key_value_store) {
-  CHECK(write_state_ == WriteState::kStartRoData);
-  compiler_driver_ = compiler_driver;
-  image_writer_ = image_writer;
-  dex_files_ = &dex_files;
-
-  // Record the ELF rodata section offset, i.e. the beginning of the OAT data.
-  if (!RecordOatDataOffset(oat_rodata)) {
-     return false;
-  }
-
-  // Record whether this is the primary oat file.
-  primary_oat_file_ = (key_value_store != nullptr);
-
-  // Initialize OAT header.
-  oat_size_ = InitOatHeader(dchecked_integral_cast<uint32_t>(oat_dex_files_.size()),
-                            key_value_store);
-
-  ChecksumUpdatingOutputStream checksum_updating_rodata(oat_rodata, this);
-
   // Write type lookup tables into the oat file.
   if (!WriteTypeLookupTables(&checksum_updating_rodata, dex_files)) {
     return false;
@@ -720,10 +705,21 @@
     return false;
   }
 
+  *opened_dex_files_map = std::move(dex_files_map);
+  *opened_dex_files = std::move(dex_files);
   write_state_ = WriteState::kPrepareLayout;
   return true;
 }
 
+// Initialize the writer with the given parameters.
+void OatWriter::Initialize(const CompilerDriver* compiler_driver,
+                           ImageWriter* image_writer,
+                           const std::vector<const DexFile*>& dex_files) {
+  compiler_driver_ = compiler_driver;
+  image_writer_ = image_writer;
+  dex_files_ = &dex_files;
+}
+
 void OatWriter::PrepareLayout(MultiOatRelativePatcher* relative_patcher) {
   CHECK(write_state_ == WriteState::kPrepareLayout);
 
@@ -3772,8 +3768,9 @@
   return true;
 }
 
-bool OatWriter::WriteTypeLookupTables(OutputStream* oat_rodata,
-                                      const std::vector<const DexFile*>& opened_dex_files) {
+bool OatWriter::WriteTypeLookupTables(
+    OutputStream* oat_rodata,
+    const std::vector<std::unique_ptr<const DexFile>>& opened_dex_files) {
   TimingLogger::ScopedTiming split("WriteTypeLookupTables", timings_);
 
   uint32_t expected_offset = oat_data_offset_ + oat_size_;
@@ -3855,11 +3852,11 @@
 
 bool OatWriter::WriteDexLayoutSections(
     OutputStream* oat_rodata,
-    const std::vector<const DexFile*>& opened_dex_files) {
+    const std::vector<std::unique_ptr<const DexFile>>& opened_dex_files) {
   TimingLogger::ScopedTiming split(__FUNCTION__, timings_);
 
   if (!kWriteDexLayoutInfo) {
-    return true;
+    return true;;
   }
 
   uint32_t expected_offset = oat_data_offset_ + oat_size_;
diff --git a/dex2oat/linker/oat_writer.h b/dex2oat/linker/oat_writer.h
index c95b4dd..84d13a8 100644
--- a/dex2oat/linker/oat_writer.h
+++ b/dex2oat/linker/oat_writer.h
@@ -132,7 +132,7 @@
   //   - AddVdexDexFilesSource().
   // Then the user must call in order
   //   - WriteAndOpenDexFiles()
-  //   - StartRoData()
+  //   - Initialize()
   //   - WriteVerifierDeps()
   //   - WriteQuickeningInfo()
   //   - WriteChecksumsAndVdexHeader()
@@ -167,23 +167,23 @@
   dchecked_vector<std::string> GetSourceLocations() const;
 
   // Write raw dex files to the vdex file, mmap the file and open the dex files from it.
+  // Supporting data structures are written into the .rodata section of the oat file.
   // The `verify` setting dictates whether the dex file verifier should check the dex files.
   // This is generally the case, and should only be false for tests.
   // If `update_input_vdex` is true, then this method won't actually write the dex files,
   // and the compiler will just re-use the existing vdex file.
   bool WriteAndOpenDexFiles(File* vdex_file,
+                            OutputStream* oat_rodata,
+                            SafeMap<std::string, std::string>* key_value_store,
                             bool verify,
                             bool update_input_vdex,
                             CopyOption copy_dex_files,
                             /*out*/ std::vector<MemMap>* opened_dex_files_map,
                             /*out*/ std::vector<std::unique_ptr<const DexFile>>* opened_dex_files);
-  // Initialize the writer with the given parameters and start writing .rodata.
-  // Supporting data structures for dex files are written into the .rodata section of the oat file.
-  bool StartRoData(const CompilerDriver* compiler_driver,
-                   ImageWriter* image_writer,
-                   const std::vector<const DexFile*>& dex_files,
-                   OutputStream* oat_rodata,
-                   SafeMap<std::string, std::string>* key_value_store);
+  // Initialize the writer with the given parameters.
+  void Initialize(const CompilerDriver* compiler_driver,
+                  ImageWriter* image_writer,
+                  const std::vector<const DexFile*>& dex_files);
   bool WriteQuickeningInfo(OutputStream* vdex_out);
   bool WriteVerifierDeps(OutputStream* vdex_out, verifier::VerifierDeps* verifier_deps);
   bool WriteChecksumsAndVdexHeader(OutputStream* vdex_out);
@@ -339,9 +339,9 @@
 
   bool RecordOatDataOffset(OutputStream* out);
   bool WriteTypeLookupTables(OutputStream* oat_rodata,
-                             const std::vector<const DexFile*>& opened_dex_files);
+                             const std::vector<std::unique_ptr<const DexFile>>& opened_dex_files);
   bool WriteDexLayoutSections(OutputStream* oat_rodata,
-                              const std::vector<const DexFile*>& opened_dex_files);
+                              const std::vector<std::unique_ptr<const DexFile>>& opened_dex_files);
   bool WriteCodeAlignment(OutputStream* out, uint32_t aligned_code_delta);
   bool WriteUpTo16BytesAlignment(OutputStream* out, uint32_t size, uint32_t* stat);
   void SetMultiOatRelativePatcherAdjustment();
@@ -355,7 +355,6 @@
 
   enum class WriteState {
     kAddingDexFileSources,
-    kStartRoData,
     kPrepareLayout,
     kWriteRoData,
     kWriteText,
diff --git a/dex2oat/linker/oat_writer_test.cc b/dex2oat/linker/oat_writer_test.cc
index 0a92baf..2142727 100644
--- a/dex2oat/linker/oat_writer_test.cc
+++ b/dex2oat/linker/oat_writer_test.cc
@@ -178,6 +178,8 @@
     std::vector<std::unique_ptr<const DexFile>> opened_dex_files;
     if (!oat_writer.WriteAndOpenDexFiles(
         vdex_file,
+        oat_rodata,
+        &key_value_store,
         verify,
         /*update_input_vdex=*/ false,
         copy,
@@ -197,13 +199,7 @@
     MultiOatRelativePatcher patcher(compiler_options_->GetInstructionSet(),
                                     compiler_options_->GetInstructionSetFeatures(),
                                     compiler_driver_->GetCompiledMethodStorage());
-    if (!oat_writer.StartRoData(compiler_driver_.get(),
-                                /*image_writer=*/ nullptr,
-                                dex_files,
-                                oat_rodata,
-                                &key_value_store)) {
-      return false;
-    }
+    oat_writer.Initialize(compiler_driver_.get(), nullptr, dex_files);
     oat_writer.PrepareLayout(&patcher);
     elf_writer->PrepareDynamicSection(oat_writer.GetOatHeader().GetExecutableOffset(),
                                       oat_writer.GetCodeSize(),