[Courgette] Using LabelManager to reduce Courgette-apply peak RAM by 25%.

courgette/disassembler_elf_32.cc

 // Copyright 2013 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "courgette/disassembler_elf_32.h"
 
 #include <algorithm>
+#include <iterator>
 #include <utility>
 
 #include "base/logging.h"
 #include "courgette/assembly_program.h"
 #include "courgette/courgette.h"
 
 namespace courgette {
 
 namespace {
 
 // Sorts |section_headers| by file offset and stores the resulting permutation
 // of section ids in |order|.
 std::vector<Elf32_Half> GetSectionHeaderFileOffsetOrder(
     const std::vector<Elf32_Shdr>& section_headers) {
   size_t size = section_headers.size();
   std::vector<Elf32_Half> order(size);
   for (size_t i = 0; i < size; ++i)
     order[i] = static_cast<Elf32_Half>(i);
 
   auto comp = [&](int idx1, int idx2) {
     return section_headers[idx1].sh_offset < section_headers[idx2].sh_offset;
   };
   std::stable_sort(order.begin(), order.end(), comp);
   return order;
 }
 
 }  // namespace
 
+DisassemblerElf32::Elf32RvaVisitor_Rel32::Elf32RvaVisitor_Rel32(
+    const std::vector<std::unique_ptr<TypedRVA>>& rva_locations)
+    : VectorRvaVisitor<std::unique_ptr<TypedRVA>>(rva_locations) {
+}
+
+RVA DisassemblerElf32::Elf32RvaVisitor_Rel32::Get() const {
+  return (*it_)->rva() + (*it_)->relative_target();
+}
+
 DisassemblerElf32::DisassemblerElf32(const void* start, size_t length)
     : Disassembler(start, length),
       header_(nullptr),
       section_header_table_size_(0),
       program_header_table_(nullptr),
       program_header_table_size_(0),
       default_string_section_(nullptr) {
 }
 
 RVA DisassemblerElf32::FileOffsetToRVA(FileOffset offset) const {
@@ skipping 111 matching lines @@
 bool DisassemblerElf32::Disassemble(AssemblyProgram* target) {
   if (!ok())
     return false;
 
   // The Image Base is always 0 for ELF Executables
   target->set_image_base(0);
 
   if (!ParseAbs32Relocs())
     return false;
 
-  if (!ParseRel32RelocsFromSections())
+  if (!ParseRel32RelocsFromSections())  // Does not sort rel32 locations.
     return false;
 
+  PrecomputeLabels(target);
+  RemoveUnusedRel32Locations(target);
+
   if (!ParseFile(target))
     return false;
 
+  // Finally sort rel32 locations.
   std::sort(rel32_locations_.begin(),
             rel32_locations_.end(),
             TypedRVA::IsLessThanByRVA);
   DCHECK(rel32_locations_.empty() ||
          rel32_locations_.back()->rva() != kUnassignedRVA);
 
   target->DefaultAssignIndexes();
   return true;
 }
 
+CheckBool DisassemblerElf32::IsValidTargetRVA(RVA rva) const {
+  if (rva == kUnassignedRVA)
+    return false;
+
+  // |rva| is valid if it's contained in any program segment.
+  for (Elf32_Half segment_id = 0; segment_id < ProgramSegmentHeaderCount();
+       ++segment_id) {
+    const Elf32_Phdr* segment_header = ProgramSegmentHeader(segment_id);
+
+    if (segment_header->p_type != PT_LOAD)
+      continue;
+
+    Elf32_Addr begin = segment_header->p_vaddr;
+    Elf32_Addr end = segment_header->p_vaddr + segment_header->p_memsz;
+
+    if (rva >= begin && rva < end)
+      return true;
+  }
+
+  return false;
+}
+
 bool DisassemblerElf32::UpdateLength() {
   Elf32_Off result = 0;
 
   // Find the end of the last section
   for (Elf32_Half section_id = 0; section_id < SectionHeaderCount();
        ++section_id) {
     const Elf32_Shdr* section_header = SectionHeader(section_id);
 
     if (section_header->sh_type == SHT_NOBITS)
       continue;
@@ skipping 38 matching lines @@
     name->clear();
   } else {
     if (string_pos >= default_string_section_size_)
       return false;
     // Safe because string section must terminate with null.
     *name = default_string_section_ + string_pos;
   }
   return true;
 }
 
-CheckBool DisassemblerElf32::IsValidTargetRVA(RVA rva) const {
-  if (rva == kUnassignedRVA)
-    return false;
-
-  // It's valid if it's contained in any program segment
-  for (Elf32_Half segment_id = 0; segment_id < ProgramSegmentHeaderCount();
-       ++segment_id) {
-    const Elf32_Phdr* segment_header = ProgramSegmentHeader(segment_id);
-
-    if (segment_header->p_type != PT_LOAD)
-      continue;
-
-    Elf32_Addr begin = segment_header->p_vaddr;
-    Elf32_Addr end = segment_header->p_vaddr + segment_header->p_memsz;
-
-    if (rva >= begin && rva < end)
-      return true;
-  }
-
-  return false;
-}
-
 CheckBool DisassemblerElf32::RVAsToFileOffsets(
     const std::vector<RVA>& rvas,
     std::vector<FileOffset>* file_offsets) {
   file_offsets->clear();
+  file_offsets->reserve(rvas.size());
   for (RVA rva : rvas) {
     FileOffset file_offset = RVAToFileOffset(rva);
     if (file_offset == kNoFileOffset)
       return false;
     file_offsets->push_back(file_offset);
   }
   return true;
 }
 
 CheckBool DisassemblerElf32::RVAsToFileOffsets(
     std::vector<std::unique_ptr<TypedRVA>>* typed_rvas) {
   for (auto& typed_rva : *typed_rvas) {
     FileOffset file_offset = RVAToFileOffset(typed_rva->rva());
     if (file_offset == kNoFileOffset)
       return false;
     typed_rva->set_file_offset(file_offset);
   }
   return true;
 }
 
+RvaVisitor* DisassemblerElf32::CreateAbs32TargetRvaVisitor() {
+  return new RvaVisitor_Abs32(abs32_locations_, *this);
+}
+
+RvaVisitor* DisassemblerElf32::CreateRel32TargetRvaVisitor() {
+  return new Elf32RvaVisitor_Rel32(rel32_locations_);
+}
+
+void DisassemblerElf32::RemoveUnusedRel32Locations(AssemblyProgram* program) {
+  auto tail_it = rel32_locations_.begin();
+  for (auto head_it = rel32_locations_.begin();
+       head_it != rel32_locations_.end(); ++head_it) {
+    RVA target_rva = (*head_it)->rva() + (*head_it)->relative_target();
+    if (program->FindRel32Label(target_rva) == nullptr) {
+      // If address does not match a Label (because it was removed), deallocate.
+      (*head_it).reset(nullptr);
+    } else {
+      // Else squeeze nullptr to end to compactify.
+      if (tail_it != head_it)
+        (*tail_it).swap(*head_it);
+      ++tail_it;
+    }
+  }
+  rel32_locations_.resize(std::distance(rel32_locations_.begin(), tail_it));
+}
+
 CheckBool DisassemblerElf32::ParseFile(AssemblyProgram* program) {
   // Walk all the bytes in the file, whether or not in a section.
   FileOffset file_offset = 0;
 
   std::vector<FileOffset> abs_offsets;
 
   // File parsing follows file offset order, and we visit abs32 and rel32
   // locations in lockstep. Therefore we need to extract and sort file offsets
   // of all abs32 and rel32 locations.
   if (!RVAsToFileOffsets(abs32_locations_, &abs_offsets))
@@ skipping 121 matching lines @@
 
       file_offset = next_relocation;
       continue;
     }
 
     if (*current_abs_offset != end_abs_offset &&
         file_offset == **current_abs_offset) {
       RVA target_rva = PointerToTargetRVA(FileOffsetToPointer(file_offset));
       DCHECK_NE(kNoRVA, target_rva);
 
-      if (!program->EmitAbs32(program->FindOrMakeAbs32Label(target_rva)))
+      Label* label = program->FindAbs32Label(target_rva);
+      CHECK(label);
+      if (!program->EmitAbs32(label))
         return false;
       file_offset += sizeof(RVA);
       ++(*current_abs_offset);
       continue;
     }
 
     if (*current_rel != end_rel &&
         file_offset == (**current_rel)->file_offset()) {
       uint32_t relative_target = (**current_rel)->relative_target();
+      CHECK_EQ(RVA(origin + (file_offset - origin_offset)),
+               (**current_rel)->rva());
       // This cast is for 64 bit systems, and is only safe because we
       // are working on 32 bit executables.
       RVA target_rva = (RVA)(origin + (file_offset - origin_offset) +
                              relative_target);
 
-      if (!(**current_rel)->EmitInstruction(program, target_rva))
+      Label* label = program->FindRel32Label(target_rva);
+      CHECK(label);
+
+      if (!(**current_rel)->EmitInstruction(program, label))
         return false;
       file_offset += (**current_rel)->op_size();
       ++(*current_rel);
       continue;
     }
   }
 
   // Rest of the section (if any)
   return ParseSimpleRegion(file_offset, section_end, program);
 }
@@ skipping 97 matching lines @@
     if (!ParseRel32RelocsFromSection(section_header))
       return false;
   }
   if (!found_rel32)
     VLOG(1) << "Warning: Found no rel32 addresses. Missing .text section?";
 
   return true;
 }
 
 }  // namespace courgette