Refactor writing type lookup tables into OAT

OatWriter used to write type lookup tables through an mmap of the
DEX section of the OAT file. This approach would not work once the
DEX section is moved into the VDEX file, hence it was rewritten to
write the data through OutputStream.

Test: m test-art-host
Bug: 30937355
Change-Id: Id89a680e6e7cbdba47fa7e211a48b92d30a82b29
diff --git a/compiler/oat_writer.cc b/compiler/oat_writer.cc
index 2095608..c605578 100644
--- a/compiler/oat_writer.cc
+++ b/compiler/oat_writer.cc
@@ -226,7 +226,6 @@
     return dex_file_location_data_;
   }
 
-  void ReserveTypeLookupTable(OatWriter* oat_writer);
   void ReserveClassOffsets(OatWriter* oat_writer);
 
   size_t SizeOf() const;
@@ -436,36 +435,35 @@
                                 instruction_set_features,
                                 dchecked_integral_cast<uint32_t>(oat_dex_files_.size()),
                                 key_value_store);
-  offset = InitOatDexFiles(offset);
-  size_ = offset;
+  size_ = InitOatDexFiles(offset);
 
   std::unique_ptr<MemMap> dex_files_map;
   std::vector<std::unique_ptr<const DexFile>> dex_files;
-  if (!WriteDexFiles(rodata, file)) {
+  if (!WriteDexFiles(rodata, file) ||
+      !OpenDexFiles(file, verify, &dex_files_map, &dex_files)) {
     return false;
   }
-  // Reserve space for type lookup tables and update type_lookup_table_offset_.
-  for (OatDexFile& oat_dex_file : oat_dex_files_) {
-    oat_dex_file.ReserveTypeLookupTable(this);
+
+  // Do a bulk checksum update for Dex[]. Doing it piece by piece would be
+  // difficult because we're not using the OutputStream directly.
+  if (!oat_dex_files_.empty()) {
+    size_t size = size_ - oat_dex_files_[0].dex_file_offset_;
+    oat_header_->UpdateChecksum(dex_files_map->Begin(), size);
   }
-  size_t size_after_type_lookup_tables = size_;
+
+  ChecksumUpdatingOutputStream checksum_updating_rodata(rodata, oat_header_.get());
+
+  if (!WriteTypeLookupTables(&checksum_updating_rodata, dex_files)) {
+    return false;
+  }
+
   // Reserve space for class offsets and update class_offsets_offset_.
   for (OatDexFile& oat_dex_file : oat_dex_files_) {
     oat_dex_file.ReserveClassOffsets(this);
   }
-  ChecksumUpdatingOutputStream checksum_updating_rodata(rodata, oat_header_.get());
-  if (!WriteOatDexFiles(&checksum_updating_rodata) ||
-      !ExtendForTypeLookupTables(rodata, file, size_after_type_lookup_tables) ||
-      !OpenDexFiles(file, verify, &dex_files_map, &dex_files) ||
-      !WriteTypeLookupTables(dex_files_map.get(), dex_files)) {
-    return false;
-  }
 
-  // Do a bulk checksum update for Dex[] and TypeLookupTable[]. Doing it piece by
-  // piece would be difficult because we're not using the OutpuStream directly.
-  if (!oat_dex_files_.empty()) {
-    size_t size = size_after_type_lookup_tables - oat_dex_files_[0].dex_file_offset_;
-    oat_header_->UpdateChecksum(dex_files_map->Begin(), size);
+  if (!WriteOatDexFiles(&checksum_updating_rodata)) {
+    return false;
   }
 
   *opened_dex_files_map = std::move(dex_files_map);
@@ -2122,30 +2120,6 @@
   return true;
 }
 
-bool OatWriter::ExtendForTypeLookupTables(OutputStream* rodata, File* file, size_t offset) {
-  TimingLogger::ScopedTiming split("ExtendForTypeLookupTables", timings_);
-
-  int64_t new_length = oat_data_offset_ + dchecked_integral_cast<int64_t>(offset);
-  if (file->SetLength(new_length) != 0) {
-    PLOG(ERROR) << "Failed to extend file for type lookup tables. new_length: " << new_length
-        << "File: " << file->GetPath();
-    return false;
-  }
-  off_t actual_offset = rodata->Seek(new_length, kSeekSet);
-  if (actual_offset != static_cast<off_t>(new_length)) {
-    PLOG(ERROR) << "Failed to seek stream after extending file for type lookup tables."
-                << " Actual: " << actual_offset << " Expected: " << new_length
-                << " File: " << rodata->GetLocation();
-    return false;
-  }
-  if (!rodata->Flush()) {
-    PLOG(ERROR) << "Failed to flush stream after extending for type lookup tables."
-                << " File: " << rodata->GetLocation();
-    return false;
-  }
-  return true;
-}
-
 bool OatWriter::OpenDexFiles(
     File* file,
     bool verify,
@@ -2223,26 +2197,66 @@
 }
 
 bool OatWriter::WriteTypeLookupTables(
-    MemMap* opened_dex_files_map,
+    OutputStream* rodata,
     const std::vector<std::unique_ptr<const DexFile>>& opened_dex_files) {
   TimingLogger::ScopedTiming split("WriteTypeLookupTables", timings_);
 
   DCHECK_EQ(opened_dex_files.size(), oat_dex_files_.size());
   for (size_t i = 0, size = opened_dex_files.size(); i != size; ++i) {
     OatDexFile* oat_dex_file = &oat_dex_files_[i];
-    if (oat_dex_file->lookup_table_offset_ != 0u) {
-      DCHECK(oat_dex_file->create_type_lookup_table_ == CreateTypeLookupTable::kCreate);
-      DCHECK_NE(oat_dex_file->class_offsets_.size(), 0u);
-      size_t map_offset = oat_dex_files_[0].dex_file_offset_;
-      size_t lookup_table_offset = oat_dex_file->lookup_table_offset_;
-      uint8_t* lookup_table = opened_dex_files_map->Begin() + (lookup_table_offset - map_offset);
-      opened_dex_files[i]->CreateTypeLookupTable(lookup_table);
+    DCHECK_EQ(oat_dex_file->lookup_table_offset_, 0u);
+
+    if (oat_dex_file->create_type_lookup_table_ != CreateTypeLookupTable::kCreate ||
+        oat_dex_file->class_offsets_.empty()) {
+      continue;
     }
+
+    size_t table_size = TypeLookupTable::RawDataLength(oat_dex_file->class_offsets_.size());
+    if (table_size == 0u) {
+      continue;
+    }
+
+    // Create the lookup table. When `nullptr` is given as the storage buffer,
+    // TypeLookupTable allocates its own and DexFile takes ownership.
+    opened_dex_files[i]->CreateTypeLookupTable(/* storage */ nullptr);
+    TypeLookupTable* table = opened_dex_files[i]->GetTypeLookupTable();
+
+    // Type tables are required to be 4 byte aligned.
+    size_t original_offset = size_;
+    size_t rodata_offset = RoundUp(original_offset, 4);
+    size_t padding_size = rodata_offset - original_offset;
+
+    if (padding_size != 0u) {
+      std::vector<uint8_t> buffer(padding_size, 0u);
+      if (!rodata->WriteFully(buffer.data(), padding_size)) {
+        PLOG(ERROR) << "Failed to write lookup table alignment padding."
+                    << " File: " << oat_dex_file->GetLocation()
+                    << " Output: " << rodata->GetLocation();
+        return false;
+      }
+    }
+
+    DCHECK_EQ(oat_data_offset_ + rodata_offset,
+              static_cast<size_t>(rodata->Seek(0u, kSeekCurrent)));
+    DCHECK_EQ(table_size, table->RawDataLength());
+
+    if (!rodata->WriteFully(table->RawData(), table_size)) {
+      PLOG(ERROR) << "Failed to write lookup table."
+                  << " File: " << oat_dex_file->GetLocation()
+                  << " Output: " << rodata->GetLocation();
+      return false;
+    }
+
+    oat_dex_file->lookup_table_offset_ = rodata_offset;
+
+    size_ += padding_size + table_size;
+    size_oat_lookup_table_ += table_size;
+    size_oat_lookup_table_alignment_ += padding_size;
   }
 
-  DCHECK_EQ(opened_dex_files_map == nullptr, opened_dex_files.empty());
-  if (opened_dex_files_map != nullptr && !opened_dex_files_map->Sync()) {
-    PLOG(ERROR) << "Failed to Sync() type lookup tables. Map: " << opened_dex_files_map->GetName();
+  if (!rodata->Flush()) {
+    PLOG(ERROR) << "Failed to flush stream after writing type lookup tables."
+                << " File: " << rodata->GetLocation();
     return false;
   }
 
@@ -2298,22 +2312,6 @@
           + sizeof(lookup_table_offset_);
 }
 
-void OatWriter::OatDexFile::ReserveTypeLookupTable(OatWriter* oat_writer) {
-  DCHECK_EQ(lookup_table_offset_, 0u);
-  if (create_type_lookup_table_ == CreateTypeLookupTable::kCreate && !class_offsets_.empty()) {
-    size_t table_size = TypeLookupTable::RawDataLength(class_offsets_.size());
-    if (table_size != 0u) {
-      // Type tables are required to be 4 byte aligned.
-      size_t original_offset = oat_writer->size_;
-      size_t offset = RoundUp(original_offset, 4);
-      oat_writer->size_oat_lookup_table_alignment_ += offset - original_offset;
-      lookup_table_offset_ = offset;
-      oat_writer->size_ = offset + table_size;
-      oat_writer->size_oat_lookup_table_ += table_size;
-    }
-  }
-}
-
 void OatWriter::OatDexFile::ReserveClassOffsets(OatWriter* oat_writer) {
   DCHECK_EQ(class_offsets_offset_, 0u);
   if (!class_offsets_.empty()) {
diff --git a/compiler/oat_writer.h b/compiler/oat_writer.h
index decb7db..93e2e44 100644
--- a/compiler/oat_writer.h
+++ b/compiler/oat_writer.h
@@ -262,12 +262,11 @@
   bool WriteDexFile(OutputStream* rodata, File* file, OatDexFile* oat_dex_file, File* dex_file);
   bool WriteDexFile(OutputStream* rodata, OatDexFile* oat_dex_file, const uint8_t* dex_file);
   bool WriteOatDexFiles(OutputStream* rodata);
-  bool ExtendForTypeLookupTables(OutputStream* rodata, File* file, size_t offset);
   bool OpenDexFiles(File* file,
                     bool verify,
                     /*out*/ std::unique_ptr<MemMap>* opened_dex_files_map,
                     /*out*/ std::vector<std::unique_ptr<const DexFile>>* opened_dex_files);
-  bool WriteTypeLookupTables(MemMap* opened_dex_files_map,
+  bool WriteTypeLookupTables(OutputStream* rodata,
                              const std::vector<std::unique_ptr<const DexFile>>& opened_dex_files);
   bool WriteCodeAlignment(OutputStream* out, uint32_t aligned_code_delta);
   void SetMultiOatRelativePatcherAdjustment();