Add logic for calculating offsets during writing

Delete all the offset related layout logic, move the logic to
writing. Layout is now controlled by changing the order of the data
vectors.

Calculate all offsets during writing instead of reading. This allows
more flexible layout optimizations without needing to be as careful
about offsets in other sections changing.

Maintained logic for using existing offsets in the case where
changing the layout is not required.

Areas to improve:
How to size the memmap, do we want 2 passes? Currently the algorithm
reserves 10% extra space in case required. This is sufficient for top
1000 apps in the play store. Will consider a two pass approach later.

Bug: 63756964
Bug: 68948395
Bug: 68867570
Bug: 68864106
Bug: 68864168

Test: test/testrunner/testrunner.py --host -j40
Test: test-art-host-gtest
Test: Dexlayout speed profile top 1000 apps

Change-Id: I7dee869da3a010559547f8cfdf93e9aa4c3f47ff

15 files changed, 1019 insertions, 693 deletions

diff --git a/dex2oat/dex2oat_test.cc b/dex2oat/dex2oat_test.cc
index bd8583bd41..ad287b0745 100644
--- a/dex2oat/dex2oat_test.cc
+++ b/dex2oat/dex2oat_test.cc
@@ -1332,10 +1332,10 @@ TEST_F(Dex2oatTest, LayoutSections) {
         code_section.parts_[static_cast<size_t>(LayoutType::kLayoutTypeUnused)];
 
     // All the sections should be non-empty.
-    EXPECT_GT(section_hot_code.size_, 0u);
-    EXPECT_GT(section_sometimes_used.size_, 0u);
-    EXPECT_GT(section_startup_only.size_, 0u);
-    EXPECT_GT(section_unused.size_, 0u);
+    EXPECT_GT(section_hot_code.Size(), 0u);
+    EXPECT_GT(section_sometimes_used.Size(), 0u);
+    EXPECT_GT(section_startup_only.Size(), 0u);
+    EXPECT_GT(section_unused.Size(), 0u);
 
     // Open the dex file since we need to peek at the code items to verify the layout matches what
     // we expect.
@@ -1364,18 +1364,18 @@ TEST_F(Dex2oatTest, LayoutSections) {
       const bool is_post_startup = ContainsElement(post_methods, method_idx);
       if (is_hot) {
         // Hot is highest precedence, check that the hot methods are in the hot section.
-        EXPECT_LT(code_item_offset - section_hot_code.offset_, section_hot_code.size_);
+        EXPECT_TRUE(section_hot_code.Contains(code_item_offset));
         ++hot_count;
       } else if (is_post_startup) {
         // Post startup is sometimes used section.
-        EXPECT_LT(code_item_offset - section_sometimes_used.offset_, section_sometimes_used.size_);
+        EXPECT_TRUE(section_sometimes_used.Contains(code_item_offset));
         ++post_startup_count;
       } else if (is_startup) {
         // Startup at this point means not hot or post startup, these must be startup only then.
-        EXPECT_LT(code_item_offset - section_startup_only.offset_, section_startup_only.size_);
+        EXPECT_TRUE(section_startup_only.Contains(code_item_offset));
         ++startup_count;
       } else {
-        if (code_item_offset - section_unused.offset_ < section_unused.size_) {
+        if (section_unused.Contains(code_item_offset)) {
           // If no flags are set, the method should be unused ...
           ++unused_count;
         } else {
diff --git a/dexlayout/compact_dex_writer.cc b/dexlayout/compact_dex_writer.cc
index b089c1d4b3..f2d46199a2 100644
--- a/dexlayout/compact_dex_writer.cc
+++ b/dexlayout/compact_dex_writer.cc
@@ -47,7 +47,7 @@ void CompactDexWriter::WriteHeader() {
   header.class_defs_off_ = collections.ClassDefsOffset();
   header.data_size_ = header_->DataSize();
   header.data_off_ = header_->DataOffset();
-  Write(reinterpret_cast<uint8_t*>(&header), sizeof(header), 0u);
+  UNUSED(Write(reinterpret_cast<uint8_t*>(&header), sizeof(header), 0u));
 }
 
 }  // namespace art
diff --git a/dexlayout/compact_dex_writer.h b/dexlayout/compact_dex_writer.h
index 1c77202c9a..e28efab5c1 100644
--- a/dexlayout/compact_dex_writer.h
+++ b/dexlayout/compact_dex_writer.h
@@ -25,9 +25,12 @@ namespace art {
 
 class CompactDexWriter : public DexWriter {
  public:
-  CompactDexWriter(dex_ir::Header* header, MemMap* mem_map, CompactDexLevel compact_dex_level)
-      : DexWriter(header, mem_map),
-        compact_dex_level_(compact_dex_level) { }
+  CompactDexWriter(dex_ir::Header* header,
+                   MemMap* mem_map,
+                   DexLayout* dex_layout,
+                   CompactDexLevel compact_dex_level)
+      : DexWriter(header, mem_map, dex_layout, /*compute_offsets*/ true),
+        compact_dex_level_(compact_dex_level) {}
 
  protected:
   void WriteHeader() OVERRIDE;
diff --git a/dexlayout/dex_ir.cc b/dexlayout/dex_ir.cc
index 3edb0a44f2..a8ba950c28 100644
--- a/dexlayout/dex_ir.cc
+++ b/dexlayout/dex_ir.cc
@@ -186,22 +186,28 @@ static bool GetIdsFromByteCode(Collections& collections,
   return has_id;
 }
 
-EncodedValue* Collections::ReadEncodedValue(const uint8_t** data) {
+EncodedValue* Collections::ReadEncodedValue(const DexFile& dex_file, const uint8_t** data) {
   const uint8_t encoded_value = *(*data)++;
   const uint8_t type = encoded_value & 0x1f;
   EncodedValue* item = new EncodedValue(type);
-  ReadEncodedValue(data, type, encoded_value >> 5, item);
+  ReadEncodedValue(dex_file, data, type, encoded_value >> 5, item);
   return item;
 }
 
-EncodedValue* Collections::ReadEncodedValue(const uint8_t** data, uint8_t type, uint8_t length) {
+EncodedValue* Collections::ReadEncodedValue(const DexFile& dex_file,
+                                            const uint8_t** data,
+                                            uint8_t type,
+                                            uint8_t length) {
   EncodedValue* item = new EncodedValue(type);
-  ReadEncodedValue(data, type, length, item);
+  ReadEncodedValue(dex_file, data, type, length, item);
   return item;
 }
 
-void Collections::ReadEncodedValue(
-    const uint8_t** data, uint8_t type, uint8_t length, EncodedValue* item) {
+void Collections::ReadEncodedValue(const DexFile& dex_file,
+                                   const uint8_t** data,
+                                   uint8_t type,
+                                   uint8_t length,
+                                   EncodedValue* item) {
   switch (type) {
     case DexFile::kDexAnnotationByte:
       item->SetByte(static_cast<int8_t>(ReadVarWidth(data, length, false)));
@@ -271,12 +277,17 @@ void Collections::ReadEncodedValue(
     }
     case DexFile::kDexAnnotationArray: {
       EncodedValueVector* values = new EncodedValueVector();
+      const uint32_t offset = *data - dex_file.Begin();
       const uint32_t size = DecodeUnsignedLeb128(data);
       // Decode all elements.
       for (uint32_t i = 0; i < size; i++) {
-        values->push_back(std::unique_ptr<EncodedValue>(ReadEncodedValue(data)));
+        values->push_back(std::unique_ptr<EncodedValue>(ReadEncodedValue(dex_file, data)));
       }
-      item->SetEncodedArray(new EncodedArrayItem(values));
+      EncodedArrayItem* array_item = new EncodedArrayItem(values);
+      if (eagerly_assign_offsets_) {
+        array_item->SetOffset(offset);
+      }
+      item->SetEncodedArray(array_item);
       break;
     }
     case DexFile::kDexAnnotationAnnotation: {
@@ -287,7 +298,7 @@ void Collections::ReadEncodedValue(
       for (uint32_t i = 0; i < size; i++) {
         const uint32_t name_index = DecodeUnsignedLeb128(data);
         elements->push_back(std::unique_ptr<AnnotationElement>(
-            new AnnotationElement(GetStringId(name_index), ReadEncodedValue(data))));
+            new AnnotationElement(GetStringId(name_index), ReadEncodedValue(dex_file, data))));
       }
       item->SetEncodedAnnotation(new EncodedAnnotation(GetTypeId(type_idx), elements));
       break;
@@ -305,16 +316,16 @@ void Collections::ReadEncodedValue(
 void Collections::CreateStringId(const DexFile& dex_file, uint32_t i) {
   const DexFile::StringId& disk_string_id = dex_file.GetStringId(dex::StringIndex(i));
   StringData* string_data = new StringData(dex_file.GetStringData(disk_string_id));
-  string_datas_.AddItem(string_data, disk_string_id.string_data_off_);
+  AddItem(string_datas_map_, string_datas_, string_data, disk_string_id.string_data_off_);
 
   StringId* string_id = new StringId(string_data);
-  string_ids_.AddIndexedItem(string_id, StringIdsOffset() + i * StringId::ItemSize(), i);
+  AddIndexedItem(string_ids_, string_id, StringIdsOffset() + i * StringId::ItemSize(), i);
 }
 
 void Collections::CreateTypeId(const DexFile& dex_file, uint32_t i) {
   const DexFile::TypeId& disk_type_id = dex_file.GetTypeId(dex::TypeIndex(i));
   TypeId* type_id = new TypeId(GetStringId(disk_type_id.descriptor_idx_.index_));
-  type_ids_.AddIndexedItem(type_id, TypeIdsOffset() + i * TypeId::ItemSize(), i);
+  AddIndexedItem(type_ids_, type_id, TypeIdsOffset() + i * TypeId::ItemSize(), i);
 }
 
 void Collections::CreateProtoId(const DexFile& dex_file, uint32_t i) {
@@ -325,7 +336,7 @@ void Collections::CreateProtoId(const DexFile& dex_file, uint32_t i) {
   ProtoId* proto_id = new ProtoId(GetStringId(disk_proto_id.shorty_idx_.index_),
                                   GetTypeId(disk_proto_id.return_type_idx_.index_),
                                   parameter_type_list);
-  proto_ids_.AddIndexedItem(proto_id, ProtoIdsOffset() + i * ProtoId::ItemSize(), i);
+  AddIndexedItem(proto_ids_, proto_id, ProtoIdsOffset() + i * ProtoId::ItemSize(), i);
 }
 
 void Collections::CreateFieldId(const DexFile& dex_file, uint32_t i) {
@@ -333,7 +344,7 @@ void Collections::CreateFieldId(const DexFile& dex_file, uint32_t i) {
   FieldId* field_id = new FieldId(GetTypeId(disk_field_id.class_idx_.index_),
                                   GetTypeId(disk_field_id.type_idx_.index_),
                                   GetStringId(disk_field_id.name_idx_.index_));
-  field_ids_.AddIndexedItem(field_id, FieldIdsOffset() + i * FieldId::ItemSize(), i);
+  AddIndexedItem(field_ids_, field_id, FieldIdsOffset() + i * FieldId::ItemSize(), i);
 }
 
 void Collections::CreateMethodId(const DexFile& dex_file, uint32_t i) {
@@ -341,7 +352,7 @@ void Collections::CreateMethodId(const DexFile& dex_file, uint32_t i) {
   MethodId* method_id = new MethodId(GetTypeId(disk_method_id.class_idx_.index_),
                                      GetProtoId(disk_method_id.proto_idx_),
                                      GetStringId(disk_method_id.name_idx_.index_));
-  method_ids_.AddIndexedItem(method_id, MethodIdsOffset() + i * MethodId::ItemSize(), i);
+  AddIndexedItem(method_ids_, method_id, MethodIdsOffset() + i * MethodId::ItemSize(), i);
 }
 
 void Collections::CreateClassDef(const DexFile& dex_file, uint32_t i) {
@@ -365,48 +376,48 @@ void Collections::CreateClassDef(const DexFile& dex_file, uint32_t i) {
   // Static field initializers.
   const uint8_t* static_data = dex_file.GetEncodedStaticFieldValuesArray(disk_class_def);
   EncodedArrayItem* static_values =
-      CreateEncodedArrayItem(static_data, disk_class_def.static_values_off_);
+      CreateEncodedArrayItem(dex_file, static_data, disk_class_def.static_values_off_);
   ClassData* class_data = CreateClassData(
       dex_file, dex_file.GetClassData(disk_class_def), disk_class_def.class_data_off_);
   ClassDef* class_def = new ClassDef(class_type, access_flags, superclass, interfaces_type_list,
                                      source_file, annotations, static_values, class_data);
-  class_defs_.AddIndexedItem(class_def, ClassDefsOffset() + i * ClassDef::ItemSize(), i);
+  AddIndexedItem(class_defs_, class_def, ClassDefsOffset() + i * ClassDef::ItemSize(), i);
 }
 
 TypeList* Collections::CreateTypeList(const DexFile::TypeList* dex_type_list, uint32_t offset) {
   if (dex_type_list == nullptr) {
     return nullptr;
   }
-  auto found_type_list = TypeLists().find(offset);
-  if (found_type_list != TypeLists().end()) {
-    return found_type_list->second.get();
-  }
-  TypeIdVector* type_vector = new TypeIdVector();
-  uint32_t size = dex_type_list->Size();
-  for (uint32_t index = 0; index < size; ++index) {
-    type_vector->push_back(GetTypeId(dex_type_list->GetTypeItem(index).type_idx_.index_));
+  TypeList* type_list = type_lists_map_.GetExistingObject(offset);
+  if (type_list == nullptr) {
+    TypeIdVector* type_vector = new TypeIdVector();
+    uint32_t size = dex_type_list->Size();
+    for (uint32_t index = 0; index < size; ++index) {
+      type_vector->push_back(GetTypeId(dex_type_list->GetTypeItem(index).type_idx_.index_));
+    }
+    type_list = new TypeList(type_vector);
+    AddItem(type_lists_map_, type_lists_, type_list, offset);
   }
-  TypeList* new_type_list = new TypeList(type_vector);
-  type_lists_.AddItem(new_type_list, offset);
-  return new_type_list;
+  return type_list;
 }
 
-EncodedArrayItem* Collections::CreateEncodedArrayItem(const uint8_t* static_data, uint32_t offset) {
+EncodedArrayItem* Collections::CreateEncodedArrayItem(const DexFile& dex_file,
+                                                      const uint8_t* static_data,
+                                                      uint32_t offset) {
   if (static_data == nullptr) {
     return nullptr;
   }
-  auto found_encoded_array_item = EncodedArrayItems().find(offset);
-  if (found_encoded_array_item != EncodedArrayItems().end()) {
-    return found_encoded_array_item->second.get();
-  }
-  uint32_t size = DecodeUnsignedLeb128(&static_data);
-  EncodedValueVector* values = new EncodedValueVector();
-  for (uint32_t i = 0; i < size; ++i) {
-    values->push_back(std::unique_ptr<EncodedValue>(ReadEncodedValue(&static_data)));
+  EncodedArrayItem* encoded_array_item = encoded_array_items_map_.GetExistingObject(offset);
+  if (encoded_array_item == nullptr) {
+    uint32_t size = DecodeUnsignedLeb128(&static_data);
+    EncodedValueVector* values = new EncodedValueVector();
+    for (uint32_t i = 0; i < size; ++i) {
+      values->push_back(std::unique_ptr<EncodedValue>(ReadEncodedValue(dex_file, &static_data)));
+    }
+    // TODO: Calculate the size of the encoded array.
+    encoded_array_item = new EncodedArrayItem(values);
+    AddItem(encoded_array_items_map_, encoded_array_items_, encoded_array_item, offset);
   }
-  // TODO: Calculate the size of the encoded array.
-  EncodedArrayItem* encoded_array_item = new EncodedArrayItem(values);
-  encoded_array_items_.AddItem(encoded_array_item, offset);
   return encoded_array_item;
 }
 
@@ -427,19 +438,16 @@ AnnotationItem* Collections::CreateAnnotationItem(const DexFile& dex_file,
                                                   const DexFile::AnnotationItem* annotation) {
   const uint8_t* const start_data = reinterpret_cast<const uint8_t*>(annotation);
   const uint32_t offset = start_data - dex_file.Begin();
-  auto found_annotation_item = AnnotationItems().find(offset);
-  if (found_annotation_item != AnnotationItems().end()) {
-    return found_annotation_item->second.get();
+  AnnotationItem* annotation_item = annotation_items_map_.GetExistingObject(offset);
+  if (annotation_item == nullptr) {
+    uint8_t visibility = annotation->visibility_;
+    const uint8_t* annotation_data = annotation->annotation_;
+    std::unique_ptr<EncodedValue> encoded_value(
+        ReadEncodedValue(dex_file, &annotation_data, DexFile::kDexAnnotationAnnotation, 0));
+    annotation_item = new AnnotationItem(visibility, encoded_value->ReleaseEncodedAnnotation());
+    annotation_item->SetSize(annotation_data - start_data);
+    AddItem(annotation_items_map_, annotation_items_, annotation_item, offset);
   }
-  uint8_t visibility = annotation->visibility_;
-  const uint8_t* annotation_data = annotation->annotation_;
-  std::unique_ptr<EncodedValue> encoded_value(
-      ReadEncodedValue(&annotation_data, DexFile::kDexAnnotationAnnotation, 0));
-  AnnotationItem* annotation_item =
-      new AnnotationItem(visibility, encoded_value->ReleaseEncodedAnnotation());
-  annotation_item->SetOffset(offset);
-  annotation_item->SetSize(annotation_data - start_data);
-  annotation_items_.AddItem(annotation_item, annotation_item->GetOffset());
   return annotation_item;
 }
 
@@ -449,30 +457,30 @@ AnnotationSetItem* Collections::CreateAnnotationSetItem(const DexFile& dex_file,
   if (disk_annotations_item == nullptr || (disk_annotations_item->size_ == 0 && offset == 0)) {
     return nullptr;
   }
-  auto found_anno_set_item = AnnotationSetItems().find(offset);
-  if (found_anno_set_item != AnnotationSetItems().end()) {
-    return found_anno_set_item->second.get();
-  }
-  std::vector<AnnotationItem*>* items = new std::vector<AnnotationItem*>();
-  for (uint32_t i = 0; i < disk_annotations_item->size_; ++i) {
-    const DexFile::AnnotationItem* annotation =
-        dex_file.GetAnnotationItem(disk_annotations_item, i);
-    if (annotation == nullptr) {
-      continue;
+  AnnotationSetItem* annotation_set_item = annotation_set_items_map_.GetExistingObject(offset);
+  if (annotation_set_item == nullptr) {
+    std::vector<AnnotationItem*>* items = new std::vector<AnnotationItem*>();
+    for (uint32_t i = 0; i < disk_annotations_item->size_; ++i) {
+      const DexFile::AnnotationItem* annotation =
+          dex_file.GetAnnotationItem(disk_annotations_item, i);
+      if (annotation == nullptr) {
+        continue;
+      }
+      AnnotationItem* annotation_item = CreateAnnotationItem(dex_file, annotation);
+      items->push_back(annotation_item);
     }
-    AnnotationItem* annotation_item = CreateAnnotationItem(dex_file, annotation);
-    items->push_back(annotation_item);
+    annotation_set_item = new AnnotationSetItem(items);
+    AddItem(annotation_set_items_map_, annotation_set_items_, annotation_set_item, offset);
   }
-  AnnotationSetItem* annotation_set_item = new AnnotationSetItem(items);
-  annotation_set_items_.AddItem(annotation_set_item, offset);
   return annotation_set_item;
 }
 
 AnnotationsDirectoryItem* Collections::CreateAnnotationsDirectoryItem(const DexFile& dex_file,
     const DexFile::AnnotationsDirectoryItem* disk_annotations_item, uint32_t offset) {
-  auto found_anno_dir_item = AnnotationsDirectoryItems().find(offset);
-  if (found_anno_dir_item != AnnotationsDirectoryItems().end()) {
-    return found_anno_dir_item->second.get();
+  AnnotationsDirectoryItem* annotations_directory_item =
+      annotations_directory_items_map_.GetExistingObject(offset);
+  if (annotations_directory_item != nullptr) {
+    return annotations_directory_item;
   }
   const DexFile::AnnotationSetItem* class_set_item =
       dex_file.GetClassAnnotationSet(disk_annotations_item);
@@ -527,20 +535,19 @@ AnnotationsDirectoryItem* Collections::CreateAnnotationsDirectoryItem(const DexF
     }
   }
   // TODO: Calculate the size of the annotations directory.
-  AnnotationsDirectoryItem* annotations_directory_item = new AnnotationsDirectoryItem(
+annotations_directory_item = new AnnotationsDirectoryItem(
       class_annotation, field_annotations, method_annotations, parameter_annotations);
-  annotations_directory_items_.AddItem(annotations_directory_item, offset);
+  AddItem(annotations_directory_items_map_,
+          annotations_directory_items_,
+          annotations_directory_item,
+          offset);
   return annotations_directory_item;
 }
 
 ParameterAnnotation* Collections::GenerateParameterAnnotation(
     const DexFile& dex_file, MethodId* method_id,
     const DexFile::AnnotationSetRefList* annotation_set_ref_list, uint32_t offset) {
-  AnnotationSetRefList* set_ref_list = nullptr;
-  auto found_set_ref_list = AnnotationSetRefLists().find(offset);
-  if (found_set_ref_list != AnnotationSetRefLists().end()) {
-    set_ref_list = found_set_ref_list->second.get();
-  }
+  AnnotationSetRefList* set_ref_list = annotation_set_ref_lists_map_.GetExistingObject(offset);
   if (set_ref_list == nullptr) {
     std::vector<AnnotationSetItem*>* annotations = new std::vector<AnnotationSetItem*>();
     for (uint32_t i = 0; i < annotation_set_ref_list->size_; ++i) {
@@ -550,7 +557,7 @@ ParameterAnnotation* Collections::GenerateParameterAnnotation(
       annotations->push_back(CreateAnnotationSetItem(dex_file, annotation_set_item, set_offset));
     }
     set_ref_list = new AnnotationSetRefList(annotations);
-    annotation_set_ref_lists_.AddItem(set_ref_list, offset);
+    AddItem(annotation_set_ref_lists_map_, annotation_set_ref_lists_, set_ref_list, offset);
   }
   return new ParameterAnnotation(method_id, set_ref_list);
 }
@@ -566,13 +573,13 @@ CodeItem* Collections::CreateCodeItem(const DexFile& dex_file,
   const uint8_t* debug_info_stream = dex_file.GetDebugInfoStream(&disk_code_item);
   DebugInfoItem* debug_info = nullptr;
   if (debug_info_stream != nullptr) {
-    debug_info = debug_info_items_.GetExistingObject(disk_code_item.debug_info_off_);
+    debug_info = debug_info_items_map_.GetExistingObject(disk_code_item.debug_info_off_);
     if (debug_info == nullptr) {
       uint32_t debug_info_size = GetDebugInfoStreamSize(debug_info_stream);
       uint8_t* debug_info_buffer = new uint8_t[debug_info_size];
       memcpy(debug_info_buffer, debug_info_stream, debug_info_size);
       debug_info = new DebugInfoItem(debug_info_size, debug_info_buffer);
-      debug_info_items_.AddItem(debug_info, disk_code_item.debug_info_off_);
+      AddItem(debug_info_items_map_, debug_info_items_, debug_info, disk_code_item.debug_info_off_);
     }
   }
 
@@ -662,7 +669,7 @@ CodeItem* Collections::CreateCodeItem(const DexFile& dex_file,
   CodeItem* code_item = new CodeItem(
       registers_size, ins_size, outs_size, debug_info, insns_size, insns, tries, handler_list);
   code_item->SetSize(size);
-  code_items_.AddItem(code_item, offset);
+  AddItem(code_items_map_, code_items_, code_item, offset);
   // Add "fixup" references to types, strings, methods, and fields.
   // This is temporary, as we will probably want more detailed parsing of the
   // instructions here.
@@ -690,7 +697,7 @@ MethodItem* Collections::GenerateMethodItem(const DexFile& dex_file, ClassDataIt
   MethodId* method_id = GetMethodId(cdii.GetMemberIndex());
   uint32_t access_flags = cdii.GetRawMemberAccessFlags();
   const DexFile::CodeItem* disk_code_item = cdii.GetMethodCodeItem();
-  CodeItem* code_item = code_items_.GetExistingObject(cdii.GetMethodCodeItemOffset());
+  CodeItem* code_item = code_items_map_.GetExistingObject(cdii.GetMethodCodeItemOffset());
   DebugInfoItem* debug_info = nullptr;
   if (disk_code_item != nullptr) {
     if (code_item == nullptr) {
@@ -705,7 +712,7 @@ ClassData* Collections::CreateClassData(
     const DexFile& dex_file, const uint8_t* encoded_data, uint32_t offset) {
   // Read the fields and methods defined by the class, resolving the circular reference from those
   // to classes by setting class at the same time.
-  ClassData* class_data = class_datas_.GetExistingObject(offset);
+  ClassData* class_data = class_datas_map_.GetExistingObject(offset);
   if (class_data == nullptr && encoded_data != nullptr) {
     ClassDataItemIterator cdii(dex_file, encoded_data);
     // Static fields.
@@ -735,7 +742,7 @@ ClassData* Collections::CreateClassData(
     }
     class_data = new ClassData(static_fields, instance_fields, direct_methods, virtual_methods);
     class_data->SetSize(cdii.EndDataPointer() - encoded_data);
-    class_datas_.AddItem(class_data, offset);
+    AddItem(class_datas_map_, class_datas_, class_data, offset);
   }
   return class_data;
 }
@@ -771,10 +778,10 @@ void Collections::CreateCallSiteId(const DexFile& dex_file, uint32_t i) {
   const DexFile::CallSiteIdItem& disk_call_site_id = dex_file.GetCallSiteId(i);
   const uint8_t* disk_call_item_ptr = dex_file.Begin() + disk_call_site_id.data_off_;
   EncodedArrayItem* call_site_item =
-      CreateEncodedArrayItem(disk_call_item_ptr, disk_call_site_id.data_off_);
+      CreateEncodedArrayItem(dex_file, disk_call_item_ptr, disk_call_site_id.data_off_);
 
   CallSiteId* call_site_id = new CallSiteId(call_site_item);
-  call_site_ids_.AddIndexedItem(call_site_id, CallSiteIdsOffset() + i * CallSiteId::ItemSize(), i);
+  AddIndexedItem(call_site_ids_, call_site_id, CallSiteIdsOffset() + i * CallSiteId::ItemSize(), i);
 }
 
 void Collections::CreateMethodHandleItem(const DexFile& dex_file, uint32_t i) {
@@ -796,8 +803,23 @@ void Collections::CreateMethodHandleItem(const DexFile& dex_file, uint32_t i) {
     field_or_method_id = GetFieldId(index);
   }
   MethodHandleItem* method_handle = new MethodHandleItem(type, field_or_method_id);
-  method_handle_items_.AddIndexedItem(
-      method_handle, MethodHandleItemsOffset() + i * MethodHandleItem::ItemSize(), i);
+  AddIndexedItem(method_handle_items_,
+                 method_handle,
+                 MethodHandleItemsOffset() + i * MethodHandleItem::ItemSize(),
+                 i);
+}
+
+void Collections::SortVectorsByMapOrder() {
+  string_datas_map_.SortVectorByMapOrder(string_datas_);
+  type_lists_map_.SortVectorByMapOrder(type_lists_);
+  encoded_array_items_map_.SortVectorByMapOrder(encoded_array_items_);
+  annotation_items_map_.SortVectorByMapOrder(annotation_items_);
+  annotation_set_items_map_.SortVectorByMapOrder(annotation_set_items_);
+  annotation_set_ref_lists_map_.SortVectorByMapOrder(annotation_set_ref_lists_);
+  annotations_directory_items_map_.SortVectorByMapOrder(annotations_directory_items_);
+  debug_info_items_map_.SortVectorByMapOrder(debug_info_items_);
+  code_items_map_.SortVectorByMapOrder(code_items_);
+  class_datas_map_.SortVectorByMapOrder(class_datas_);
 }
 
 static uint32_t HeaderOffset(const dex_ir::Collections& collections ATTRIBUTE_UNUSED) {
diff --git a/dexlayout/dex_ir.h b/dexlayout/dex_ir.h
index 179d3b96e0..61a4eae9b1 100644
--- a/dexlayout/dex_ir.h
+++ b/dexlayout/dex_ir.h
@@ -112,33 +112,55 @@ template<class T> class CollectionBase {
  public:
   CollectionBase() = default;
 
-  uint32_t GetOffset() const { return offset_; }
-  void SetOffset(uint32_t new_offset) { offset_ = new_offset; }
+  uint32_t GetOffset() const {
+    return offset_;
+  }
+  void SetOffset(uint32_t new_offset) {
+    offset_ = new_offset;
+  }
 
  private:
-  uint32_t offset_ = 0;
+  // Start out unassigned.
+  uint32_t offset_ = 0u;
 
   DISALLOW_COPY_AND_ASSIGN(CollectionBase);
 };
 
 template<class T> class CollectionVector : public CollectionBase<T> {
  public:
+  using Vector = std::vector<std::unique_ptr<T>>;
   CollectionVector() = default;
 
-  void AddIndexedItem(T* object, uint32_t offset, uint32_t index) {
-    object->SetOffset(offset);
-    object->SetIndex(index);
+  uint32_t Size() const { return collection_.size(); }
+  Vector& Collection() { return collection_; }
+
+ protected:
+  Vector collection_;
+
+  void AddItem(T* object) {
     collection_.push_back(std::unique_ptr<T>(object));
   }
-  uint32_t Size() const { return collection_.size(); }
-  std::vector<std::unique_ptr<T>>& Collection() { return collection_; }
 
  private:
-  std::vector<std::unique_ptr<T>> collection_;
-
+  friend class Collections;
   DISALLOW_COPY_AND_ASSIGN(CollectionVector);
 };
 
+template<class T> class IndexedCollectionVector : public CollectionVector<T> {
+ public:
+  using Vector = std::vector<std::unique_ptr<T>>;
+  IndexedCollectionVector() = default;
+
+ private:
+  void AddIndexedItem(T* object, uint32_t index) {
+    object->SetIndex(index);
+    CollectionVector<T>::collection_.push_back(std::unique_ptr<T>(object));
+  }
+
+  friend class Collections;
+  DISALLOW_COPY_AND_ASSIGN(IndexedCollectionVector);
+};
+
 template<class T> class CollectionMap : public CollectionBase<T> {
  public:
   CollectionMap() = default;
@@ -146,21 +168,35 @@ template<class T> class CollectionMap : public CollectionBase<T> {
   // Returns the existing item if it is already inserted, null otherwise.
   T* GetExistingObject(uint32_t offset) {
     auto it = collection_.find(offset);
-    return it != collection_.end() ? it->second.get() : nullptr;
+    return it != collection_.end() ? it->second : nullptr;
   }
 
-  void AddItem(T* object, uint32_t offset) {
-    object->SetOffset(offset);
-    auto it = collection_.emplace(offset, std::unique_ptr<T>(object));
-    CHECK(it.second) << "CollectionMap already has an object with offset " << offset << " "
-                     << " and address " << it.first->second.get();
-  }
   uint32_t Size() const { return collection_.size(); }
-  std::map<uint32_t, std::unique_ptr<T>>& Collection() { return collection_; }
+  std::map<uint32_t, T*>& Collection() { return collection_; }
+
+  // Sort the vector by copying pointers over.
+  void SortVectorByMapOrder(CollectionVector<T>& vector) {
+    auto it = collection_.begin();
+    CHECK_EQ(vector.Size(), Size());
+    for (size_t i = 0; i < Size(); ++i) {
+      // There are times when the array will temporarily contain the same pointer twice, doing the
+      // release here sure there is no double free errors.
+      vector.Collection()[i].release();
+      vector.Collection()[i].reset(it->second);
+      ++it;
+    }
+  }
 
  private:
-  std::map<uint32_t, std::unique_ptr<T>> collection_;
+  std::map<uint32_t, T*> collection_;
 
+  void AddItem(T* object, uint32_t offset) {
+    auto it = collection_.emplace(offset, object);
+    CHECK(it.second) << "CollectionMap already has an object with offset " << offset << " "
+                     << " and address " << it.first->second;
+  }
+
+  friend class Collections;
   DISALLOW_COPY_AND_ASSIGN(CollectionMap);
 };
 
@@ -168,32 +204,31 @@ class Collections {
  public:
   Collections() = default;
 
-  std::vector<std::unique_ptr<StringId>>& StringIds() { return string_ids_.Collection(); }
-  std::vector<std::unique_ptr<TypeId>>& TypeIds() { return type_ids_.Collection(); }
-  std::vector<std::unique_ptr<ProtoId>>& ProtoIds() { return proto_ids_.Collection(); }
-  std::vector<std::unique_ptr<FieldId>>& FieldIds() { return field_ids_.Collection(); }
-  std::vector<std::unique_ptr<MethodId>>& MethodIds() { return method_ids_.Collection(); }
-  std::vector<std::unique_ptr<ClassDef>>& ClassDefs() { return class_defs_.Collection(); }
-  std::vector<std::unique_ptr<CallSiteId>>& CallSiteIds() { return call_site_ids_.Collection(); }
-  std::vector<std::unique_ptr<MethodHandleItem>>& MethodHandleItems()
+  CollectionVector<StringId>::Vector& StringIds() { return string_ids_.Collection(); }
+  CollectionVector<TypeId>::Vector& TypeIds() { return type_ids_.Collection(); }
+  CollectionVector<ProtoId>::Vector& ProtoIds() { return proto_ids_.Collection(); }
+  CollectionVector<FieldId>::Vector& FieldIds() { return field_ids_.Collection(); }
+  CollectionVector<MethodId>::Vector& MethodIds() { return method_ids_.Collection(); }
+  CollectionVector<ClassDef>::Vector& ClassDefs() { return class_defs_.Collection(); }
+  CollectionVector<CallSiteId>::Vector& CallSiteIds() { return call_site_ids_.Collection(); }
+  CollectionVector<MethodHandleItem>::Vector& MethodHandleItems()
       { return method_handle_items_.Collection(); }
-  std::map<uint32_t, std::unique_ptr<StringData>>& StringDatas()
-      { return string_datas_.Collection(); }
-  std::map<uint32_t, std::unique_ptr<TypeList>>& TypeLists() { return type_lists_.Collection(); }
-  std::map<uint32_t, std::unique_ptr<EncodedArrayItem>>& EncodedArrayItems()
+  CollectionVector<StringData>::Vector& StringDatas() { return string_datas_.Collection(); }
+  CollectionVector<TypeList>::Vector& TypeLists() { return type_lists_.Collection(); }
+  CollectionVector<EncodedArrayItem>::Vector& EncodedArrayItems()
       { return encoded_array_items_.Collection(); }
-  std::map<uint32_t, std::unique_ptr<AnnotationItem>>& AnnotationItems()
+  CollectionVector<AnnotationItem>::Vector& AnnotationItems()
       { return annotation_items_.Collection(); }
-  std::map<uint32_t, std::unique_ptr<AnnotationSetItem>>& AnnotationSetItems()
+  CollectionVector<AnnotationSetItem>::Vector& AnnotationSetItems()
       { return annotation_set_items_.Collection(); }
-  std::map<uint32_t, std::unique_ptr<AnnotationSetRefList>>& AnnotationSetRefLists()
+  CollectionVector<AnnotationSetRefList>::Vector& AnnotationSetRefLists()
       { return annotation_set_ref_lists_.Collection(); }
-  std::map<uint32_t, std::unique_ptr<AnnotationsDirectoryItem>>& AnnotationsDirectoryItems()
+  CollectionVector<AnnotationsDirectoryItem>::Vector& AnnotationsDirectoryItems()
       { return annotations_directory_items_.Collection(); }
-  std::map<uint32_t, std::unique_ptr<DebugInfoItem>>& DebugInfoItems()
+  CollectionVector<DebugInfoItem>::Vector& DebugInfoItems()
       { return debug_info_items_.Collection(); }
-  std::map<uint32_t, std::unique_ptr<CodeItem>>& CodeItems() { return code_items_.Collection(); }
-  std::map<uint32_t, std::unique_ptr<ClassData>>& ClassDatas() { return class_datas_.Collection(); }
+  CollectionVector<CodeItem>::Vector& CodeItems() { return code_items_.Collection(); }
+  CollectionVector<ClassData>::Vector& ClassDatas() { return class_datas_.Collection(); }
 
   void CreateStringId(const DexFile& dex_file, uint32_t i);
   void CreateTypeId(const DexFile& dex_file, uint32_t i);
@@ -207,7 +242,9 @@ class Collections {
   void CreateCallSitesAndMethodHandles(const DexFile& dex_file);
 
   TypeList* CreateTypeList(const DexFile::TypeList* type_list, uint32_t offset);
-  EncodedArrayItem* CreateEncodedArrayItem(const uint8_t* static_data, uint32_t offset);
+  EncodedArrayItem* CreateEncodedArrayItem(const DexFile& dex_file,
+                                           const uint8_t* static_data,
+                                           uint32_t offset);
   AnnotationItem* CreateAnnotationItem(const DexFile& dex_file,
                                        const DexFile::AnnotationItem* annotation);
   AnnotationSetItem* CreateAnnotationSetItem(const DexFile& dex_file,
@@ -326,37 +363,99 @@ class Collections {
   uint32_t CodeItemsSize() const { return code_items_.Size(); }
   uint32_t ClassDatasSize() const { return class_datas_.Size(); }
 
+  // Sort the vectors buy map order (same order that was used in the input file).
+  void SortVectorsByMapOrder();
+
+  template <typename Type>
+  void AddItem(CollectionMap<Type>& map,
+               CollectionVector<Type>& vector,
+               Type* item,
+               uint32_t offset) {
+    DCHECK(!map.GetExistingObject(offset));
+    DCHECK(!item->OffsetAssigned());
+    if (eagerly_assign_offsets_) {
+      item->SetOffset(offset);
+    }
+    map.AddItem(item, offset);
+    vector.AddItem(item);
+  }
+
+  template <typename Type>
+  void AddIndexedItem(IndexedCollectionVector<Type>& vector,
+                      Type* item,
+                      uint32_t offset,
+                      uint32_t index) {
+    DCHECK(!item->OffsetAssigned());
+    if (eagerly_assign_offsets_) {
+      item->SetOffset(offset);
+    }
+    vector.AddIndexedItem(item, index);
+  }
+
+  void SetEagerlyAssignOffsets(bool eagerly_assign_offsets) {
+    eagerly_assign_offsets_ = eagerly_assign_offsets;
+  }
+
  private:
-  EncodedValue* ReadEncodedValue(const uint8_t** data);
-  EncodedValue* ReadEncodedValue(const uint8_t** data, uint8_t type, uint8_t length);
-  void ReadEncodedValue(const uint8_t** data, uint8_t type, uint8_t length, EncodedValue* item);
+  EncodedValue* ReadEncodedValue(const DexFile& dex_file, const uint8_t** data);
+  EncodedValue* ReadEncodedValue(const DexFile& dex_file,
+                                 const uint8_t** data,
+                                 uint8_t type,
+                                 uint8_t length);
+  void ReadEncodedValue(const DexFile& dex_file,
+                        const uint8_t** data,
+                        uint8_t type,
+                        uint8_t length,
+                        EncodedValue* item);
 
   ParameterAnnotation* GenerateParameterAnnotation(const DexFile& dex_file, MethodId* method_id,
       const DexFile::AnnotationSetRefList* annotation_set_ref_list, uint32_t offset);
   MethodItem* GenerateMethodItem(const DexFile& dex_file, ClassDataItemIterator& cdii);
 
-  CollectionVector<StringId> string_ids_;
-  CollectionVector<TypeId> type_ids_;
-  CollectionVector<ProtoId> proto_ids_;
-  CollectionVector<FieldId> field_ids_;
-  CollectionVector<MethodId> method_ids_;
-  CollectionVector<ClassDef> class_defs_;
-  CollectionVector<CallSiteId> call_site_ids_;
-  CollectionVector<MethodHandleItem> method_handle_items_;
-
-  CollectionMap<StringData> string_datas_;
-  CollectionMap<TypeList> type_lists_;
-  CollectionMap<EncodedArrayItem> encoded_array_items_;
-  CollectionMap<AnnotationItem> annotation_items_;
-  CollectionMap<AnnotationSetItem> annotation_set_items_;
-  CollectionMap<AnnotationSetRefList> annotation_set_ref_lists_;
-  CollectionMap<AnnotationsDirectoryItem> annotations_directory_items_;
-  CollectionMap<DebugInfoItem> debug_info_items_;
-  CollectionMap<CodeItem> code_items_;
-  CollectionMap<ClassData> class_datas_;
+  // Collection vectors own the IR data.
+  IndexedCollectionVector<StringId> string_ids_;
+  IndexedCollectionVector<TypeId> type_ids_;
+  IndexedCollectionVector<ProtoId> proto_ids_;
+  IndexedCollectionVector<FieldId> field_ids_;
+  IndexedCollectionVector<MethodId> method_ids_;
+  IndexedCollectionVector<CallSiteId> call_site_ids_;
+  IndexedCollectionVector<MethodHandleItem> method_handle_items_;
+  IndexedCollectionVector<StringData> string_datas_;
+  IndexedCollectionVector<TypeList> type_lists_;
+  IndexedCollectionVector<EncodedArrayItem> encoded_array_items_;
+  IndexedCollectionVector<AnnotationItem> annotation_items_;
+  IndexedCollectionVector<AnnotationSetItem> annotation_set_items_;
+  IndexedCollectionVector<AnnotationSetRefList> annotation_set_ref_lists_;
+  IndexedCollectionVector<AnnotationsDirectoryItem> annotations_directory_items_;
+  IndexedCollectionVector<ClassDef> class_defs_;
+  // The order of the vectors controls the layout of the output file by index order, to change the
+  // layout just sort the vector. Note that you may only change the order of the non indexed vectors
+  // below. Indexed vectors are accessed by indices in other places, changing the sorting order will
+  // invalidate the existing indices and is not currently supported.
+  CollectionVector<DebugInfoItem> debug_info_items_;
+  CollectionVector<CodeItem> code_items_;
+  CollectionVector<ClassData> class_datas_;
+
+  // Note that the maps do not have ownership, the vectors do.
+  // TODO: These maps should only be required for building the IR and should be put in a separate
+  // IR builder class.
+  CollectionMap<StringData> string_datas_map_;
+  CollectionMap<TypeList> type_lists_map_;
+  CollectionMap<EncodedArrayItem> encoded_array_items_map_;
+  CollectionMap<AnnotationItem> annotation_items_map_;
+  CollectionMap<AnnotationSetItem> annotation_set_items_map_;
+  CollectionMap<AnnotationSetRefList> annotation_set_ref_lists_map_;
+  CollectionMap<AnnotationsDirectoryItem> annotations_directory_items_map_;
+  CollectionMap<DebugInfoItem> debug_info_items_map_;
+  CollectionMap<CodeItem> code_items_map_;
+  CollectionMap<ClassData> class_datas_map_;
 
   uint32_t map_list_offset_ = 0;
 
+  // If we eagerly assign offsets during IR building or later after layout. Must be false if
+  // changing the layout is enabled.
+  bool eagerly_assign_offsets_;
+
   DISALLOW_COPY_AND_ASSIGN(Collections);
 };
 
@@ -365,15 +464,26 @@ class Item {
   Item() { }
   virtual ~Item() { }
 
-  uint32_t GetOffset() const { return offset_; }
+  // Return the assigned offset.
+  uint32_t GetOffset() const {
+    CHECK(OffsetAssigned());
+    return offset_;
+  }
   uint32_t GetSize() const { return size_; }
   void SetOffset(uint32_t offset) { offset_ = offset; }
   void SetSize(uint32_t size) { size_ = size; }
+  bool OffsetAssigned() const {
+    return offset_ != kOffsetUnassigned;
+  }
 
  protected:
   Item(uint32_t offset, uint32_t size) : offset_(offset), size_(size) { }
 
-  uint32_t offset_ = 0;
+  // 0 is the dex file header and shouldn't be a valid offset for any part of the dex file.
+  static constexpr uint32_t kOffsetUnassigned = 0u;
+
+  // Start out unassigned.
+  uint32_t offset_ = kOffsetUnassigned;
   uint32_t size_ = 0;
 };
 
diff --git a/dexlayout/dex_ir_builder.cc b/dexlayout/dex_ir_builder.cc
index bd3e1fa718..924dfe076a 100644
--- a/dexlayout/dex_ir_builder.cc
+++ b/dexlayout/dex_ir_builder.cc
@@ -26,7 +26,7 @@ namespace dex_ir {
 
 static void CheckAndSetRemainingOffsets(const DexFile& dex_file, Collections* collections);
 
-Header* DexIrBuilder(const DexFile& dex_file) {
+Header* DexIrBuilder(const DexFile& dex_file, bool eagerly_assign_offsets) {
   const DexFile::Header& disk_header = dex_file.GetHeader();
   Header* header = new Header(disk_header.magic_,
                               disk_header.checksum_,
@@ -39,6 +39,7 @@ Header* DexIrBuilder(const DexFile& dex_file) {
                               disk_header.data_size_,
                               disk_header.data_off_);
   Collections& collections = header->GetCollections();
+  collections.SetEagerlyAssignOffsets(eagerly_assign_offsets);
   // Walk the rest of the header fields.
   // StringId table.
   collections.SetStringIdsOffset(disk_header.string_ids_off_);
@@ -74,9 +75,11 @@ Header* DexIrBuilder(const DexFile& dex_file) {
   collections.SetMapListOffset(disk_header.map_off_);
   // CallSiteIds and MethodHandleItems.
   collections.CreateCallSitesAndMethodHandles(dex_file);
-
   CheckAndSetRemainingOffsets(dex_file, &collections);
 
+  // Sort the vectors by the map order (same order as the file).
+  collections.SortVectorsByMapOrder();
+
   return header;
 }
 
diff --git a/dexlayout/dex_ir_builder.h b/dexlayout/dex_ir_builder.h
index c53157b5fc..4d4b4e8699 100644
--- a/dexlayout/dex_ir_builder.h
+++ b/dexlayout/dex_ir_builder.h
@@ -24,7 +24,9 @@
 namespace art {
 namespace dex_ir {
 
-dex_ir::Header* DexIrBuilder(const DexFile& dex_file);
+// Eagerly assign offsets assigns offsets based on the original offsets in the input dex file. If
+// this not done, dex_ir::Item::GetOffset will abort when reading uninitialized offsets.
+dex_ir::Header* DexIrBuilder(const DexFile& dex_file, bool eagerly_assign_offsets);
 
 }  // namespace dex_ir
 }  // namespace art
diff --git a/dexlayout/dex_writer.cc b/dexlayout/dex_writer.cc
index 4895ab6957..c85bca0d92 100644
--- a/dexlayout/dex_writer.cc
+++ b/dexlayout/dex_writer.cc
@@ -18,17 +18,37 @@
 
 #include <stdint.h>
 
-#include <queue>
 #include <vector>
 
 #include "cdex/compact_dex_file.h"
 #include "compact_dex_writer.h"
+#include "dex_file_layout.h"
 #include "dex_file_types.h"
+#include "dexlayout.h"
 #include "standard_dex_file.h"
 #include "utf.h"
 
 namespace art {
 
+static constexpr uint32_t kDataSectionAlignment = sizeof(uint32_t) * 2;
+static constexpr uint32_t kDexSectionWordAlignment = 4;
+
+static constexpr uint32_t SectionAlignment(DexFile::MapItemType type) {
+  switch (type) {
+    case DexFile::kDexTypeClassDataItem:
+    case DexFile::kDexTypeStringDataItem:
+    case DexFile::kDexTypeDebugInfoItem:
+    case DexFile::kDexTypeAnnotationItem:
+    case DexFile::kDexTypeEncodedArrayItem:
+      return alignof(uint8_t);
+
+    default:
+      // All other sections are kDexAlignedSection.
+      return kDexSectionWordAlignment;
+  }
+}
+
+
 size_t EncodeIntValue(int32_t value, uint8_t* buffer) {
   size_t length = 0;
   if (value >= 0) {
@@ -245,130 +265,213 @@ size_t DexWriter::WriteEncodedMethods(dex_ir::MethodItemVector* methods, size_t
   return offset - original_offset;
 }
 
-void DexWriter::WriteStrings() {
-  uint32_t string_data_off[1];
+// TODO: Refactor this to remove duplicated boiler plate. One way to do this is adding
+// function that takes a CollectionVector<T> and uses overloading.
+uint32_t DexWriter::WriteStringIds(uint32_t offset, bool reserve_only) {
+  const uint32_t start = offset;
   for (std::unique_ptr<dex_ir::StringId>& string_id : header_->GetCollections().StringIds()) {
-    string_data_off[0] = string_id->DataItem()->GetOffset();
-    Write(string_data_off, string_id->GetSize(), string_id->GetOffset());
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeStringIdItem));
+    if (reserve_only) {
+      offset += string_id->GetSize();
+    } else {
+      uint32_t string_data_off = string_id->DataItem()->GetOffset();
+      offset += Write(&string_data_off, string_id->GetSize(), offset);
+    }
   }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetStringIdsOffset(start);
+  }
+  return offset - start;
+}
 
-  for (auto& string_data_pair : header_->GetCollections().StringDatas()) {
-    std::unique_ptr<dex_ir::StringData>& string_data = string_data_pair.second;
-    uint32_t offset = string_data->GetOffset();
+uint32_t DexWriter::WriteStringDatas(uint32_t offset) {
+  const uint32_t start = offset;
+  for (std::unique_ptr<dex_ir::StringData>& string_data : header_->GetCollections().StringDatas()) {
+    ProcessOffset(&offset, string_data.get());
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeStringDataItem));
     offset += WriteUleb128(CountModifiedUtf8Chars(string_data->Data()), offset);
-    Write(string_data->Data(), strlen(string_data->Data()), offset);
+    // Skip null terminator (already zeroed out, no need to write).
+    offset += Write(string_data->Data(), strlen(string_data->Data()), offset) + 1u;
+  }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetStringDatasOffset(start);
   }
+  return offset - start;
 }
 
-void DexWriter::WriteTypes() {
+uint32_t DexWriter::WriteTypeIds(uint32_t offset) {
   uint32_t descriptor_idx[1];
+  const uint32_t start = offset;
   for (std::unique_ptr<dex_ir::TypeId>& type_id : header_->GetCollections().TypeIds()) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeTypeIdItem));
+    ProcessOffset(&offset, type_id.get());
     descriptor_idx[0] = type_id->GetStringId()->GetIndex();
-    Write(descriptor_idx, type_id->GetSize(), type_id->GetOffset());
+    offset += Write(descriptor_idx, type_id->GetSize(), offset);
   }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetTypeIdsOffset(start);
+  }
+  return offset - start;
 }
 
-void DexWriter::WriteTypeLists() {
+uint32_t DexWriter::WriteTypeLists(uint32_t offset) {
   uint32_t size[1];
   uint16_t list[1];
-  for (auto& type_list_pair : header_->GetCollections().TypeLists()) {
-    std::unique_ptr<dex_ir::TypeList>& type_list = type_list_pair.second;
+  const uint32_t start = offset;
+  for (std::unique_ptr<dex_ir::TypeList>& type_list : header_->GetCollections().TypeLists()) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeTypeList));
     size[0] = type_list->GetTypeList()->size();
-    uint32_t offset = type_list->GetOffset();
+    ProcessOffset(&offset, type_list.get());
     offset += Write(size, sizeof(uint32_t), offset);
     for (const dex_ir::TypeId* type_id : *type_list->GetTypeList()) {
       list[0] = type_id->GetIndex();
       offset += Write(list, sizeof(uint16_t), offset);
     }
   }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetTypeListsOffset(start);
+  }
+  return offset - start;
 }
 
-void DexWriter::WriteProtos() {
+uint32_t DexWriter::WriteProtoIds(uint32_t offset, bool reserve_only) {
   uint32_t buffer[3];
+  const uint32_t start = offset;
   for (std::unique_ptr<dex_ir::ProtoId>& proto_id : header_->GetCollections().ProtoIds()) {
-    buffer[0] = proto_id->Shorty()->GetIndex();
-    buffer[1] = proto_id->ReturnType()->GetIndex();
-    buffer[2] = proto_id->Parameters() == nullptr ? 0 : proto_id->Parameters()->GetOffset();
-    Write(buffer, proto_id->GetSize(), proto_id->GetOffset());
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeProtoIdItem));
+    ProcessOffset(&offset, proto_id.get());
+    if (reserve_only) {
+      offset += proto_id->GetSize();
+    } else {
+      buffer[0] = proto_id->Shorty()->GetIndex();
+      buffer[1] = proto_id->ReturnType()->GetIndex();
+      buffer[2] = proto_id->Parameters() == nullptr ? 0 : proto_id->Parameters()->GetOffset();
+      offset += Write(buffer, proto_id->GetSize(), offset);
+    }
   }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetProtoIdsOffset(start);
+  }
+  return offset - start;
 }
 
-void DexWriter::WriteFields() {
+uint32_t DexWriter::WriteFieldIds(uint32_t offset) {
   uint16_t buffer[4];
+  const uint32_t start = offset;
   for (std::unique_ptr<dex_ir::FieldId>& field_id : header_->GetCollections().FieldIds()) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeFieldIdItem));
+    ProcessOffset(&offset, field_id.get());
     buffer[0] = field_id->Class()->GetIndex();
     buffer[1] = field_id->Type()->GetIndex();
     buffer[2] = field_id->Name()->GetIndex();
     buffer[3] = field_id->Name()->GetIndex() >> 16;
-    Write(buffer, field_id->GetSize(), field_id->GetOffset());
+    offset += Write(buffer, field_id->GetSize(), offset);
+  }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetFieldIdsOffset(start);
   }
+  return offset - start;
 }
 
-void DexWriter::WriteMethods() {
+uint32_t DexWriter::WriteMethodIds(uint32_t offset) {
   uint16_t buffer[4];
+  const uint32_t start = offset;
   for (std::unique_ptr<dex_ir::MethodId>& method_id : header_->GetCollections().MethodIds()) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeMethodIdItem));
+    ProcessOffset(&offset, method_id.get());
     buffer[0] = method_id->Class()->GetIndex();
     buffer[1] = method_id->Proto()->GetIndex();
     buffer[2] = method_id->Name()->GetIndex();
     buffer[3] = method_id->Name()->GetIndex() >> 16;
-    Write(buffer, method_id->GetSize(), method_id->GetOffset());
+    offset += Write(buffer, method_id->GetSize(), offset);
+  }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetMethodIdsOffset(start);
   }
+  return offset - start;
 }
 
-void DexWriter::WriteEncodedArrays() {
-  for (auto& encoded_array_pair : header_->GetCollections().EncodedArrayItems()) {
-    std::unique_ptr<dex_ir::EncodedArrayItem>& encoded_array = encoded_array_pair.second;
-    WriteEncodedArray(encoded_array->GetEncodedValues(), encoded_array->GetOffset());
+uint32_t DexWriter::WriteEncodedArrays(uint32_t offset) {
+  const uint32_t start = offset;
+  for (std::unique_ptr<dex_ir::EncodedArrayItem>& encoded_array :
+      header_->GetCollections().EncodedArrayItems()) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeEncodedArrayItem));
+    ProcessOffset(&offset, encoded_array.get());
+    offset += WriteEncodedArray(encoded_array->GetEncodedValues(), offset);
   }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetEncodedArrayItemsOffset(start);
+  }
+  return offset - start;
 }
 
-void DexWriter::WriteAnnotations() {
+uint32_t DexWriter::WriteAnnotations(uint32_t offset) {
   uint8_t visibility[1];
-  for (auto& annotation_pair : header_->GetCollections().AnnotationItems()) {
-    std::unique_ptr<dex_ir::AnnotationItem>& annotation = annotation_pair.second;
+  const uint32_t start = offset;
+  for (std::unique_ptr<dex_ir::AnnotationItem>& annotation :
+      header_->GetCollections().AnnotationItems()) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeAnnotationItem));
     visibility[0] = annotation->GetVisibility();
-    size_t offset = annotation->GetOffset();
+    ProcessOffset(&offset, annotation.get());
     offset += Write(visibility, sizeof(uint8_t), offset);
-    WriteEncodedAnnotation(annotation->GetAnnotation(), offset);
+    offset += WriteEncodedAnnotation(annotation->GetAnnotation(), offset);
+  }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetAnnotationItemsOffset(start);
   }
+  return offset - start;
 }
 
-void DexWriter::WriteAnnotationSets() {
+uint32_t DexWriter::WriteAnnotationSets(uint32_t offset) {
   uint32_t size[1];
   uint32_t annotation_off[1];
-  for (auto& annotation_set_pair : header_->GetCollections().AnnotationSetItems()) {
-    std::unique_ptr<dex_ir::AnnotationSetItem>& annotation_set = annotation_set_pair.second;
+  const uint32_t start = offset;
+  for (std::unique_ptr<dex_ir::AnnotationSetItem>& annotation_set :
+      header_->GetCollections().AnnotationSetItems()) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeAnnotationSetItem));
     size[0] = annotation_set->GetItems()->size();
-    size_t offset = annotation_set->GetOffset();
+    ProcessOffset(&offset, annotation_set.get());
     offset += Write(size, sizeof(uint32_t), offset);
     for (dex_ir::AnnotationItem* annotation : *annotation_set->GetItems()) {
       annotation_off[0] = annotation->GetOffset();
       offset += Write(annotation_off, sizeof(uint32_t), offset);
     }
   }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetAnnotationSetItemsOffset(start);
+  }
+  return offset - start;
 }
 
-void DexWriter::WriteAnnotationSetRefs() {
+uint32_t DexWriter::WriteAnnotationSetRefs(uint32_t offset) {
   uint32_t size[1];
   uint32_t annotations_off[1];
-  for (auto& anno_set_ref_pair : header_->GetCollections().AnnotationSetRefLists()) {
-    std::unique_ptr<dex_ir::AnnotationSetRefList>& annotation_set_ref = anno_set_ref_pair.second;
+  const uint32_t start = offset;
+  for (std::unique_ptr<dex_ir::AnnotationSetRefList>& annotation_set_ref :
+      header_->GetCollections().AnnotationSetRefLists()) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeAnnotationSetRefList));
     size[0] = annotation_set_ref->GetItems()->size();
-    size_t offset = annotation_set_ref->GetOffset();
+    ProcessOffset(&offset, annotation_set_ref.get());
     offset += Write(size, sizeof(uint32_t), offset);
     for (dex_ir::AnnotationSetItem* annotation_set : *annotation_set_ref->GetItems()) {
       annotations_off[0] = annotation_set == nullptr ? 0 : annotation_set->GetOffset();
       offset += Write(annotations_off, sizeof(uint32_t), offset);
     }
   }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetAnnotationSetRefListsOffset(start);
+  }
+  return offset - start;
 }
 
-void DexWriter::WriteAnnotationsDirectories() {
+uint32_t DexWriter::WriteAnnotationsDirectories(uint32_t offset) {
   uint32_t directory_buffer[4];
   uint32_t annotation_buffer[2];
-  for (auto& annotations_directory_pair : header_->GetCollections().AnnotationsDirectoryItems()) {
-    std::unique_ptr<dex_ir::AnnotationsDirectoryItem>& annotations_directory =
-        annotations_directory_pair.second;
+  const uint32_t start = offset;
+  for (std::unique_ptr<dex_ir::AnnotationsDirectoryItem>& annotations_directory :
+      header_->GetCollections().AnnotationsDirectoryItems()) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeAnnotationsDirectoryItem));
+    ProcessOffset(&offset, annotations_directory.get());
     directory_buffer[0] = annotations_directory->GetClassAnnotation() == nullptr ? 0 :
         annotations_directory->GetClassAnnotation()->GetOffset();
     directory_buffer[1] = annotations_directory->GetFieldAnnotations() == nullptr ? 0 :
@@ -377,7 +480,6 @@ void DexWriter::WriteAnnotationsDirectories() {
         annotations_directory->GetMethodAnnotations()->size();
     directory_buffer[3] = annotations_directory->GetParameterAnnotations() == nullptr ? 0 :
         annotations_directory->GetParameterAnnotations()->size();
-    uint32_t offset = annotations_directory->GetOffset();
     offset += Write(directory_buffer, 4 * sizeof(uint32_t), offset);
     if (annotations_directory->GetFieldAnnotations() != nullptr) {
       for (std::unique_ptr<dex_ir::FieldAnnotation>& field :
@@ -404,27 +506,55 @@ void DexWriter::WriteAnnotationsDirectories() {
       }
     }
   }
-}
-
-void DexWriter::WriteDebugInfoItems() {
-  for (auto& debug_info_pair : header_->GetCollections().DebugInfoItems()) {
-    std::unique_ptr<dex_ir::DebugInfoItem>& debug_info = debug_info_pair.second;
-    Write(debug_info->GetDebugInfo(), debug_info->GetDebugInfoSize(), debug_info->GetOffset());
-  }
-}
-
-void DexWriter::WriteCodeItems() {
-  uint16_t uint16_buffer[4];
-  uint32_t uint32_buffer[2];
-  for (auto& code_item_pair : header_->GetCollections().CodeItems()) {
-    std::unique_ptr<dex_ir::CodeItem>& code_item = code_item_pair.second;
-    uint16_buffer[0] = code_item->RegistersSize();
-    uint16_buffer[1] = code_item->InsSize();
-    uint16_buffer[2] = code_item->OutsSize();
-    uint16_buffer[3] = code_item->TriesSize();
-    uint32_buffer[0] = code_item->DebugInfo() == nullptr ? 0 : code_item->DebugInfo()->GetOffset();
-    uint32_buffer[1] = code_item->InsnsSize();
-    size_t offset = code_item->GetOffset();
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetAnnotationsDirectoryItemsOffset(start);
+  }
+  return offset - start;
+}
+
+uint32_t DexWriter::WriteDebugInfoItems(uint32_t offset) {
+  const uint32_t start = offset;
+  for (std::unique_ptr<dex_ir::DebugInfoItem>& debug_info :
+      header_->GetCollections().DebugInfoItems()) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeDebugInfoItem));
+    ProcessOffset(&offset, debug_info.get());
+    offset += Write(debug_info->GetDebugInfo(), debug_info->GetDebugInfoSize(), offset);
+  }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetDebugInfoItemsOffset(start);
+  }
+  return offset - start;
+}
+
+uint32_t DexWriter::WriteCodeItems(uint32_t offset, bool reserve_only) {
+  DexLayoutSection* code_section = nullptr;
+  if (!reserve_only && dex_layout_ != nullptr) {
+    code_section = &dex_layout_->GetSections().sections_[static_cast<size_t>(
+        DexLayoutSections::SectionType::kSectionTypeCode)];
+  }
+  uint16_t uint16_buffer[4] = {};
+  uint32_t uint32_buffer[2] = {};
+  uint32_t start = offset;
+  for (auto& code_item : header_->GetCollections().CodeItems()) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeCodeItem));
+    ProcessOffset(&offset, code_item.get());
+    if (!reserve_only) {
+      uint16_buffer[0] = code_item->RegistersSize();
+      uint16_buffer[1] = code_item->InsSize();
+      uint16_buffer[2] = code_item->OutsSize();
+      uint16_buffer[3] = code_item->TriesSize();
+      uint32_buffer[0] = code_item->DebugInfo() == nullptr ? 0 :
+          code_item->DebugInfo()->GetOffset();
+      uint32_buffer[1] = code_item->InsnsSize();
+      // Only add the section hotness info once.
+      if (code_section != nullptr) {
+        auto it = dex_layout_->LayoutHotnessInfo().code_item_layout_.find(code_item.get());
+        if (it != dex_layout_->LayoutHotnessInfo().code_item_layout_.end()) {
+          code_section->parts_[static_cast<size_t>(it->second)].CombineSection(
+              code_item->GetOffset(), code_item->GetOffset() + code_item->GetSize());
+        }
+      }
+    }
     offset += Write(uint16_buffer, 4 * sizeof(uint16_t), offset);
     offset += Write(uint32_buffer, 2 * sizeof(uint32_t), offset);
     offset += Write(code_item->Insns(), code_item->InsnsSize() * sizeof(uint16_t), offset);
@@ -443,7 +573,7 @@ void DexWriter::WriteCodeItems() {
         offset += Write(insn_count_and_handler_off, 2 * sizeof(uint16_t), offset);
       }
       // Leave offset pointing to the end of the try items.
-      WriteUleb128(code_item->Handlers()->size(), offset);
+      UNUSED(WriteUleb128(code_item->Handlers()->size(), offset));
       for (std::unique_ptr<const dex_ir::CatchHandler>& handlers : *code_item->Handlers()) {
         size_t list_offset = offset + handlers->GetListOffset();
         uint32_t size = handlers->HasCatchAll() ? (handlers->GetHandlers()->size() - 1) * -1 :
@@ -457,32 +587,52 @@ void DexWriter::WriteCodeItems() {
         }
       }
     }
+    // TODO: Clean this up to properly calculate the size instead of assuming it doesn't change.
+    offset = code_item->GetOffset() + code_item->GetSize();
   }
+
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetCodeItemsOffset(start);
+  }
+  return offset - start;
 }
 
-void DexWriter::WriteClasses() {
+uint32_t DexWriter::WriteClassDefs(uint32_t offset, bool reserve_only) {
+  const uint32_t start = offset;
   uint32_t class_def_buffer[8];
   for (std::unique_ptr<dex_ir::ClassDef>& class_def : header_->GetCollections().ClassDefs()) {
-    class_def_buffer[0] = class_def->ClassType()->GetIndex();
-    class_def_buffer[1] = class_def->GetAccessFlags();
-    class_def_buffer[2] = class_def->Superclass() == nullptr ? dex::kDexNoIndex :
-        class_def->Superclass()->GetIndex();
-    class_def_buffer[3] = class_def->InterfacesOffset();
-    class_def_buffer[4] = class_def->SourceFile() == nullptr ? dex::kDexNoIndex :
-        class_def->SourceFile()->GetIndex();
-    class_def_buffer[5] = class_def->Annotations() == nullptr ? 0 :
-        class_def->Annotations()->GetOffset();
-    class_def_buffer[6] = class_def->GetClassData() == nullptr ? 0 :
-        class_def->GetClassData()->GetOffset();
-    class_def_buffer[7] = class_def->StaticValues() == nullptr ? 0 :
-        class_def->StaticValues()->GetOffset();
-    size_t offset = class_def->GetOffset();
-    Write(class_def_buffer, class_def->GetSize(), offset);
-  }
-
-  for (auto& class_data_pair : header_->GetCollections().ClassDatas()) {
-    std::unique_ptr<dex_ir::ClassData>& class_data = class_data_pair.second;
-    size_t offset = class_data->GetOffset();
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeClassDefItem));
+    if (reserve_only) {
+      offset += class_def->GetSize();
+    } else {
+      class_def_buffer[0] = class_def->ClassType()->GetIndex();
+      class_def_buffer[1] = class_def->GetAccessFlags();
+      class_def_buffer[2] = class_def->Superclass() == nullptr ? dex::kDexNoIndex :
+          class_def->Superclass()->GetIndex();
+      class_def_buffer[3] = class_def->InterfacesOffset();
+      class_def_buffer[4] = class_def->SourceFile() == nullptr ? dex::kDexNoIndex :
+          class_def->SourceFile()->GetIndex();
+      class_def_buffer[5] = class_def->Annotations() == nullptr ? 0 :
+          class_def->Annotations()->GetOffset();
+      class_def_buffer[6] = class_def->GetClassData() == nullptr ? 0 :
+          class_def->GetClassData()->GetOffset();
+      class_def_buffer[7] = class_def->StaticValues() == nullptr ? 0 :
+          class_def->StaticValues()->GetOffset();
+      offset += Write(class_def_buffer, class_def->GetSize(), offset);
+    }
+  }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetClassDefsOffset(start);
+  }
+  return offset - start;
+}
+
+uint32_t DexWriter::WriteClassDatas(uint32_t offset) {
+  const uint32_t start = offset;
+  for (const std::unique_ptr<dex_ir::ClassData>& class_data :
+      header_->GetCollections().ClassDatas()) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeClassDataItem));
+    ProcessOffset(&offset, class_data.get());
     offset += WriteUleb128(class_data->StaticFields()->size(), offset);
     offset += WriteUleb128(class_data->InstanceFields()->size(), offset);
     offset += WriteUleb128(class_data->DirectMethods()->size(), offset);
@@ -492,139 +642,134 @@ void DexWriter::WriteClasses() {
     offset += WriteEncodedMethods(class_data->DirectMethods(), offset);
     offset += WriteEncodedMethods(class_data->VirtualMethods(), offset);
   }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetClassDatasOffset(start);
+  }
+  return offset - start;
 }
 
-void DexWriter::WriteCallSites() {
+uint32_t DexWriter::WriteCallSiteIds(uint32_t offset, bool reserve_only) {
+  const uint32_t start = offset;
   uint32_t call_site_off[1];
   for (std::unique_ptr<dex_ir::CallSiteId>& call_site_id :
       header_->GetCollections().CallSiteIds()) {
-    call_site_off[0] = call_site_id->CallSiteItem()->GetOffset();
-    Write(call_site_off, call_site_id->GetSize(), call_site_id->GetOffset());
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeCallSiteIdItem));
+    if (reserve_only) {
+      offset += call_site_id->GetSize();
+    } else {
+      call_site_off[0] = call_site_id->CallSiteItem()->GetOffset();
+      offset += Write(call_site_off, call_site_id->GetSize(), offset);
+    }
   }
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetCallSiteIdsOffset(start);
+  }
+  return offset - start;
 }
 
-void DexWriter::WriteMethodHandles() {
+uint32_t DexWriter::WriteMethodHandles(uint32_t offset) {
+  const uint32_t start = offset;
   uint16_t method_handle_buff[4];
   for (std::unique_ptr<dex_ir::MethodHandleItem>& method_handle :
       header_->GetCollections().MethodHandleItems()) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeMethodHandleItem));
     method_handle_buff[0] = static_cast<uint16_t>(method_handle->GetMethodHandleType());
     method_handle_buff[1] = 0;  // unused.
     method_handle_buff[2] = method_handle->GetFieldOrMethodId()->GetIndex();
     method_handle_buff[3] = 0;  // unused.
-    Write(method_handle_buff, method_handle->GetSize(), method_handle->GetOffset());
+    offset += Write(method_handle_buff, method_handle->GetSize(), offset);
   }
-}
-
-struct MapItemContainer {
-  MapItemContainer(uint32_t type, uint32_t size, uint32_t offset)
-      : type_(type), size_(size), offset_(offset) { }
-
-  bool operator<(const MapItemContainer& other) const {
-    return offset_ > other.offset_;
-  }
-
-  uint32_t type_;
-  uint32_t size_;
-  uint32_t offset_;
-};
-
-void DexWriter::WriteMapItem() {
-  dex_ir::Collections& collection = header_->GetCollections();
-  std::priority_queue<MapItemContainer> queue;
-
-  // Header and index section.
-  queue.push(MapItemContainer(DexFile::kDexTypeHeaderItem, 1, 0));
-  if (collection.StringIdsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeStringIdItem, collection.StringIdsSize(),
-        collection.StringIdsOffset()));
-  }
-  if (collection.TypeIdsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeTypeIdItem, collection.TypeIdsSize(),
-        collection.TypeIdsOffset()));
-  }
-  if (collection.ProtoIdsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeProtoIdItem, collection.ProtoIdsSize(),
-        collection.ProtoIdsOffset()));
-  }
-  if (collection.FieldIdsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeFieldIdItem, collection.FieldIdsSize(),
-        collection.FieldIdsOffset()));
-  }
-  if (collection.MethodIdsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeMethodIdItem, collection.MethodIdsSize(),
-        collection.MethodIdsOffset()));
-  }
-  if (collection.ClassDefsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeClassDefItem, collection.ClassDefsSize(),
-        collection.ClassDefsOffset()));
-  }
-  if (collection.CallSiteIdsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeCallSiteIdItem, collection.CallSiteIdsSize(),
-        collection.CallSiteIdsOffset()));
-  }
-  if (collection.MethodHandleItemsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeMethodHandleItem,
-        collection.MethodHandleItemsSize(), collection.MethodHandleItemsOffset()));
-  }
-
-  // Data section.
-  queue.push(MapItemContainer(DexFile::kDexTypeMapList, 1, collection.MapListOffset()));
-  if (collection.TypeListsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeTypeList, collection.TypeListsSize(),
-        collection.TypeListsOffset()));
-  }
-  if (collection.AnnotationSetRefListsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeAnnotationSetRefList,
-        collection.AnnotationSetRefListsSize(), collection.AnnotationSetRefListsOffset()));
-  }
-  if (collection.AnnotationSetItemsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeAnnotationSetItem,
-        collection.AnnotationSetItemsSize(), collection.AnnotationSetItemsOffset()));
-  }
-  if (collection.ClassDatasSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeClassDataItem, collection.ClassDatasSize(),
-        collection.ClassDatasOffset()));
-  }
-  if (collection.CodeItemsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeCodeItem, collection.CodeItemsSize(),
-        collection.CodeItemsOffset()));
-  }
-  if (collection.StringDatasSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeStringDataItem, collection.StringDatasSize(),
-        collection.StringDatasOffset()));
-  }
-  if (collection.DebugInfoItemsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeDebugInfoItem, collection.DebugInfoItemsSize(),
-        collection.DebugInfoItemsOffset()));
-  }
-  if (collection.AnnotationItemsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeAnnotationItem, collection.AnnotationItemsSize(),
-        collection.AnnotationItemsOffset()));
-  }
-  if (collection.EncodedArrayItemsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeEncodedArrayItem,
-        collection.EncodedArrayItemsSize(), collection.EncodedArrayItemsOffset()));
-  }
-  if (collection.AnnotationsDirectoryItemsSize() != 0) {
-    queue.push(MapItemContainer(DexFile::kDexTypeAnnotationsDirectoryItem,
-        collection.AnnotationsDirectoryItemsSize(), collection.AnnotationsDirectoryItemsOffset()));
+  if (compute_offsets_ && start != offset) {
+    header_->GetCollections().SetMethodHandleItemsOffset(start);
   }
+  return offset - start;
+}
 
