[BOLT] Refactoring of section handling code

Summary:
This is a big refactoring of the section handling code.  I've removed the SectionInfoMap and NoteSectionInfo and stored all the associated info about sections in BinaryContext and BinarySection classes.  BinarySections should now hold all the info we care about for each section.  They can be initialized from SectionRefs but don't necessarily require one to be created.  There are only one or two spots that needed access to the original SectionRef to work properly.

The trickiest part was making sure RewriteInstance.cpp iterated over the proper sets of sections for each of it's different types of processing.  The different sets are broken down roughly as allocatable and non-alloctable and "registered" (I couldn't think up a better name).  "Registered" means that the section has been updated to include output information, i.e. contents, file offset/address, new size, etc.  It may help to have special iterators on BinaryContext to iterate over the different classes to make things easier.  I can do that if you guys think it is worthwhile.

I found pointee_iterator in the llvm ADT code.  Use that for iterating over BBs in BinaryFunction rather than the custom iterator class.

(cherry picked from FBD6879086)

bolt/BinaryContext.cpp

@@ -23,6 +23,9 @@
 using namespace llvm;
 using namespace bolt;
 
+#undef  DEBUG_TYPE
+#define DEBUG_TYPE "bolt"
+
 namespace opts {
 
 extern cl::OptionCategory BoltCategory;
@@ -31,23 +34,30 @@ static cl::opt<bool>
 PrintDebugInfo("print-debug-info",
   cl::desc("print debug info when printing functions"),
   cl::Hidden,
+  cl::ZeroOrMore,
   cl::cat(BoltCategory));
 
-static cl::opt<bool>
+cl::opt<bool>
 PrintRelocations("print-relocations",
-  cl::desc("print relocations when printing functions"),
+  cl::desc("print relocations when printing functions/objects"),
   cl::Hidden,
+  cl::ZeroOrMore,
   cl::cat(BoltCategory));
 
 static cl::opt<bool>
 PrintMemData("print-mem-data",
   cl::desc("print memory data annotations when printing functions"),
   cl::Hidden,
+  cl::ZeroOrMore,
   cl::cat(BoltCategory));
 
 } // namespace opts
 