-  uint32_t offset = collection.MapListOffset();
+uint32_t DexWriter::WriteMapItems(uint32_t offset, MapItemQueue* queue) {
+  // All the sections should already have been added.
   uint16_t uint16_buffer[2];
   uint32_t uint32_buffer[2];
   uint16_buffer[1] = 0;
-  uint32_buffer[0] = queue.size();
+  uint32_buffer[0] = queue->size();
+  const uint32_t start = offset;
   offset += Write(uint32_buffer, sizeof(uint32_t), offset);
-  while (!queue.empty()) {
-    const MapItemContainer& map_item = queue.top();
+  while (!queue->empty()) {
+    const MapItem& map_item = queue->top();
     uint16_buffer[0] = map_item.type_;
     uint32_buffer[0] = map_item.size_;
     uint32_buffer[1] = map_item.offset_;
     offset += Write(uint16_buffer, 2 * sizeof(uint16_t), offset);
     offset += Write(uint32_buffer, 2 * sizeof(uint32_t), offset);
-    queue.pop();
+    queue->pop();
   }
+  return offset - start;
+}
+
+uint32_t DexWriter::GenerateAndWriteMapItems(uint32_t offset) {
+  dex_ir::Collections& collection = header_->GetCollections();
+  MapItemQueue queue;
+
+  // Header and index section.
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeHeaderItem, 1, 0));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeStringIdItem,
+                              collection.StringIdsSize(),
+                              collection.StringIdsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeTypeIdItem,
+                              collection.TypeIdsSize(),
+                              collection.TypeIdsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeProtoIdItem,
+                              collection.ProtoIdsSize(),
+                              collection.ProtoIdsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeFieldIdItem,
+                              collection.FieldIdsSize(),
+                              collection.FieldIdsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeMethodIdItem,
+                              collection.MethodIdsSize(),
+                              collection.MethodIdsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeClassDefItem,
+                              collection.ClassDefsSize(),
+                              collection.ClassDefsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeCallSiteIdItem,
+                              collection.CallSiteIdsSize(),
+                              collection.CallSiteIdsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeMethodHandleItem,
+                              collection.MethodHandleItemsSize(),
+                              collection.MethodHandleItemsOffset()));
+  // Data section.
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeMapList, 1, collection.MapListOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeTypeList,
+                              collection.TypeListsSize(),
+                              collection.TypeListsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeAnnotationSetRefList,
+                              collection.AnnotationSetRefListsSize(),
+                              collection.AnnotationSetRefListsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeAnnotationSetItem,
+                              collection.AnnotationSetItemsSize(),
+                              collection.AnnotationSetItemsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeClassDataItem,
+                              collection.ClassDatasSize(),
+                              collection.ClassDatasOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeCodeItem,
+                              collection.CodeItemsSize(),
+                              collection.CodeItemsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeStringDataItem,
+                              collection.StringDatasSize(),
+                              collection.StringDatasOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeDebugInfoItem,
+                              collection.DebugInfoItemsSize(),
+                              collection.DebugInfoItemsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeAnnotationItem,
+                              collection.AnnotationItemsSize(),
+                              collection.AnnotationItemsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeEncodedArrayItem,
+                              collection.EncodedArrayItemsSize(),
+                              collection.EncodedArrayItemsOffset()));
+  queue.AddIfNotEmpty(MapItem(DexFile::kDexTypeAnnotationsDirectoryItem,
+                              collection.AnnotationsDirectoryItemsSize(),
+                              collection.AnnotationsDirectoryItemsOffset()));
+
+  // Write the map items.
+  return WriteMapItems(offset, &queue);
 }
 
 void DexWriter::WriteHeader() {
@@ -657,38 +802,110 @@ void DexWriter::WriteHeader() {
   header.data_off_ = header_->DataOffset();
 
   static_assert(sizeof(header) == 0x70, "Size doesn't match dex spec");
-  Write(reinterpret_cast<uint8_t*>(&header), sizeof(header), 0u);
+  UNUSED(Write(reinterpret_cast<uint8_t*>(&header), sizeof(header), 0u));
 }
 
 void DexWriter::WriteMemMap() {
-  WriteStrings();
-  WriteTypes();
-  WriteTypeLists();
-  WriteProtos();
-  WriteFields();
-  WriteMethods();
-  WriteEncodedArrays();
-  WriteAnnotations();
-  WriteAnnotationSets();
-  WriteAnnotationSetRefs();
-  WriteAnnotationsDirectories();
-  WriteDebugInfoItems();
-  WriteCodeItems();
-  WriteClasses();
-  WriteCallSites();
-  WriteMethodHandles();
-  WriteMapItem();
+  // Starting offset is right after the header.
+  uint32_t offset = sizeof(StandardDexFile::Header);
+
+  dex_ir::Collections& collection = header_->GetCollections();
+
+  // Based on: https://source.android.com/devices/tech/dalvik/dex-format
+  // Since the offsets may not be calculated already, the writing must be done in the correct order.
+  const uint32_t string_ids_offset = offset;
+  offset += WriteStringIds(offset, /*reserve_only*/ true);
+  offset += WriteTypeIds(offset);
+  const uint32_t proto_ids_offset = offset;
+  offset += WriteProtoIds(offset, /*reserve_only*/ true);
+  offset += WriteFieldIds(offset);
+  offset += WriteMethodIds(offset);
+  const uint32_t class_defs_offset = offset;
+  offset += WriteClassDefs(offset, /*reserve_only*/ true);
+  const uint32_t call_site_ids_offset = offset;
+  offset += WriteCallSiteIds(offset, /*reserve_only*/ true);
+  offset += WriteMethodHandles(offset);
+
+  uint32_t data_offset_ = 0u;
+  if (compute_offsets_) {
+    // Data section.
+    offset = RoundUp(offset, kDataSectionAlignment);
+    data_offset_ = offset;
+  }
+
+  // Write code item first to minimize the space required for encoded methods.
+  // Reserve code item space since we need the debug offsets to actually write them.
+  const uint32_t code_items_offset = offset;
+  offset += WriteCodeItems(offset, /*reserve_only*/ true);
+  // Write debug info section.
+  offset += WriteDebugInfoItems(offset);
+  // Actually write code items since debug info offsets are calculated now.
+  WriteCodeItems(code_items_offset, /*reserve_only*/ false);
+
+  offset += WriteEncodedArrays(offset);
+  offset += WriteAnnotations(offset);
+  offset += WriteAnnotationSets(offset);
+  offset += WriteAnnotationSetRefs(offset);
+  offset += WriteAnnotationsDirectories(offset);
+  offset += WriteTypeLists(offset);
+  offset += WriteClassDatas(offset);
+  offset += WriteStringDatas(offset);
+
+  // Write delayed id sections that depend on data sections.
+  WriteStringIds(string_ids_offset, /*reserve_only*/ false);
+  WriteProtoIds(proto_ids_offset, /*reserve_only*/ false);
+  WriteClassDefs(class_defs_offset, /*reserve_only*/ false);
+  WriteCallSiteIds(call_site_ids_offset, /*reserve_only*/ false);
+
+  // Write the map list.
+  if (compute_offsets_) {
+    offset = RoundUp(offset, SectionAlignment(DexFile::kDexTypeMapList));
+    collection.SetMapListOffset(offset);
+  } else {
+    offset = collection.MapListOffset();
+  }
+  offset += GenerateAndWriteMapItems(offset);
+  offset = RoundUp(offset, kDataSectionAlignment);
+
+  // Map items are included in the data section.
+  if (compute_offsets_) {
+    header_->SetDataSize(offset - data_offset_);
+    if (header_->DataSize() != 0) {
+      // Offset must be zero when the size is zero.
+      header_->SetDataOffset(data_offset_);
+    } else {
+      header_->SetDataOffset(0u);
+    }
+  }
+
+  // TODO: Write link data?
+
+  // Write header last.
+  if (compute_offsets_) {
+    header_->SetFileSize(offset);
+  }
   WriteHeader();
 }
 
-void DexWriter::Output(dex_ir::Header* header, MemMap* mem_map, CompactDexLevel compact_dex_level) {
+void DexWriter::Output(dex_ir::Header* header,
+                       MemMap* mem_map,
+                       DexLayout* dex_layout,
+                       bool compute_offsets,
+                       CompactDexLevel compact_dex_level) {
+  CHECK(dex_layout != nullptr);
   std::unique_ptr<DexWriter> writer;
   if (compact_dex_level != CompactDexLevel::kCompactDexLevelNone) {
-    writer.reset(new CompactDexWriter(header, mem_map, compact_dex_level));
+    writer.reset(new CompactDexWriter(header, mem_map, dex_layout, compact_dex_level));
   } else {
-    writer.reset(new DexWriter(header, mem_map));
+    writer.reset(new DexWriter(header, mem_map, dex_layout, compute_offsets));
   }
   writer->WriteMemMap();
 }
 
+void MapItemQueue::AddIfNotEmpty(const MapItem& item) {
+  if (item.size_ != 0) {
+    push(item);
+  }
+}
+
 }  // namespace art
diff --git a/dexlayout/dex_writer.h b/dexlayout/dex_writer.h
index 85d3e7ebf3..c47898e533 100644
--- a/dexlayout/dex_writer.h
+++ b/dexlayout/dex_writer.h
@@ -19,56 +19,119 @@
 #ifndef ART_DEXLAYOUT_DEX_WRITER_H_
 #define ART_DEXLAYOUT_DEX_WRITER_H_
 
+#include <functional>
+
 #include "base/unix_file/fd_file.h"
 #include "cdex/compact_dex_level.h"
 #include "dex_ir.h"
 #include "mem_map.h"
 #include "os.h"
 
+#include <queue>
+
 namespace art {
 
+class DexLayout;
+class DexLayoutHotnessInfo;
+
+struct MapItem {
+  // Not using DexFile::MapItemType since compact dex and standard dex file may have different
+  // sections.
+  MapItem() = default;
+  MapItem(uint32_t type, uint32_t size, uint32_t offset)
+      : type_(type), size_(size), offset_(offset) { }
+
+  // Sort by decreasing order since the priority_queue puts largest elements first.
+  bool operator>(const MapItem& other) const {
+    return offset_ > other.offset_;
+  }
+
+  uint32_t type_ = 0u;
+  uint32_t size_ = 0u;
+  uint32_t offset_ = 0u;
+};
+
+class MapItemQueue : public
+    std::priority_queue<MapItem, std::vector<MapItem>, std::greater<MapItem>> {
+ public:
+  void AddIfNotEmpty(const MapItem& item);
+};
+
 class DexWriter {
  public:
-  DexWriter(dex_ir::Header* header, MemMap* mem_map) : header_(header), mem_map_(mem_map) {}
+  DexWriter(dex_ir::Header* header,
+            MemMap* mem_map,
+            DexLayout* dex_layout,
+            bool compute_offsets)
+      : header_(header),
+        mem_map_(mem_map),
+        dex_layout_(dex_layout),
+        compute_offsets_(compute_offsets) {}
 
-  static void Output(dex_ir::Header* header, MemMap* mem_map, CompactDexLevel compact_dex_level);
+  static void Output(dex_ir::Header* header,
+                     MemMap* mem_map,
+                     DexLayout* dex_layout,
+                     bool compute_offsets,
+                     CompactDexLevel compact_dex_level);
 
   virtual ~DexWriter() {}
 
  protected:
   void WriteMemMap();
 
-  size_t Write(const void* buffer, size_t length, size_t offset);
-  size_t WriteSleb128(uint32_t value, size_t offset);
-  size_t WriteUleb128(uint32_t value, size_t offset);
-  size_t WriteEncodedValue(dex_ir::EncodedValue* encoded_value, size_t offset);
-  size_t WriteEncodedValueHeader(int8_t value_type, size_t value_arg, size_t offset);
-  size_t WriteEncodedArray(dex_ir::EncodedValueVector* values, size_t offset);
-  size_t WriteEncodedAnnotation(dex_ir::EncodedAnnotation* annotation, size_t offset);
-  size_t WriteEncodedFields(dex_ir::FieldItemVector* fields, size_t offset);
-  size_t WriteEncodedMethods(dex_ir::MethodItemVector* methods, size_t offset);
-
-  void WriteStrings();
-  void WriteTypes();
-  void WriteTypeLists();
-  void WriteProtos();
-  void WriteFields();
-  void WriteMethods();
-  void WriteEncodedArrays();
-  void WriteAnnotations();
-  void WriteAnnotationSets();
-  void WriteAnnotationSetRefs();
-  void WriteAnnotationsDirectories();
-  void WriteDebugInfoItems();
-  void WriteCodeItems();
-  void WriteClasses();
-  void WriteCallSites();
-  void WriteMethodHandles();
-  void WriteMapItem();
+  size_t Write(const void* buffer, size_t length, size_t offset) WARN_UNUSED;
+  size_t WriteSleb128(uint32_t value, size_t offset) WARN_UNUSED;
+  size_t WriteUleb128(uint32_t value, size_t offset) WARN_UNUSED;
+  size_t WriteEncodedValue(dex_ir::EncodedValue* encoded_value, size_t offset) WARN_UNUSED;
+  size_t WriteEncodedValueHeader(int8_t value_type, size_t value_arg, size_t offset) WARN_UNUSED;
+  size_t WriteEncodedArray(dex_ir::EncodedValueVector* values, size_t offset) WARN_UNUSED;
+  size_t WriteEncodedAnnotation(dex_ir::EncodedAnnotation* annotation, size_t offset) WARN_UNUSED;
+  size_t WriteEncodedFields(dex_ir::FieldItemVector* fields, size_t offset) WARN_UNUSED;
+  size_t WriteEncodedMethods(dex_ir::MethodItemVector* methods, size_t offset) WARN_UNUSED;
+
+  // Header and id section
   virtual void WriteHeader();
+  // reserve_only means don't write, only reserve space. This is required since the string data
+  // offsets must be assigned.
+  uint32_t WriteStringIds(uint32_t offset, bool reserve_only);
+  uint32_t WriteTypeIds(uint32_t offset);
+  uint32_t WriteProtoIds(uint32_t offset, bool reserve_only);
+  uint32_t WriteFieldIds(uint32_t offset);
+  uint32_t WriteMethodIds(uint32_t offset);
+  uint32_t WriteClassDefs(uint32_t offset, bool reserve_only);
+  uint32_t WriteCallSiteIds(uint32_t offset, bool reserve_only);
+
+  uint32_t WriteEncodedArrays(uint32_t offset);
+  uint32_t WriteAnnotations(uint32_t offset);
+  uint32_t WriteAnnotationSets(uint32_t offset);
+  uint32_t WriteAnnotationSetRefs(uint32_t offset);
+  uint32_t WriteAnnotationsDirectories(uint32_t offset);
+
+  // Data section.
+  uint32_t WriteDebugInfoItems(uint32_t offset);
+  uint32_t WriteCodeItems(uint32_t offset, bool reserve_only);
+  uint32_t WriteTypeLists(uint32_t offset);
+  uint32_t WriteStringDatas(uint32_t offset);
+  uint32_t WriteClassDatas(uint32_t offset);
+  uint32_t WriteMethodHandles(uint32_t offset);
+  uint32_t WriteMapItems(uint32_t offset, MapItemQueue* queue);
+  uint32_t GenerateAndWriteMapItems(uint32_t offset);
+
+  // Process an offset, if compute_offset is set, write into the dex ir item, otherwise read the
+  // existing offset and use that for writing.
+  void ProcessOffset(uint32_t* const offset, dex_ir::Item* item) {
+    if (compute_offsets_) {
+      item->SetOffset(*offset);
+    } else {
+      // Not computing offsets, just use the one in the item.
+      *offset = item->GetOffset();
+    }
+  }
 
   dex_ir::Header* const header_;
   MemMap* const mem_map_;
+  DexLayout* const dex_layout_;
+  bool compute_offsets_;
 
  private:
   DISALLOW_COPY_AND_ASSIGN(DexWriter);
diff --git a/dexlayout/dexdiag.cc b/dexlayout/dexdiag.cc
index c14cd5f807..e83f98ee6b 100644
--- a/dexlayout/dexdiag.cc
+++ b/dexlayout/dexdiag.cc
@@ -289,7 +289,8 @@ static void ProcessOneDexMapping(uint64_t* pagemap,
   // Build a list of the dex file section types, sorted from highest offset to lowest.
   std::vector<dex_ir::DexFileSection> sections;
   {
-    std::unique_ptr<dex_ir::Header> header(dex_ir::DexIrBuilder(*dex_file));
+    std::unique_ptr<dex_ir::Header> header(dex_ir::DexIrBuilder(*dex_file,
+                                                                /*eagerly_assign_offsets*/ true));
     sections = dex_ir::GetSortedDexFileSections(header.get(),
                                                 dex_ir::SortDirection::kSortDescending);
   }
diff --git a/dexlayout/dexlayout.cc b/dexlayout/dexlayout.cc
index dd2e809a92..d904a52f0c 100644
--- a/dexlayout/dexlayout.cc
+++ b/dexlayout/dexlayout.cc
@@ -56,9 +56,6 @@ using android::base::StringPrintf;
 // necessary to ensure the partial order w.r.t. class derivation. TODO: Re-enable (b/68317550).
 static constexpr bool kChangeClassDefOrder = false;
 
-static constexpr uint32_t kDataSectionAlignment = sizeof(uint32_t) * 2;
-static constexpr uint32_t kDexCodeItemAlignment = 4;
-
 /*
  * Flags for use with createAccessFlagStr().
  */
@@ -1564,7 +1561,7 @@ void DexLayout::DumpDexFile() {
   }
 }
 
-std::vector<dex_ir::ClassData*> DexLayout::LayoutClassDefsAndClassData(const DexFile* dex_file) {
+void DexLayout::LayoutClassDefsAndClassData(const DexFile* dex_file) {
   std::vector<dex_ir::ClassDef*> new_class_def_order;
   for (std::unique_ptr<dex_ir::ClassDef>& class_def : header_->GetCollections().ClassDefs()) {
     dex::TypeIndex type_idx(class_def->ClassType()->GetIndex());
@@ -1578,31 +1575,41 @@ std::vector<dex_ir::ClassData*> DexLayout::LayoutClassDefsAndClassData(const Dex
       new_class_def_order.push_back(class_def.get());
     }
   }
-  uint32_t class_defs_offset = header_->GetCollections().ClassDefsOffset();
-  uint32_t class_data_offset = header_->GetCollections().ClassDatasOffset();
   std::unordered_set<dex_ir::ClassData*> visited_class_data;
-  std::vector<dex_ir::ClassData*> new_class_data_order;
-  for (uint32_t i = 0; i < new_class_def_order.size(); ++i) {
-    dex_ir::ClassDef* class_def = new_class_def_order[i];
-    if (kChangeClassDefOrder) {
-      // This produces dex files that violate the spec since the super class class_def is supposed
-      // to occur before any subclasses.
-      class_def->SetIndex(i);
-      class_def->SetOffset(class_defs_offset);
-      class_defs_offset += dex_ir::ClassDef::ItemSize();
-    }
+  size_t class_data_index = 0;
+  dex_ir::CollectionVector<dex_ir::ClassData>::Vector& class_datas =
+      header_->GetCollections().ClassDatas();
+  for (dex_ir::ClassDef* class_def : new_class_def_order) {
     dex_ir::ClassData* class_data = class_def->GetClassData();
     if (class_data != nullptr && visited_class_data.find(class_data) == visited_class_data.end()) {
-      class_data->SetOffset(class_data_offset);
-      class_data_offset += class_data->GetSize();
       visited_class_data.insert(class_data);
-      new_class_data_order.push_back(class_data);
+      // Overwrite the existing vector with the new ordering, note that the sets of objects are
+      // equivalent, but the order changes. This is why this is not a memory leak.
+      // TODO: Consider cleaning this up with a shared_ptr.
+      class_datas[class_data_index].release();
+      class_datas[class_data_index].reset(class_data);
+      ++class_data_index;
+    }
+  }
+  CHECK_EQ(class_data_index, class_datas.size());
+
+  if (kChangeClassDefOrder) {
+    // This currently produces dex files that violate the spec since the super class class_def is
+    // supposed to occur before any subclasses.
+    dex_ir::CollectionVector<dex_ir::ClassDef>::Vector& class_defs =
+        header_->GetCollections().ClassDefs();
+    CHECK_EQ(new_class_def_order.size(), class_defs.size());
+    for (size_t i = 0; i < class_defs.size(); ++i) {
+      // Overwrite the existing vector with the new ordering, note that the sets of objects are
+      // equivalent, but the order changes. This is why this is not a memory leak.
+      // TODO: Consider cleaning this up with a shared_ptr.
+      class_defs[i].release();
+      class_defs[i].reset(new_class_def_order[i]);
     }
   }
-  return new_class_data_order;
 }
 
-int32_t DexLayout::LayoutStringData(const DexFile* dex_file) {
+void DexLayout::LayoutStringData(const DexFile* dex_file) {
   const size_t num_strings = header_->GetCollections().StringIds().size();
   std::vector<bool> is_shorty(num_strings, false);
   std::vector<bool> from_hot_method(num_strings, false);
@@ -1672,23 +1679,9 @@ int32_t DexLayout::LayoutStringData(const DexFile* dex_file) {
   }
   // Sort string data by specified order.
   std::vector<dex_ir::StringId*> string_ids;
-  size_t min_offset = std::numeric_limits<size_t>::max();
-  size_t max_offset = 0;
-  size_t hot_bytes = 0;
   for (auto& string_id : header_->GetCollections().StringIds()) {
     string_ids.push_back(string_id.get());
-    const size_t cur_offset = string_id->DataItem()->GetOffset();
-    CHECK_NE(cur_offset, 0u);
-    min_offset = std::min(min_offset, cur_offset);
-    dex_ir::StringData* data = string_id->DataItem();
-    const size_t element_size = data->GetSize() + 1;  // Add one extra for null.
-    size_t end_offset = cur_offset + element_size;
-    if (is_shorty[string_id->GetIndex()] || from_hot_method[string_id->GetIndex()]) {
-      hot_bytes += element_size;
-    }
-    max_offset = std::max(max_offset, end_offset);
-  }
-  VLOG(compiler) << "Hot string data bytes " << hot_bytes << "/" << max_offset - min_offset;
+  }
   std::sort(string_ids.begin(),
             string_ids.end(),
             [&is_shorty, &from_hot_method](const dex_ir::StringId* a,
@@ -1704,59 +1697,41 @@ int32_t DexLayout::LayoutStringData(const DexFile* dex_file) {
     if (a_is_shorty != b_is_shorty) {
       return a_is_shorty < b_is_shorty;
     }
-    // Preserve order.
-    return a->DataItem()->GetOffset() < b->DataItem()->GetOffset();
+    // Order by index by default.
+    return a->GetIndex() < b->GetIndex();
   });
-  // Now we know what order we want the string data, reorder the offsets.
-  size_t offset = min_offset;
+  dex_ir::CollectionVector<dex_ir::StringData>::Vector& string_datas =
+      header_->GetCollections().StringDatas();
+  // Now we know what order we want the string data, reorder them.
+  size_t data_index = 0;
   for (dex_ir::StringId* string_id : string_ids) {
-    dex_ir::StringData* data = string_id->DataItem();
-    data->SetOffset(offset);
-    offset += data->GetSize() + 1;  // Add one extra for null.
+    string_datas[data_index].release();
+    string_datas[data_index].reset(string_id->DataItem());
+    ++data_index;
   }
-  if (offset > max_offset) {
-    return offset - max_offset;
-    // If we expanded the string data section, we need to update the offsets or else we will
-    // corrupt the next section when writing out.
+  if (kIsDebugBuild) {
+    std::unordered_set<dex_ir::StringData*> visited;
+    for (const std::unique_ptr<dex_ir::StringData>& data : string_datas) {
+      visited.insert(data.get());
+    }
+    for (auto& string_id : header_->GetCollections().StringIds()) {
+      CHECK(visited.find(string_id->DataItem()) != visited.end());
+    }
   }
-  return 0;
+  CHECK_EQ(data_index, string_datas.size());
 }
 
 // Orders code items according to specified class data ordering.
-// NOTE: If the section following the code items is byte aligned, the last code item is left in
-// place to preserve alignment. Layout needs an overhaul to handle movement of other sections.
-int32_t DexLayout::LayoutCodeItems(const DexFile* dex_file,
-                                   std::vector<dex_ir::ClassData*> new_class_data_order) {
-  // Do not move code items if class data section precedes code item section.
-  // ULEB encoding is variable length, causing problems determining the offset of the code items.
-  // TODO: We should swap the order of these sections in the future to avoid this issue.
-  uint32_t class_data_offset = header_->GetCollections().ClassDatasOffset();
-  uint32_t code_item_offset = header_->GetCollections().CodeItemsOffset();
-  if (class_data_offset < code_item_offset) {
-    return 0;
-  }
-
-  // Find the last code item so we can leave it in place if the next section is not 4 byte aligned.
-  dex_ir::CodeItem* last_code_item = nullptr;
-  std::unordered_set<dex_ir::CodeItem*> visited_code_items;
-  bool is_code_item_aligned = IsNextSectionCodeItemAligned(code_item_offset);
-  if (!is_code_item_aligned) {
-    for (auto& code_item_pair : header_->GetCollections().CodeItems()) {
-      std::unique_ptr<dex_ir::CodeItem>& code_item = code_item_pair.second;
-      if (last_code_item == nullptr
-          || last_code_item->GetOffset() < code_item->GetOffset()) {
-        last_code_item = code_item.get();
-      }
-    }
-  }
-
+void DexLayout::LayoutCodeItems(const DexFile* dex_file) {
   static constexpr InvokeType invoke_types[] = {
     kDirect,
     kVirtual
   };
 
-  const size_t num_layout_types = static_cast<size_t>(LayoutType::kLayoutTypeCount);
-  std::unordered_set<dex_ir::CodeItem*> code_items[num_layout_types];
+  std::unordered_map<dex_ir::CodeItem*, LayoutType>& code_item_layout =
+      layout_hotness_info_.code_item_layout_;
+
+  // Assign hotness flags to all code items.
   for (InvokeType invoke_type : invoke_types) {
     for (std::unique_ptr<dex_ir::ClassDef>& class_def : header_->GetCollections().ClassDefs()) {
       const bool is_profile_class =
@@ -1772,7 +1747,7 @@ int32_t DexLayout::LayoutCodeItems(const DexFile* dex_file,
                                 : class_data->VirtualMethods())) {
         const dex_ir::MethodId *method_id = method->GetMethodId();
         dex_ir::CodeItem *code_item = method->GetCodeItem();
-        if (code_item == last_code_item || code_item == nullptr) {
+        if (code_item == nullptr) {
           continue;
         }
         // Separate executed methods (clinits and profiled methods) from unexecuted methods.
@@ -1794,194 +1769,61 @@ int32_t DexLayout::LayoutCodeItems(const DexFile* dex_file,
         } else if (hotness.IsInProfile()) {
           state = LayoutType::kLayoutTypeSometimesUsed;
         }
-        code_items[static_cast<size_t>(state)].insert(code_item);
-      }
-    }
-  }
-
-  // Removing duplicate CodeItems may expose other issues with downstream
-  // optimizations such as quickening.  But we need to ensure at least the weak
-  // forms of it currently in use do not break layout optimizations.
-  std::map<dex_ir::CodeItem*, uint32_t> original_code_item_offset;
-  // Total_diff includes diffs generated by clinits, executed, and non-executed methods.
-  int32_t total_diff = 0;
-  // The relative placement has no effect on correctness; it is used to ensure
-  // the layout is deterministic
-  for (size_t index = 0; index < num_layout_types; ++index) {
-    const std::unordered_set<dex_ir::CodeItem*>& code_items_set = code_items[index];
-    // diff is reset for each class of code items.
-    int32_t diff = 0;
-    const uint32_t start_offset = code_item_offset;
-    for (dex_ir::ClassData* data : new_class_data_order) {
-      data->SetOffset(data->GetOffset() + diff);
-      for (InvokeType invoke_type : invoke_types) {
-        for (auto &method : *(invoke_type == InvokeType::kDirect
-                                  ? data->DirectMethods()
-                                  : data->VirtualMethods())) {
-          dex_ir::CodeItem* code_item = method->GetCodeItem();
-          if (code_item != nullptr &&
-              code_items_set.find(code_item) != code_items_set.end()) {
-            // Compute where the CodeItem was originally laid out.
-            uint32_t original_offset = code_item->GetOffset();
-            auto it = original_code_item_offset.find(code_item);
-            if (it != original_code_item_offset.end()) {
-              original_offset = it->second;
-            } else {
-              original_code_item_offset[code_item] = code_item->GetOffset();
-              // Assign the new offset and move the pointer to allocate space.
-              code_item->SetOffset(code_item_offset);
-              code_item_offset +=
-                  RoundUp(code_item->GetSize(), kDexCodeItemAlignment);
-            }
-            // Update the size of the encoded methods to reflect that the offset difference
-            // may have changed the ULEB128 length.
-            diff +=
-                UnsignedLeb128Size(code_item->GetOffset()) - UnsignedLeb128Size(original_offset);
-          }
+        auto it = code_item_layout.emplace(code_item, state);
+        if (!it.second) {
+          LayoutType& layout_type = it.first->second;
+          // Already exists, merge the hotness.
+          layout_type = MergeLayoutType(layout_type, state);
         }
       }
     }
-    DexLayoutSection& code_section = dex_sections_.sections_[static_cast<size_t>(
-        DexLayoutSections::SectionType::kSectionTypeCode)];
-    code_section.parts_[index].offset_ = start_offset;
-    code_section.parts_[index].size_ = code_item_offset - start_offset;
-    for (size_t i = 0; i < num_layout_types; ++i) {
-      VLOG(dex) << "Code item layout bucket " << i << " count=" << code_items[i].size()
-                << " bytes=" << code_section.parts_[i].size_;
-    }
-    total_diff += diff;
   }
-  // Adjust diff to be 4-byte aligned.
-  return RoundUp(total_diff, kDexCodeItemAlignment);
-}
 
-bool DexLayout::IsNextSectionCodeItemAligned(uint32_t offset) {
-  dex_ir::Collections& collections = header_->GetCollections();
-  std::set<uint32_t> section_offsets;
-  section_offsets.insert(collections.MapListOffset());
-  section_offsets.insert(collections.TypeListsOffset());
-  section_offsets.insert(collections.AnnotationSetRefListsOffset());
-  section_offsets.insert(collections.AnnotationSetItemsOffset());
-  section_offsets.insert(collections.ClassDatasOffset());
-  section_offsets.insert(collections.CodeItemsOffset());
-  section_offsets.insert(collections.StringDatasOffset());
-  section_offsets.insert(collections.DebugInfoItemsOffset());
-  section_offsets.insert(collections.AnnotationItemsOffset());
-  section_offsets.insert(collections.EncodedArrayItemsOffset());
-  section_offsets.insert(collections.AnnotationsDirectoryItemsOffset());
-
-  auto found = section_offsets.find(offset);
-  if (found != section_offsets.end()) {
-    found++;
-    if (found != section_offsets.end()) {
-      return *found % kDexCodeItemAlignment == 0;
+  dex_ir::CollectionVector<dex_ir::CodeItem>::Vector& code_items =
+        header_->GetCollections().CodeItems();
+  if (VLOG_IS_ON(dex)) {
+    size_t layout_count[static_cast<size_t>(LayoutType::kLayoutTypeCount)] = {};
+    for (const std::unique_ptr<dex_ir::CodeItem>& code_item : code_items) {
+      auto it = code_item_layout.find(code_item.get());
+      DCHECK(it != code_item_layout.end());
+      ++layout_count[static_cast<size_t>(it->second)];
+    }
+    for (size_t i = 0; i < static_cast<size_t>(LayoutType::kLayoutTypeCount); ++i) {
+      LOG(INFO) << "Code items in category " << i << " count=" << layout_count[i];
     }
-  }
-  return false;
-}
-
-// Adjust offsets of every item in the specified section by diff bytes.
-template<class T> void DexLayout::FixupSection(std::map<uint32_t, std::unique_ptr<T>>& map,
-                                               uint32_t diff) {
-  for (auto& pair : map) {
-    std::unique_ptr<T>& item = pair.second;
-    item->SetOffset(item->GetOffset() + diff);
-  }
-}
-
-// Adjust offsets of all sections with an address after the specified offset by diff bytes.
-void DexLayout::FixupSections(uint32_t offset, uint32_t diff) {
-  dex_ir::Collections& collections = header_->GetCollections();
-  uint32_t map_list_offset = collections.MapListOffset();
-  if (map_list_offset > offset) {
-    collections.SetMapListOffset(map_list_offset + diff);
-  }
-
-  uint32_t type_lists_offset = collections.TypeListsOffset();
-  if (type_lists_offset > offset) {
-    collections.SetTypeListsOffset(type_lists_offset + diff);
-    FixupSection(collections.TypeLists(), diff);
-  }
-
-  uint32_t annotation_set_ref_lists_offset = collections.AnnotationSetRefListsOffset();
-  if (annotation_set_ref_lists_offset > offset) {
-    collections.SetAnnotationSetRefListsOffset(annotation_set_ref_lists_offset + diff);
-    FixupSection(collections.AnnotationSetRefLists(), diff);
-  }
-
-  uint32_t annotation_set_items_offset = collections.AnnotationSetItemsOffset();
-  if (annotation_set_items_offset > offset) {
-    collections.SetAnnotationSetItemsOffset(annotation_set_items_offset + diff);
-    FixupSection(collections.AnnotationSetItems(), diff);
-  }
-
-  uint32_t class_datas_offset = collections.ClassDatasOffset();
-  if (class_datas_offset > offset) {
-    collections.SetClassDatasOffset(class_datas_offset + diff);
-    FixupSection(collections.ClassDatas(), diff);
-  }
-
-  uint32_t code_items_offset = collections.CodeItemsOffset();
-  if (code_items_offset > offset) {
-    collections.SetCodeItemsOffset(code_items_offset + diff);
-    FixupSection(collections.CodeItems(), diff);
-  }
-
-  uint32_t string_datas_offset = collections.StringDatasOffset();
-  if (string_datas_offset > offset) {
-    collections.SetStringDatasOffset(string_datas_offset + diff);
-    FixupSection(collections.StringDatas(), diff);
-  }
-
-  uint32_t debug_info_items_offset = collections.DebugInfoItemsOffset();
-  if (debug_info_items_offset > offset) {
-    collections.SetDebugInfoItemsOffset(debug_info_items_offset + diff);
-    FixupSection(collections.DebugInfoItems(), diff);
-  }
-
-  uint32_t annotation_items_offset = collections.AnnotationItemsOffset();
-  if (annotation_items_offset > offset) {
-    collections.SetAnnotationItemsOffset(annotation_items_offset + diff);
-    FixupSection(collections.AnnotationItems(), diff);
-  }
-
-  uint32_t encoded_array_items_offset = collections.EncodedArrayItemsOffset();
-  if (encoded_array_items_offset > offset) {
-    collections.SetEncodedArrayItemsOffset(encoded_array_items_offset + diff);
-    FixupSection(collections.EncodedArrayItems(), diff);
   }
 
-  uint32_t annotations_directory_items_offset = collections.AnnotationsDirectoryItemsOffset();
-  if (annotations_directory_items_offset > offset) {
-    collections.SetAnnotationsDirectoryItemsOffset(annotations_directory_items_offset + diff);
-    FixupSection(collections.AnnotationsDirectoryItems(), diff);
-  }
+  // Sort the code items vector by new layout. The writing process will take care of calculating
+  // all the offsets. Stable sort to preserve any existing locality that might be there.
+  std::stable_sort(code_items.begin(),
+                   code_items.end(),
+                   [&](const std::unique_ptr<dex_ir::CodeItem>& a,
+                       const std::unique_ptr<dex_ir::CodeItem>& b) {
+    auto it_a = code_item_layout.find(a.get());
+    auto it_b = code_item_layout.find(b.get());
+    DCHECK(it_a != code_item_layout.end());
+    DCHECK(it_b != code_item_layout.end());
+    const LayoutType layout_type_a = it_a->second;
+    const LayoutType layout_type_b = it_b->second;
+    return layout_type_a < layout_type_b;
+  });
 }
 
 void DexLayout::LayoutOutputFile(const DexFile* dex_file) {
-  const int32_t string_diff = LayoutStringData(dex_file);
-  // If we expanded the string data section, we need to update the offsets or else we will
-  // corrupt the next section when writing out.
-  FixupSections(header_->GetCollections().StringDatasOffset(), string_diff);
-  // Update file size.
-  header_->SetFileSize(header_->FileSize() + string_diff);
-
-  std::vector<dex_ir::ClassData*> new_class_data_order = LayoutClassDefsAndClassData(dex_file);
-  const int32_t code_item_diff = LayoutCodeItems(dex_file, new_class_data_order);
-  // Move sections after ClassData by diff bytes.
-  FixupSections(header_->GetCollections().ClassDatasOffset(), code_item_diff);
-
-  // Update file and data size.
-  // The data size must be aligned to kDataSectionAlignment.
-  const int32_t total_diff = code_item_diff + string_diff;
-  header_->SetDataSize(RoundUp(header_->DataSize() + total_diff, kDataSectionAlignment));
-  header_->SetFileSize(header_->FileSize() + total_diff);
+  LayoutStringData(dex_file);
+  LayoutClassDefsAndClassData(dex_file);
+  LayoutCodeItems(dex_file);
 }
 
-void DexLayout::OutputDexFile(const DexFile* dex_file) {
+void DexLayout::OutputDexFile(const DexFile* dex_file, bool compute_offsets) {
   const std::string& dex_file_location = dex_file->GetLocation();
   std::string error_msg;
   std::unique_ptr<File> new_file;
+  // Since we allow dex growth, we need to size the map larger than the original input to be safe.
+  // Reserve an extra 10% to add some buffer room. Note that this is probably more than
+  // necessary.
+  constexpr size_t kReserveFraction = 10;
+  const size_t max_size = header_->FileSize() + header_->FileSize() / kReserveFraction;
   if (!options_.output_to_memmap_) {
     std::string output_location(options_.output_dex_directory_);
     size_t last_slash = dex_file_location.rfind('/');
@@ -1998,15 +1840,15 @@ void DexLayout::OutputDexFile(const DexFile* dex_file) {
       LOG(ERROR) << "Could not create dex writer output file: " << output_location;
       return;
     }
-    if (ftruncate(new_file->Fd(), header_->FileSize()) != 0) {
+    if (ftruncate(new_file->Fd(), max_size) != 0) {
       LOG(ERROR) << "Could not grow dex writer output file: " << output_location;;
       new_file->Erase();
       return;
     }
-    mem_map_.reset(MemMap::MapFile(header_->FileSize(), PROT_READ | PROT_WRITE, MAP_SHARED,
+    mem_map_.reset(MemMap::MapFile(max_size, PROT_READ | PROT_WRITE, MAP_SHARED,
         new_file->Fd(), 0, /*low_4gb*/ false, output_location.c_str(), &error_msg));
   } else {
-    mem_map_.reset(MemMap::MapAnonymous("layout dex", nullptr, header_->FileSize(),
+    mem_map_.reset(MemMap::MapAnonymous("layout dex", nullptr, max_size,
         PROT_READ | PROT_WRITE, /* low_4gb */ false, /* reuse */ false, &error_msg));
   }
   if (mem_map_ == nullptr) {
@@ -2016,8 +1858,14 @@ void DexLayout::OutputDexFile(const DexFile* dex_file) {
     }
     return;
   }
-  DexWriter::Output(header_, mem_map_.get(), options_.compact_dex_level_);
+  DexWriter::Output(header_, mem_map_.get(), this, compute_offsets, options_.compact_dex_level_);
   if (new_file != nullptr) {
+    // Since we make the memmap larger than needed, shrink the file back down to not leave extra
+    // padding.
+    int res = new_file->SetLength(header_->FileSize());
+    if (res != 0) {
+      LOG(ERROR) << "Truncating file resulted in " << res;
+    }
     UNUSED(new_file->FlushCloseOrErase());
   }
 }
@@ -2028,7 +1876,15 @@ void DexLayout::OutputDexFile(const DexFile* dex_file) {
 void DexLayout::ProcessDexFile(const char* file_name,
                                const DexFile* dex_file,
                                size_t dex_file_index) {
-  std::unique_ptr<dex_ir::Header> header(dex_ir::DexIrBuilder(*dex_file));
+  const bool output = options_.output_dex_directory_ != nullptr || options_.output_to_memmap_;
+  // Try to avoid eagerly assigning offsets to find bugs since GetOffset will abort if the offset
+  // is unassigned.
+  bool eagerly_assign_offsets = false;
+  if (options_.visualize_pattern_ || options_.show_section_statistics_ || options_.dump_) {
+    // These options required the offsets for dumping purposes.
+    eagerly_assign_offsets = true;
+  }
+  std::unique_ptr<dex_ir::Header> header(dex_ir::DexIrBuilder(*dex_file, eagerly_assign_offsets));
   SetHeader(header.get());
 
   if (options_.verbose_) {
@@ -2052,13 +1908,17 @@ void DexLayout::ProcessDexFile(const char* file_name,
   }
 
   // In case we are outputting to a file, keep it open so we can verify.
-  if (options_.output_dex_directory_ != nullptr || options_.output_to_memmap_) {
-    if (info_ != nullptr) {
+  if (output) {
+    // Layout information about what strings and code items are hot. Used by the writing process
+    // to generate the sections that are stored in the oat file.
+    bool do_layout = info_ != nullptr;
+    if (do_layout) {
       LayoutOutputFile(dex_file);
     }
-    OutputDexFile(dex_file);
+    OutputDexFile(dex_file, do_layout);
 
     // Clear header before verifying to reduce peak RAM usage.
+    const size_t file_size = header_->FileSize();
     header.reset();
 
     // Verify the output dex file's structure, only enabled by default for debug builds.
@@ -2066,7 +1926,7 @@ void DexLayout::ProcessDexFile(const char* file_name,
       std::string error_msg;
       std::string location = "memory mapped file for " + std::string(file_name);
       std::unique_ptr<const DexFile> output_dex_file(DexFileLoader::Open(mem_map_->Begin(),
-                                                                         mem_map_->Size(),
+                                                                         file_size,
                                                                          location,
                                                                          /* checksum */ 0,
                                                                          /*oat_dex_file*/ nullptr,
@@ -2076,11 +1936,16 @@ void DexLayout::ProcessDexFile(const char* file_name,
       CHECK(output_dex_file != nullptr) << "Failed to re-open output file:" << error_msg;
 
       // Do IR-level comparison between input and output. This check ignores potential differences
-      // due to layout, so offsets are not checked. Instead, it checks the data contents of each item.
+      // due to layout, so offsets are not checked. Instead, it checks the data contents of each
+      // item.
       //
       // Regenerate output IR to catch any bugs that might happen during writing.
-      std::unique_ptr<dex_ir::Header> output_header(dex_ir::DexIrBuilder(*output_dex_file));
-      std::unique_ptr<dex_ir::Header> orig_header(dex_ir::DexIrBuilder(*dex_file));
+      std::unique_ptr<dex_ir::Header> output_header(
+          dex_ir::DexIrBuilder(*output_dex_file,
+                               /*eagerly_assign_offsets*/ true));
+      std::unique_ptr<dex_ir::Header> orig_header(
+          dex_ir::DexIrBuilder(*dex_file,
+                               /*eagerly_assign_offsets*/ true));
       CHECK(VerifyOutputDexFile(output_header.get(), orig_header.get(), &error_msg)) << error_msg;
     }
   }
diff --git a/dexlayout/dexlayout.h b/dexlayout/dexlayout.h
index 2e897739cc..8a277b7afe 100644
--- a/dexlayout/dexlayout.h
+++ b/dexlayout/dexlayout.h
@@ -25,6 +25,7 @@
 
 #include <stdint.h>
 #include <stdio.h>
+#include <unordered_map>
 
 #include "cdex/compact_dex_level.h"
 #include "dex_file_layout.h"
@@ -69,6 +70,13 @@ class Options {
   const char* profile_file_name_ = nullptr;
 };
 
+// Hotness info
+class DexLayoutHotnessInfo {
+ public:
+  // Store layout information so that the offset calculation can specify the section sizes.
+  std::unordered_map<dex_ir::CodeItem*, LayoutType> code_item_layout_;
+};
+
 class DexLayout {
  public:
   DexLayout(Options& options,
@@ -86,10 +94,14 @@ class DexLayout {
 
   MemMap* GetAndReleaseMemMap() { return mem_map_.release(); }
 
-  const DexLayoutSections& GetSections() const {
+  DexLayoutSections& GetSections() {
     return dex_sections_;
   }
 
+  const DexLayoutHotnessInfo& LayoutHotnessInfo() const {
+    return layout_hotness_info_;
+  }
+
  private:
   void DumpAnnotationSetItem(dex_ir::AnnotationSetItem* set_item);
   void DumpBytecodes(uint32_t idx, const dex_ir::CodeItem* code, uint32_t code_offset);
@@ -120,18 +132,14 @@ class DexLayout {
   void DumpSField(uint32_t idx, uint32_t flags, int i, dex_ir::EncodedValue* init);
   void DumpDexFile();
 
-  std::vector<dex_ir::ClassData*> LayoutClassDefsAndClassData(const DexFile* dex_file);
-  int32_t LayoutCodeItems(const DexFile* dex_file,
-                          std::vector<dex_ir::ClassData*> new_class_data_order);
-  int32_t LayoutStringData(const DexFile* dex_file);
-  bool IsNextSectionCodeItemAligned(uint32_t offset);
-  template<class T> void FixupSection(std::map<uint32_t, std::unique_ptr<T>>& map, uint32_t diff);
-  void FixupSections(uint32_t offset, uint32_t diff);
+  void LayoutClassDefsAndClassData(const DexFile* dex_file);
+  void LayoutCodeItems(const DexFile* dex_file);
+  void LayoutStringData(const DexFile* dex_file);
 
   // Creates a new layout for the dex file based on profile info.
   // Currently reorders ClassDefs, ClassDataItems, and CodeItems.
   void LayoutOutputFile(const DexFile* dex_file);
-  void OutputDexFile(const DexFile* dex_file);
+  void OutputDexFile(const DexFile* dex_file, bool compute_offsets);
 
   void DumpCFG(const DexFile* dex_file, int idx);
   void DumpCFG(const DexFile* dex_file, uint32_t dex_method_idx, const DexFile::CodeItem* code);
@@ -142,6 +150,8 @@ class DexLayout {
   dex_ir::Header* header_;
   std::unique_ptr<MemMap> mem_map_;
   DexLayoutSections dex_sections_;
+  // Layout hotness information is only calculated when dexlayout is enabled.
+  DexLayoutHotnessInfo layout_hotness_info_;
 
   DISALLOW_COPY_AND_ASSIGN(DexLayout);
 };
diff --git a/runtime/dex_file_layout.cc b/runtime/dex_file_layout.cc
index c3fae15b14..1973440d55 100644
--- a/runtime/dex_file_layout.cc
+++ b/runtime/dex_file_layout.cc
@@ -26,10 +26,10 @@ namespace art {
 
 void DexLayoutSection::Subsection::Madvise(const DexFile* dex_file, int advice) const {
   DCHECK(dex_file != nullptr);
-  DCHECK_LE(size_, dex_file->Size());
-  DCHECK_LE(offset_ + size_, dex_file->Size());
-  MadviseLargestPageAlignedRegion(dex_file->Begin() + offset_,
-                                  dex_file->Begin() + offset_ + size_,
+  DCHECK_LT(start_offset_, dex_file->Size());
+  DCHECK_LE(end_offset_, dex_file->Size());
+  MadviseLargestPageAlignedRegion(dex_file->Begin() + start_offset_,
+                                  dex_file->Begin() + end_offset_,
                                   advice);
 }
 
@@ -69,7 +69,7 @@ std::ostream& operator<<(std::ostream& os, const DexLayoutSection& section) {
   for (size_t i = 0; i < static_cast<size_t>(LayoutType::kLayoutTypeCount); ++i) {
     const DexLayoutSection::Subsection& part = section.parts_[i];
     os << static_cast<LayoutType>(i) << "("
-       << part.offset_ << "-" << part.offset_ + part.size_ << ") ";
+       << part.start_offset_ << "-" << part.end_offset_ << ") ";
   }
   return os;
 }
diff --git a/runtime/dex_file_layout.h b/runtime/dex_file_layout.h
index 40cc91232e..4c960c3ff5 100644
--- a/runtime/dex_file_layout.h
+++ b/runtime/dex_file_layout.h
@@ -17,22 +17,25 @@
 #ifndef ART_RUNTIME_DEX_FILE_LAYOUT_H_
 #define ART_RUNTIME_DEX_FILE_LAYOUT_H_
 
+#include <algorithm>
 #include <cstdint>
 #include <iosfwd>
 
+#include "base/logging.h"
+
 namespace art {
 
 class DexFile;
 
 enum class LayoutType : uint8_t {
+  // Layout of things that are hot (commonly accessed), these should be pinned or madvised will
+  // need.
+  kLayoutTypeHot,
   // Layout of things that are randomly used. These should be advised to random access.
   // Without layout, this is the default mode when loading a dex file.
   kLayoutTypeSometimesUsed,
   // Layout of things that are only used during startup, these can be madvised after launch.
   kLayoutTypeStartupOnly,
-  // Layout of things that are hot (commonly accessed), these should be pinned or madvised will
-  // need.
-  kLayoutTypeHot,
   // Layout of things that are needed probably only once (class initializers). These can be
   // madvised during trim events.
   kLayoutTypeUsedOnce,
@@ -44,6 +47,11 @@ enum class LayoutType : uint8_t {
 };
 std::ostream& operator<<(std::ostream& os, const LayoutType& collector_type);
 
+// Return the "best" layout option if the same item has multiple different layouts.
+static inline LayoutType MergeLayoutType(LayoutType a, LayoutType b) {
+  return std::min(a, b);
+}
+
 enum class MadviseState : uint8_t {
   // Madvise based on a file that was just loaded.
   kMadviseStateAtLoad,
@@ -55,15 +63,35 @@ enum class MadviseState : uint8_t {
 std::ostream& operator<<(std::ostream& os, const MadviseState& collector_type);
 
 // A dex layout section such as code items or strings. Each section is composed of subsections
-// that are layed out ajacently to each other such as (hot, unused, startup, etc...).
+// that are laid out adjacently to each other such as (hot, unused, startup, etc...).
 class DexLayoutSection {
  public:
   // A subsection is a a continuous range of dex file that is all part of the same layout hint.
   class Subsection {
    public:
     // Use uint32_t to handle 32/64 bit cross compilation.
-    uint32_t offset_ = 0u;
-    uint32_t size_ = 0u;
+    uint32_t start_offset_ = 0u;
+    uint32_t end_offset_ = 0u;
+
+    bool Contains(uint32_t offset) const {
+      return start_offset_ <= offset && offset < end_offset_;
+    }
+
+    bool Size() const {
+      DCHECK_LE(start_offset_, end_offset_);
+      return end_offset_ - start_offset_;
+    }
+
+    void CombineSection(uint32_t start_offset, uint32_t end_offset) {
+      DCHECK_LT(start_offset, end_offset);
+      if (start_offset_ == end_offset_) {
+        start_offset_ = start_offset;
+        end_offset_ = end_offset;
+      } else  {
+        start_offset_ = std::min(start_offset_, start_offset);
+        end_offset_ = std::max(end_offset_, end_offset);
+      }
+    }
 
     void Madvise(const DexFile* dex_file, int advice) const;
   };
diff --git a/runtime/dex_file_verifier.cc b/runtime/dex_file_verifier.cc
index 025952f7a9..edf5650df1 100644
--- a/runtime/dex_file_verifier.cc
+++ b/runtime/dex_file_verifier.cc
@@ -471,7 +471,9 @@ bool DexFileVerifier::CheckMap() {
     if (IsDataSectionType(item_type)) {
       uint32_t icount = item->size_;
       if (UNLIKELY(icount > data_items_left)) {
-        ErrorStringPrintf("Too many items in data section: %ud", data_item_count + icount);
+        ErrorStringPrintf("Too many items in data section: %ud item_type %zx",
+                          data_item_count + icount,
+                          static_cast<size_t>(item_type));
         return false;
       }
       data_items_left -= icount;