-BinaryContext::~BinaryContext() { }
+BinaryContext::~BinaryContext() {
+  for (auto *Section : Sections) {
+    delete Section;
+  }
+}
 
 std::unique_ptr<MCObjectWriter>
 BinaryContext::createObjectWriter(raw_pwrite_stream &OS) {
@@ -474,20 +484,82 @@ BinaryContext::getSectionForAddress(uint64_t Address) const {
   return std::make_error_code(std::errc::bad_address);
 }
 
-BinarySection &BinaryContext::registerSection(SectionRef Section) {
+BinarySection &BinaryContext::registerSection(BinarySection *Section) {
-  StringRef Name;
+  assert(!Section->getName().empty() &&
-  Section.getName(Name);
+         "can't register sections without a name");
-  auto Res = Sections.insert(BinarySection(Section));
+  auto Res = Sections.insert(Section);
   assert(Res.second && "can't register the same section twice.");
-  // Cast away const here because std::set always stores values by
-  // const.  It's ok to do this because we can never change the
-  // BinarySection properties that affect set ordering.
-  auto *BS = const_cast<BinarySection *>(&*Res.first);
   // Only register sections with addresses in the AddressToSection map.
-  if (Section.getAddress())
+  if (Section->getAddress())
-    AddressToSection.insert(std::make_pair(Section.getAddress(), BS));
+    AddressToSection.insert(std::make_pair(Section->getAddress(), Section));
-  NameToSection.insert(std::make_pair(Name, BS));
+  NameToSection.insert(std::make_pair(Section->getName(), Section));
-  return *BS;
+  DEBUG(dbgs() << "BOLT-DEBUG: registering " << *Section << "\n");
+  return *Section;
+}
+
+BinarySection &BinaryContext::registerSection(SectionRef Section) {
+  return registerSection(new BinarySection(Section));
+}
+
+BinarySection &BinaryContext::registerOrUpdateSection(StringRef Name,
+                                                      unsigned ELFType,
+                                                      unsigned ELFFlags,
+                                                      uint8_t *Data,
+                                                      uint64_t Size,
+                                                      unsigned Alignment,
+                                                      bool IsLocal) {
+  auto NamedSections = getSectionByName(Name);
+  if (NamedSections.begin() != NamedSections.end()) {
+    assert(std::next(NamedSections.begin()) == NamedSections.end() &&
+           "can only update unique sections");
+    auto *Section = NamedSections.begin()->second;
+
+    DEBUG(dbgs() << "BOLT-DEBUG: updating " << *Section << " -> ");
+    const auto Flag = Section->isAllocatable();
+    Section->update(Data, Size, Alignment, ELFType, ELFFlags, IsLocal);
+    DEBUG(dbgs() << *Section << "\n");
+    assert(Flag == Section->isAllocatable() &&
+           "can't change section allocation status");
+    return *Section;
+  }
+
+  return registerSection(new BinarySection(Name, Data, Size, Alignment,
+                                           ELFType, ELFFlags, IsLocal));
+}
+
+bool BinaryContext::deregisterSection(BinarySection &Section) {
+  auto *SectionPtr = &Section;
+  auto Itr = Sections.find(SectionPtr);
+  if (Itr != Sections.end()) {
+    auto Range = AddressToSection.equal_range(SectionPtr->getAddress());
+    while (Range.first != Range.second) {
+      if (Range.first->second == SectionPtr) {
+        AddressToSection.erase(Range.first);
+        break;
+      }
+      ++Range.first;
+    }
+
+    auto NameRange = NameToSection.equal_range(SectionPtr->getName());
+    while (NameRange.first != NameRange.second) {
+      if (NameRange.first->second == SectionPtr) {
+        NameToSection.erase(NameRange.first);
+        break;
+      }
+      ++NameRange.first;
+    }
+
+    Sections.erase(Itr);
+    delete SectionPtr;
+    return true;
+  }
+  return false;
+}
+
+void BinaryContext::printSections(raw_ostream &OS) const {
+  for (auto &Section : Sections) {
+    OS << "BOLT-INFO: " << *Section << "\n";
+  }
 }
 
 ErrorOr<uint64_t>
@@ -504,27 +576,24 @@ BinaryContext::extractPointerAtAddress(uint64_t Address) const {
   return DE.getAddress(&SectionOffset);
 }
 
-void BinaryContext::addSectionRelocation(BinarySection &Section,
-                                         uint64_t Offset,
-                                         MCSymbol *Symbol,
-                                         uint64_t Type,
-                                         uint64_t Addend) {
-  Section.addRelocation(Offset, Symbol, Type, Addend);
-}
-
 void BinaryContext::addRelocation(uint64_t Address,
                                   MCSymbol *Symbol,
                                   uint64_t Type,
-                                  uint64_t Addend) {
+                                  uint64_t Addend,
+                                  uint64_t Value) {
   auto Section = getSectionForAddress(Address);
   assert(Section && "cannot find section for address");
-  Section->addRelocation(Address - Section->getAddress(), Symbol, Type, Addend);
+  Section->addRelocation(Address - Section->getAddress(),
+                         Symbol,
+                         Type,
+                         Addend,
+                         Value);
 }
 
-void BinaryContext::removeRelocationAt(uint64_t Address) {
+bool BinaryContext::removeRelocationAt(uint64_t Address) {
   auto Section = getSectionForAddress(Address);
   assert(Section && "cannot find section for address");
-  Section->removeRelocationAt(Address - Section->getAddress());
+  return Section->removeRelocationAt(Address - Section->getAddress());
 }
 
 const Relocation *BinaryContext::getRelocationAt(uint64_t Address) {

bolt/BinaryContext.h

@@ -16,6 +16,7 @@
 
 #include "BinarySection.h"
 #include "DebugData.h"
+#include "llvm/ADT/iterator.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/DebugInfo/DWARF/DWARFCompileUnit.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
@@ -58,9 +59,17 @@ class BinaryContext {
   BinaryContext() = delete;
 
   /// Set of all sections.
-  using SectionSetType = std::set<BinarySection>;
+  struct CompareSections {
+    bool operator()(const BinarySection *A, const BinarySection *B) const {
+      return *A < *B;
+    }
+  };
+  using SectionSetType = std::set<BinarySection *, CompareSections>;
   SectionSetType Sections;
 
+  using SectionIterator = pointee_iterator<SectionSetType::iterator>;
+  using SectionConstIterator = pointee_iterator<SectionSetType::const_iterator>;
+
   /// Map virtual address to a section.  It is possible to have more than one
   /// section mapped to the same address, e.g. non-allocatable sections.
   using AddressToSectionMapType = std::multimap<uint64_t, BinarySection *>;
@@ -70,6 +79,9 @@ class BinaryContext {
   /// have multiple sections with the same name.
   using NameToSectionMapType = std::multimap<std::string, BinarySection *>;
   NameToSectionMapType NameToSection;
+
+  /// Low level section registration.
+  BinarySection &registerSection(BinarySection *Section);
 public:
 
   /// [name] -> [address] map used for global symbol resolution.
@@ -125,8 +137,6 @@ public:
 
   std::unique_ptr<MCAsmBackend> MAB;
 
-  std::function<void(std::error_code)> ErrorCheck;
-
   DataReader &DR;
 
   /// Indicates if relocations are availabe for usage.
@@ -224,18 +234,53 @@ public:
   ErrorOr<ArrayRef<uint8_t>>
   getFunctionData(const BinaryFunction &Function) const;
 
-  /// Register information about the given section so we can look up
+  /// Register information about the given \p Section so we can look up
-  /// sections for addresses.
+  /// sections by address.
   BinarySection &registerSection(SectionRef Section);
 
-  iterator_range<SectionSetType::iterator> sections() {
+  /// Register or update the information for the section with the given
+  /// /p Name.  If the section already exists, the information in the
+  /// section will be updated with the new data.
+  BinarySection &registerOrUpdateSection(StringRef Name,
+                                         unsigned ELFType,
+                                         unsigned ELFFlags,
+                                         uint8_t *Data = nullptr,
+                                         uint64_t Size = 0,
+                                         unsigned Alignment = 1,
+                                         bool IsLocal = false);
+
+  /// Register the information for the note (non-allocatable) section
+  /// with the given /p Name.  If the section already exists, the
+  /// information in the section will be updated with the new data.
+  BinarySection &registerOrUpdateNoteSection(StringRef Name,
+                                             uint8_t *Data = nullptr,
+                                             uint64_t Size = 0,
+                                             unsigned Alignment = 1,
+                                             bool IsReadOnly = true,
+                                             unsigned ELFType = ELF::SHT_PROGBITS,
+                                             bool IsLocal = false) {
+    return registerOrUpdateSection(Name, ELFType,
+                                   BinarySection::getFlags(IsReadOnly),
+                                   Data, Size, Alignment, IsLocal);
+  }
+
+  /// Remove the given /p Section from the set of all sections.  Return
+  /// true if the section was removed (and deleted), otherwise false.
+  bool deregisterSection(BinarySection &Section);
+
+  /// Iterate over all registered sections.
+  iterator_range<SectionIterator> sections() {
     return make_range(Sections.begin(), Sections.end());
   }
 
-  iterator_range<SectionSetType::const_iterator> sections() const {
+  /// Iterate over all registered sections.
+  iterator_range<SectionConstIterator> sections() const {
     return make_range(Sections.begin(), Sections.end());
   }
 
+  /// Print all sections.
+  void printSections(raw_ostream& OS) const;
+
   /// Return largest section containing the given \p Address.  These
   /// functions only work for allocatable sections, i.e. ones with non-zero
   /// addresses.
@@ -305,17 +350,12 @@ public:
                     BinaryFunction &ParentBF,
                     std::map<uint64_t, BinaryFunction> &BFs);
 
-  /// Add relocation for \p Section at a given \p Offset.
+  /// Add a Section relocation at a given \p Address.
-  void addSectionRelocation(BinarySection &Section, uint64_t Offset,
-                            MCSymbol *Symbol, uint64_t Type,
-                            uint64_t Addend = 0);
-
-  /// Add a relocation at a given \p Address.
   void addRelocation(uint64_t Address, MCSymbol *Symbol, uint64_t Type,
-                     uint64_t Addend = 0);
+                     uint64_t Addend = 0, uint64_t Value = 0);
 
   /// Remove registered relocation at a given \p Address.
-  void removeRelocationAt(uint64_t Address);
+  bool removeRelocationAt(uint64_t Address);
 
   /// Return a relocation registered at a given \p Address, or nullptr if there
   /// is no relocation at such address.

bolt/BinaryFunction.cpp

@@ -1312,8 +1312,18 @@ void BinaryFunction::postProcessJumpTables() {
         TakenBranches.emplace_back(JTSiteOffset, TargetOffset);
 
       // Take ownership of jump table relocations.
-      if (BC.HasRelocations)
+      if (BC.HasRelocations) {
-        BC.removeRelocationAt(JT->Address + EntryOffset);
+        auto EntryAddress = JT->Address + EntryOffset;
+        auto Res = BC.removeRelocationAt(EntryAddress);
+        (void)Res;
+        DEBUG(
+          auto Section = BC.getSectionForAddress(EntryAddress);
+          auto Offset = EntryAddress - Section->getAddress();
+          dbgs() << "BOLT-DEBUG: removing relocation from section "
+                 << Section->getName() << " at offset 0x"
+                 << Twine::utohexstr(Offset) << " = "
+                 << Res << '\n');
+      }
 
       EntryOffset += JT->EntrySize;
 
@@ -3363,7 +3373,7 @@ void BinaryFunction::JumpTable::updateOriginal(BinaryContext &BC) {
                  << " at offset " << Twine::utohexstr(Offset) << " for symbol "
                  << Entry->getName() << " with addend "
                  << Twine::utohexstr(RelAddend) << '\n');
-    BC.addSectionRelocation(*Section, Offset, Entry, RelType, RelAddend);
+    Section->addRelocation(Offset, Entry, RelType, RelAddend);
     Offset += EntrySize;
   }
 }

bolt/BinaryFunction.h

@@ -24,6 +24,7 @@
 #include "DebugData.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/ilist.h"
+#include "llvm/ADT/iterator.h"
 #include "llvm/BinaryFormat/Dwarf.h"
 #include "llvm/MC/MCCodeEmitter.h"
 #include "llvm/MC/MCContext.h"
@@ -600,7 +601,7 @@ public:
     std::map<unsigned, MCSymbol *> Labels;
 
     /// Corresponding section if any.
-    SectionInfo *SecInfo{nullptr};
+    ErrorOr<BinarySection &> Section{std::errc::bad_address};
 
     /// Corresponding section name if any.
     std::string SectionName;
@@ -747,21 +748,6 @@ private:
   /// Count the number of functions created.
   static uint64_t Count;
 
-  template <typename Itr, typename T>
-  class Iterator : public std::iterator<std::bidirectional_iterator_tag, T> {
-   public:
-    Iterator &operator++() { ++itr; return *this; }
-    Iterator &operator--() { --itr; return *this; }
-    Iterator operator++(int) { auto tmp(itr); itr++; return tmp; }
-    Iterator operator--(int) { auto tmp(itr); itr--; return tmp; }
-    bool operator==(const Iterator& other) const { return itr == other.itr; }
-    bool operator!=(const Iterator& other) const { return itr != other.itr; }
-    T& operator*() { return **itr; }
-    Iterator(Itr itr) : itr(itr) { }
-   private:
-    Itr itr;
-  };
-
   /// Register alternative function name.
   void addAlternativeName(std::string NewName) {
     Names.emplace_back(NewName);
@@ -842,13 +828,12 @@ public:
 
   BinaryFunction(BinaryFunction &&) = default;
 
-  typedef Iterator<BasicBlockListType::iterator, BinaryBasicBlock> iterator;
+  using iterator = pointee_iterator<BasicBlockListType::iterator>;
-  typedef Iterator<BasicBlockListType::const_iterator,
+  using const_iterator = pointee_iterator<BasicBlockListType::const_iterator>;
-                   const BinaryBasicBlock> const_iterator;
+  using reverse_iterator =
-  typedef Iterator<BasicBlockListType::reverse_iterator,
+    pointee_iterator<BasicBlockListType::reverse_iterator>;
-                   BinaryBasicBlock> reverse_iterator;
+  using const_reverse_iterator =
-  typedef Iterator<BasicBlockListType::const_reverse_iterator,
+    pointee_iterator<BasicBlockListType::const_reverse_iterator>;
-                   const BinaryBasicBlock> const_reverse_iterator;
 
   typedef BasicBlockOrderType::iterator order_iterator;
   typedef BasicBlockOrderType::const_iterator const_order_iterator;

bolt/BinarySection.cpp

@@ -12,10 +12,18 @@
 #include "BinarySection.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCStreamer.h"
+#include "llvm/Support/CommandLine.h"
+
+#undef  DEBUG_TYPE
+#define DEBUG_TYPE "bolt"
 
 using namespace llvm;
 using namespace bolt;
 
+namespace opts {
+extern cl::opt<bool> PrintRelocations;
+}
+
 Triple::ArchType Relocation::Arch;
 
 bool Relocation::isSupported(uint64_t Type) {
@@ -324,3 +332,29 @@ void Relocation::print(raw_ostream &OS) const {
     OS << ", 0x" << Twine::utohexstr(Addend);
   OS << ", 0x" << Twine::utohexstr(Value);
 }
+
+BinarySection::~BinarySection() {
+  if (!isAllocatable() &&
+      (!hasSectionRef() ||
+       OutputContents.data() != getContents(Section).data())) {
+    delete[] getOutputData();
+  }
+}
+
+void BinarySection::print(raw_ostream &OS) const {
+  OS << getName() << ", "
+     << "0x" << Twine::utohexstr(getAddress()) << ", "
+     << getSize()
+     << " (0x" << Twine::utohexstr(getFileAddress()) << ", "
+     << getOutputSize() << ")"
+     << ", data = " << getData()
+     << ", output data = " << getOutputData();
+
+  if (isAllocatable())
+    OS << " (allocatable)";
+
+  if (opts::PrintRelocations) {
+    for (auto &R : relocations())
+      OS << "\n  " << R;
+  }
+}

bolt/BinarySection.h

@@ -1,4 +1,4 @@
-//===--- BinarySection.h  - Interface for object file section -------------===//
+//===--- BinarySection.h - Interface for object file section --------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
@@ -12,6 +12,7 @@
 #ifndef LLVM_TOOLS_LLVM_BOLT_BINARY_SECTION_H
 #define LLVM_TOOLS_LLVM_BOLT_BINARY_SECTION_H
 
+#include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/MC/MCSymbol.h"
@@ -91,79 +92,167 @@ inline raw_ostream &operator<<(raw_ostream &OS, const Relocation &Rel) {
   return OS;
 }
 
-/// A wrapper around SectionRef that also manages related relocations
+inline uint8_t *copyByteArray(const uint8_t *Data, uint64_t Size) {
-class BinarySection {
+  auto Array = new uint8_t[Size];
-  SectionRef Section;
+  memcpy(Array, Data, Size);
-  std::set<Relocation> Relocations;
+  return Array;
-public:
+}
-  explicit BinarySection(SectionRef Section) : Section(Section) { }
 
-  StringRef getName() const {
+inline uint8_t *copyByteArray(StringRef Buffer) {
+  return copyByteArray(reinterpret_cast<const uint8_t*>(Buffer.data()),
+                       Buffer.size());
+}
+
+inline uint8_t *copyByteArray(ArrayRef<char> Buffer) {
+  return copyByteArray(reinterpret_cast<const uint8_t*>(Buffer.data()),
+                       Buffer.size());
+}
+
+/// A class to manage binary sections that also manages related relocations
+class BinarySection {
+  friend class BinaryContext;
+
+  const std::string Name;     // Section name
+  const SectionRef Section;   // SectionRef (may be null)
+  StringRef Contents;         // input section contents
+  const uint64_t Address;     // address of section in input binary (may be 0)
+  const uint64_t Size;        // input section size
+  unsigned Alignment;         // alignment in bytes (must be > 0)
+  unsigned ELFType;           // ELF section type
+  unsigned ELFFlags;          // ELF section flags
+  bool IsLocal;               // Is this a local section?
+
+  // Relocations associated with this section.  Relocation offsets are
+  // wrt. to the original section address and size.
+  using RelocationSetType = std::set<Relocation>;
+  RelocationSetType Relocations;
+
+  // Pending relocations for this section.  For the moment, just used by
+  // the .debug_info section.  TODO: it would be nice to get rid of this.
+  RelocationSetType PendingRelocations;
+
+  // Output info
+  bool IsFinalized{false};    // Has this section had output information
+                              // finalized?
+  uint64_t FileAddress{0};    // Section address for the rewritten binary.
+  uint64_t OutputSize{0};     // Section size in the rewritten binary.
+  uint64_t FileOffset{0};     // File offset in the rewritten binary file.
+  StringRef OutputContents;   // Rewritten section contents.
+  unsigned SectionID{-1u};    // Unique ID used for address mapping.
+                              // Set by ExecutableFileMemoryManager.
+
+  // non-copyable
+  BinarySection(const BinarySection &) = delete;
+  BinarySection(BinarySection &&) = delete;
+  BinarySection &operator=(const BinarySection &) = delete;
+  BinarySection &operator=(BinarySection &&) = delete;
+
+  static StringRef getName(SectionRef Section) {
     StringRef Name;
     Section.getName(Name);
     return Name;
   }
-  uint64_t getAddress() const { return Section.getAddress(); }
+  static StringRef getContents(SectionRef Section) {
-  uint64_t getEndAddress() const { return getAddress() + getSize(); }
-  uint64_t getSize() const { return Section.getSize(); }
-  uint64_t getAlignment() const { return Section.getAlignment(); }
-  bool containsAddress(uint64_t Address) const {
-    return getAddress() <= Address && Address < getEndAddress();
-  }
-  bool containsRange(uint64_t Address, uint64_t Size) const {
-    return getAddress() <= Address && Address + Size <= getEndAddress();
-  }
-  bool isReadOnly() const { return Section.isReadOnly(); }
-  bool isVirtual() const { return Section.isVirtual(); }
-  bool isText() const { return Section.isText(); }
-  bool isAllocatable() const { return getFlags() & ELF::SHF_ALLOC; }
-  StringRef getContents() const {
     StringRef Contents;
-    if (auto EC = Section.getContents(Contents)) {
+    if (ELFSectionRef(Section).getType() != ELF::SHT_NOBITS) {
-      errs() << "BOLT-ERROR: cannot get section contents for "
+      if (auto EC = Section.getContents(Contents)) {
-             << getName() << ": " << EC.message() << ".\n";
+        errs() << "BOLT-ERROR: cannot get section contents for "
-      exit(1);
+               << getName(Section) << ": " << EC.message() << ".\n";
+        exit(1);
+      }
     }
     return Contents;
   }
-  unsigned getFlags() const { return ELFSectionRef(Section).getFlags(); }
-  unsigned getType() const { return ELFSectionRef(Section).getType(); }
-  SectionRef getSectionRef() const { return Section; }
 
-  iterator_range<std::set<Relocation>::iterator> relocations() {
+  // Set output info for this section.
-    return make_range(Relocations.begin(), Relocations.end());
+  void update(uint8_t *NewData,
+              uint64_t NewSize,
+              unsigned NewAlignment,
+              unsigned NewELFType,
+              unsigned NewELFFlags,
+              bool NewIsLocal) {
+    assert(NewAlignment > 0 && "section alignment must be > 0");
+    OutputSize = NewSize;
+    Alignment = NewAlignment;
+    ELFType = NewELFType;
+    ELFFlags = NewELFFlags;
+    IsLocal = NewIsLocal || StringRef(Name).startswith(".local.");
+    OutputContents = StringRef(reinterpret_cast<const char*>(NewData),
+                               NewData ? NewSize : 0);
+    IsFinalized = true;
+  }
+public:
+  explicit BinarySection(SectionRef Section, bool IsLocal = false)
+    : Name(getName(Section)),
+      Section(Section),
+      Contents(getContents(Section)),
+      Address(Section.getAddress()),
+      Size(Section.getSize()),
+      Alignment(Section.getAlignment()),
+      ELFType(ELFSectionRef(Section).getType()),
+      ELFFlags(ELFSectionRef(Section).getFlags()),
+      IsLocal(IsLocal || StringRef(Name).startswith(".local.")),
+      OutputSize(0) {
   }
 
-  iterator_range<std::set<Relocation>::const_iterator> relocations() const {
+  // TODO: pass Data as StringRef/ArrayRef? use StringRef::copy method.
-    return make_range(Relocations.begin(), Relocations.end());
+  BinarySection(StringRef Name,
+                uint8_t *Data,
+                uint64_t Size,
+                unsigned Alignment,
+                unsigned ELFType,
+                unsigned ELFFlags,
+                bool IsLocal)
+    : Name(Name),
+      Contents(reinterpret_cast<const char*>(Data), Data ? Size : 0),
+      Address(0),
+      Size(Size),
+      Alignment(Alignment),
+      ELFType(ELFType),
+      ELFFlags(ELFFlags),
+      IsLocal(IsLocal || Name.startswith(".local.")),
+      IsFinalized(true),
+      OutputSize(Size),
+      OutputContents(Contents) {
+    assert(Alignment > 0 && "section alignment must be > 0");
   }
 
-  bool hasRelocations() const {
+  ~BinarySection();
-    return !Relocations.empty();
+
+  /// Helper function to generate the proper ELF flags from section properties.
+  static unsigned getFlags(bool IsReadOnly = true,
+                           bool IsText = false,
+                           bool IsAllocatable = false) {
+    unsigned Flags = 0;
+    if (IsAllocatable)
+      Flags |= ELF::SHF_ALLOC;
+    if (!IsReadOnly)
+      Flags |= ELF::SHF_WRITE;
+    if (IsText)
+      Flags |= ELF::SHF_EXECINSTR;
+    return Flags;
   }
 
-  void removeRelocationAt(uint64_t Offset) {
+  operator bool() const {
-    Relocation Key{Offset, 0, 0, 0, 0};
+    return ELFType != ELF::SHT_NULL;
-    auto Itr = Relocations.find(Key);
-    if (Itr != Relocations.end())
-      Relocations.erase(Itr);
   }
 
-  void addRelocation(uint64_t Offset,
+  bool operator==(const BinarySection &Other) const {
-                     MCSymbol *Symbol,
+    return (Name == Other.Name &&
-                     uint64_t Type,
+            Address == Other.Address &&
-                     uint64_t Addend,
+            Size == Other.Size &&
-                     uint64_t Value = 0) {
+            getData() == Other.getData() &&
-    assert(Offset < getSize());
+            Alignment == Other.Alignment &&
-    Relocations.emplace(Relocation{Offset, Symbol, Type, Addend, Value});
+            ELFType == Other.ELFType &&
+            ELFFlags == Other.ELFFlags &&
+            IsLocal == Other.IsLocal);
   }
 
-  const Relocation *getRelocationAt(uint64_t Offset) const {
+  bool operator!=(const BinarySection &Other) const {
-    Relocation Key{Offset, 0, 0, 0, 0};
+    return !operator==(Other);
-    auto Itr = Relocations.find(Key);
-    return Itr != Relocations.end() ? &*Itr : nullptr;
   }
 
+  // Order sections by their immutable properties.
   bool operator<(const BinarySection &Other) const {
     return (getAddress() < Other.getAddress() ||
             (getAddress() == Other.getAddress() &&
@@ -171,8 +260,170 @@ public:
               (getSize() == Other.getSize() &&
                getName() < Other.getName()))));
   }
+
+  ///
+  /// Basic proprety access.
+  ///
+  StringRef getName() const { return Name; }
+  uint64_t getAddress() const { return Address; }
+  uint64_t getEndAddress() const { return Address + Size; }
+  uint64_t getSize() const { return Size; }
+  uint64_t getAlignment() const { return Alignment; }
+  bool isText() const {
+    return (ELFFlags & ELF::SHF_EXECINSTR);
+  }
+  bool isData() const {
+    return (ELFType == ELF::SHT_PROGBITS &&
+            (ELFFlags & (ELF::SHF_ALLOC | ELF::SHF_WRITE)));
+  }
+  bool isBSS() const {
+    return (ELFType == ELF::SHT_NOBITS &&
+            (ELFFlags & (ELF::SHF_ALLOC | ELF::SHF_WRITE)));
+  }
+  bool isNote() const { return ELFType == ELF::SHT_NOTE; }
+  bool isStrTab() const { return ELFType == ELF::SHT_STRTAB; }
+  bool isSymTab() const { return ELFType == ELF::SHT_SYMTAB; }
+  bool isVirtual() const { return ELFType == ELF::SHT_NOBITS; }
+  bool isRela() const { return ELFType == ELF::SHT_RELA; }
+  bool isReadOnly() const {
+    return ((ELFFlags & ELF::SHF_ALLOC) &&
+            !(ELFFlags & ELF::SHF_WRITE) &&
+            ELFType == ELF::SHT_PROGBITS);
+  }
+  bool isAllocatable() const {
+    return (ELFFlags & ELF::SHF_ALLOC);
+  }
+  bool isLocal() const { return IsLocal; }
+  unsigned getELFType() const { return ELFType; }
+  unsigned getELFFlags() const { return ELFFlags; }
+
+  uint8_t *getData() {
+    return reinterpret_cast<uint8_t *>(const_cast<char *>(getContents().data()));
+  }
+  const uint8_t *getData() const {
+    return reinterpret_cast<const uint8_t *>(getContents().data());
+  }
+  StringRef getContents() const { return Contents; }
+  bool hasSectionRef() const { return Section != SectionRef(); }
+  SectionRef getSectionRef() const { return Section; }
+
+  /// Does this section contain the given /p Addr?
+  /// Note: this is in terms of the original mapped binary addresses.
+  bool containsAddress(uint64_t Addr) const {
+    return getAddress() <= Addr && Addr < getEndAddress();
+  }
+  /// Does this section contain the range given by /p Addr and /p Sz?
+  /// Note: this is in terms of the original mapped binary addresses.
+  bool containsRange(uint64_t Addr, uint64_t Sz) const {
+    return getAddress() <= Addr && Addr + Sz <= getEndAddress();
+  }
+
+  /// Iterate over all non-pending relocations for this section.
+  iterator_range<RelocationSetType::iterator> relocations() {
+    return make_range(Relocations.begin(), Relocations.end());
+  }
+
+  /// Iterate over all non-pending relocations for this section.
+  iterator_range<RelocationSetType::const_iterator> relocations() const {
+    return make_range(Relocations.begin(), Relocations.end());
+  }
+
+  /// Does this section have any non-pending relocations?
+  bool hasRelocations() const {
+    return !Relocations.empty();
+  }
+
+  /// Iterate over all pending relocations in this section.
+  iterator_range<RelocationSetType::const_iterator> pendingRelocations() const {
+    return make_range(PendingRelocations.begin(), PendingRelocations.end());
+  }
+
+  /// Does this section have any pending relocations?
+  bool hasPendingRelocations() const {
+    return !PendingRelocations.empty();
+  }
+
+  /// Remove non-pending relocation with the given /p Offset.
+  bool removeRelocationAt(uint64_t Offset) {
+    Relocation Key{Offset, 0, 0, 0, 0};
+    auto Itr = Relocations.find(Key);
+    if (Itr != Relocations.end()) {
+      Relocations.erase(Itr);
+      return true;
+    }
+    return false;
+  }
+
+  /// Add a new relocation at the given /p Offset.  Note: pending relocations
+  /// are only used by .debug_info and should eventually go away.
+  void addRelocation(uint64_t Offset,
+                     MCSymbol *Symbol,
+                     uint64_t Type,
+                     uint64_t Addend,
+                     uint64_t Value = 0,
+                     bool Pending = false) {
+    assert(Offset < getSize() && "offset not within section bounds");
+    if (!Pending) {
+      Relocations.emplace(Relocation{Offset, Symbol, Type, Addend, Value});
+    } else {
+      PendingRelocations.emplace(Relocation{Offset, Symbol, Type, Addend, Value});
+    }
+  }
+
+  /// Lookup the relocation (if any) at the given /p Offset.
+  const Relocation *getRelocationAt(uint64_t Offset) const {
+    Relocation Key{Offset, 0, 0, 0, 0};
+    auto Itr = Relocations.find(Key);
+    return Itr != Relocations.end() ? &*Itr : nullptr;
+  }
+
+  ///
+  /// Property accessors related to output data.
+  ///
+
+  bool isFinalized() const { return IsFinalized; }
+  void setIsFinalized() { IsFinalized = true; }
+  uint64_t getOutputSize() const { return OutputSize; }
+  uint8_t *getOutputData() {
+    return reinterpret_cast<uint8_t *>(const_cast<char *>(getOutputContents().data()));
+  }
+  const uint8_t *getOutputData() const {
+    return reinterpret_cast<const uint8_t *>(getOutputContents().data());
+  }
+  StringRef getOutputContents() const { return OutputContents; }
+  uint64_t getAllocAddress() const {
+    return reinterpret_cast<uint64_t>(getOutputData());
+  }
+  uint64_t getFileAddress() const { return FileAddress; }
+  uint64_t getFileOffset() const { return FileOffset; }
+  unsigned getSectionID() const {
+    assert(hasValidSectionID() && "trying to use uninitialized section id");
+    return SectionID;
+  }
+  bool hasValidSectionID() const {
+    return SectionID != -1u;
+  }
+
+  // mutation
+  void setFileAddress(uint64_t Address) {
+    FileAddress = Address;
+  }
+  void setFileOffset(uint64_t Offset) {
+    FileOffset = Offset;
+  }
+  void setSectionID(unsigned ID) {
+    assert(!hasValidSectionID() && "trying to set section id twice");
+    SectionID = ID;
+  }
+
+  void print(raw_ostream &OS) const;
 };
 
+inline raw_ostream &operator<<(raw_ostream &OS, const BinarySection &Section) {
+  Section.print(OS);
+  return OS;
+}
+
 } // namespace bolt
 } // namespace llvm
 

bolt/DWARFRewriter.cpp

@@ -444,9 +444,19 @@ void RewriteInstance::updateLineTableOffsets() {
     Offset += Label->getOffset() - CurrentOffset;
     CurrentOffset = Label->getOffset();
 
-    auto &SI = EFMM->NoteSectionInfo[".debug_info"];
+    auto DbgInfoSection = BC->getUniqueSectionByName(".debug_info");
-    SI.PendingRelocs.emplace_back(
+    assert(DbgInfoSection && ".debug_info section must exist");
-        SectionInfo::Reloc{LTOffset, 4, 0, Offset});
+    auto *Zero = BC->registerNameAtAddress("Zero", 0);
+    DbgInfoSection->addRelocation(LTOffset,
+                                  Zero,
+                                  ELF::R_X86_64_32,
+                                  Offset,
+                                  0,
+                                  /*Pending=*/true);
+    // Set .debug_info as finalized so it won't be skipped over when
+    // we process sections while writing out the new binary.  This ensures
+    // that the pending relocations will be processed and not ignored.
+    DbgInfoSection->setIsFinalized();
 
     DEBUG(dbgs() << "BOLT-DEBUG: CU " << CUIDLineTablePair.first
                 << " has line table at " << Offset << "\n");
@@ -466,41 +476,20 @@ void RewriteInstance::finalizeDebugSections() {
     RangesSectionsWriter->writeArangesSection(Writer.get());
     const auto &ARangesContents = OS.str();
 
-    // Freed by ExecutableFileMemoryManager.
+    BC->registerOrUpdateNoteSection(".debug_aranges",
-    uint8_t *SectionData = new uint8_t[ARangesContents.size()];
+                                    copyByteArray(ARangesContents),
-    memcpy(SectionData, ARangesContents.data(), ARangesContents.size());
+                                    ARangesContents.size());
-    EFMM->NoteSectionInfo[".debug_aranges"] = SectionInfo(
-        reinterpret_cast<uint64_t>(SectionData),
-        ARangesContents.size(),
-        /*Alignment=*/0,
-        /*IsCode=*/false,
-        /*IsReadOnly=*/true,
-        /*IsLocal=*/false);
   }
 
   auto RangesSectionContents = RangesSectionsWriter->finalize();
-  auto SectionSize = RangesSectionContents->size();
+  BC->registerOrUpdateNoteSection(".debug_ranges",
-  uint8_t *SectionData = new uint8_t[SectionSize];
+                                  copyByteArray(*RangesSectionContents),
-  memcpy(SectionData, RangesSectionContents->data(), SectionSize);
+                                  RangesSectionContents->size());
-  EFMM->NoteSectionInfo[".debug_ranges"] = SectionInfo(
-      reinterpret_cast<uint64_t>(SectionData),
-      SectionSize,
-      /*Alignment=*/1,
-      /*IsCode=*/false,
-      /*IsReadOnly=*/true,
-      /*IsLocal=*/false);
 
   auto LocationListSectionContents = LocationListWriter->finalize();
-  SectionSize = LocationListSectionContents->size();
+  BC->registerOrUpdateNoteSection(".debug_loc",
-  SectionData = new uint8_t[SectionSize];
+                                  copyByteArray(*LocationListSectionContents),
-  memcpy(SectionData, LocationListSectionContents->data(), SectionSize);
+                                  LocationListSectionContents->size());
-  EFMM->NoteSectionInfo[".debug_loc"] = SectionInfo(
-      reinterpret_cast<uint64_t>(SectionData),
-      SectionSize,
-      /*Alignment=*/1,
-      /*IsCode=*/false,
-      /*IsReadOnly=*/true,
-      /*IsLocal=*/false);
 }
 
 void RewriteInstance::updateGdbIndexSection() {
@@ -569,7 +558,7 @@ void RewriteInstance::updateGdbIndexSection() {
   size_t NewGdbIndexSize = GdbIndexContents.size() + Delta;
 
   // Free'd by ExecutableFileMemoryManager.
-  auto * const NewGdbIndexContents = new uint8_t[NewGdbIndexSize];
+  auto *NewGdbIndexContents = new uint8_t[NewGdbIndexSize];
   auto *Buffer = NewGdbIndexContents;
 
   write32le(Buffer, Version);
@@ -606,11 +595,7 @@ void RewriteInstance::updateGdbIndexSection() {
   memcpy(Buffer, Data, TrailingSize);
 
   // Register the new section.
-  EFMM->NoteSectionInfo[".gdb_index"] = SectionInfo(
+  BC->registerOrUpdateNoteSection(".gdb_index",
-      reinterpret_cast<uint64_t>(NewGdbIndexContents),
+                                  NewGdbIndexContents,
-      NewGdbIndexSize,
+                                  NewGdbIndexSize);
-      /*Alignment=*/0,
-      /*IsCode=*/false,
-      /*IsReadOnly=*/true,
-      /*IsLocal=*/false);
 }

bolt/RewriteInstance.h

@@ -39,43 +39,6 @@ class CFIReaderWriter;
 class DataAggregator;
 class DataReader;
 
-/// Section information for mapping and re-writing.
-struct SectionInfo {
-  uint64_t AllocAddress{0};   /// Current location of the section in memory.
-  uint64_t Size{0};           /// Section size.
-  unsigned Alignment{0};      /// Alignment of the section.
-  bool     IsCode{false};     /// Does this section contain code?
-  bool     IsReadOnly{false}; /// Is the section read-only?
-  bool     IsLocal{false};    /// Is this section local to a function, and
-                              /// should only be emitted with the function?
-  bool     IsStrTab{false};   /// Is this a string table section.
-  uint64_t FileAddress{0};    /// Address for the output file (final address).
-  uint64_t FileOffset{0};     /// Offset in the output file.
-  unsigned SectionID{0};      /// Unique ID used for address mapping.
-  bool     IsELFNote{false};  /// Is ELF note section?
-
-  struct Reloc {
-    uint32_t Offset;
-    uint8_t  Size;
-    uint8_t  Type; // unused atm
-    uint32_t Value;
-  };
-
-  /// Pending relocations for the section.
-  std::vector<Reloc> PendingRelocs;
-
-  SectionInfo(uint64_t Address, uint64_t Size, unsigned Alignment, bool IsCode,
-              bool IsReadOnly, bool IsLocal, uint64_t FileAddress = 0,
-              uint64_t FileOffset = 0, unsigned SectionID = 0,
-              bool IsELFNote = false)
-
-      : AllocAddress(Address), Size(Size), Alignment(Alignment), IsCode(IsCode),
-        IsReadOnly(IsReadOnly), IsLocal(IsLocal), FileAddress(FileAddress),
-        FileOffset(FileOffset), SectionID(SectionID), IsELFNote(IsELFNote) {}
-
-  SectionInfo() {}
-};
-
 struct SegmentInfo {
   uint64_t Address;           /// Address of the segment in memory.
   uint64_t Size;              /// Size of the segment in memory.
@@ -105,20 +68,15 @@ private:
                            StringRef SectionName,
                            bool IsCode,
                            bool IsReadOnly);
-
+  BinaryContext &BC;
   bool AllowStubs;
 
 public:
   /// [start memory address] -> [segment info] mapping.
   std::map<uint64_t, SegmentInfo> SegmentMapInfo;
 
-  /// Keep [section name] -> [section info] map for later remapping.
+  ExecutableFileMemoryManager(BinaryContext &BC, bool AllowStubs)
-  std::map<std::string, SectionInfo> SectionMapInfo;
+    : BC(BC), AllowStubs(AllowStubs) {}
-
-  /// Information about non-allocatable sections.
-  std::map<std::string, SectionInfo> NoteSectionInfo;
-
-  ExecutableFileMemoryManager(bool AllowStubs) : AllowStubs(AllowStubs) {}
 
   ~ExecutableFileMemoryManager();
 
@@ -202,7 +160,7 @@ public:
   /// non-empty.
   void emitDataSection(MCStreamer *Streamer,
                        const BinarySection &Section,
-                       std::string Name = "");
+                       StringRef Name = StringRef());
 
   /// Emit data sections that have code references in them.
   void emitDataSections(MCStreamer *Streamer);
@@ -312,7 +270,7 @@ private:
   void rewriteNoteSections();
 
   /// Write .eh_frame_hdr.
-  void writeEHFrameHeader(SectionInfo &EHFrameSecInfo);
+  void writeEHFrameHeader();
 
   /// Disassemble and create function entries for PLT.
   void disassemblePLT();