[BOLT][DWARF][NFC] Refactor address writers (#98094)

Refactors address writers to create an instance for each CU and its DWO
CU.

Author: sayhaan

bolt/include/bolt/Core/DebugData.h

@@ -256,7 +256,7 @@ class DebugRangeListsSectionWriter : public DebugRangesSectionWriter {
   };
   virtual ~DebugRangeListsSectionWriter(){};
 
-  static void setAddressWriter(DebugAddrWriter *AddrW) { AddrWriter = AddrW; }
+  void setAddressWriter(DebugAddrWriter *AddrW) { AddrWriter = AddrW; }
 
   /// Add ranges with caching.
   uint64_t addRanges(
@@ -284,7 +284,7 @@ class DebugRangeListsSectionWriter : public DebugRangesSectionWriter {
   }
 
 private:
-  static DebugAddrWriter *AddrWriter;
+  DebugAddrWriter *AddrWriter = nullptr;
   /// Used to find unique CU ID.
   DWARFUnit *CU;
   /// Current relative offset of range list entry within this CUs rangelist
@@ -336,21 +336,36 @@ using AddressSectionBuffer = SmallVector<char, 4>;
 class DebugAddrWriter {
 public:
   DebugAddrWriter() = delete;
-  DebugAddrWriter(BinaryContext *BC_);
+  DebugAddrWriter(BinaryContext *BC_) : DebugAddrWriter(BC_, UCHAR_MAX) {};
+  DebugAddrWriter(BinaryContext *BC_, uint8_t AddressByteSize);
   virtual ~DebugAddrWriter(){};
   /// Given an address returns an index in .debug_addr.
   /// Adds Address to map.
   uint32_t getIndexFromAddress(uint64_t Address, DWARFUnit &CU);
 
   /// Write out entries in to .debug_addr section for CUs.
-  virtual void update(DIEBuilder &DIEBlder, DWARFUnit &CUs);
+  virtual std::optional<uint64_t> finalize(const size_t BufferSize);
 
   /// Return buffer with all the entries in .debug_addr already writen out using
   /// update(...).
-  virtual AddressSectionBuffer &finalize() { return *Buffer; }
+  virtual std::unique_ptr<AddressSectionBuffer> releaseBuffer() {
+    return std::move(Buffer);
+  }
+
+  /// Returns buffer size.
+  virtual size_t getBufferSize() const { return Buffer->size(); }
+
+  /// Returns True if Buffer is not empty.
+  bool isInitialized() const { return !Buffer->empty(); }
 
-  /// Returns False if .debug_addr section was created..
-  bool isInitialized() const { return !AddressMaps.empty(); }
+  /// Updates address base with the given Offset.
+  virtual void updateAddrBase(DIEBuilder &DIEBlder, DWARFUnit &CU,
+                              const uint64_t Offset);
+
+  /// Appends an AddressSectionBuffer to the address writer's buffer.
+  void appendToAddressBuffer(const AddressSectionBuffer &Buffer) {
+    *AddressStream << Buffer;
+  }
 
 protected:
   class AddressForDWOCU {
@@ -407,23 +422,32 @@ class DebugAddrWriter {
   }
 
   BinaryContext *BC;
-  /// Maps DWOID to AddressForDWOCU.
-  std::unordered_map<uint64_t, AddressForDWOCU> AddressMaps;
+  /// Address for the DWO CU associated with the address writer.
+  AddressForDWOCU Map;
+  uint8_t AddressByteSize;
   /// Mutex used for parallel processing of debug info.
   std::mutex WriterMutex;
   std::unique_ptr<AddressSectionBuffer> Buffer;
   std::unique_ptr<raw_svector_ostream> AddressStream;
   /// Used to track sections that were not modified so that they can be re-used.
-  DenseMap<uint64_t, uint64_t> UnmodifiedAddressOffsets;
+  static DenseMap<uint64_t, uint64_t> UnmodifiedAddressOffsets;
 };
 
 class DebugAddrWriterDwarf5 : public DebugAddrWriter {
 public:
   DebugAddrWriterDwarf5() = delete;
   DebugAddrWriterDwarf5(BinaryContext *BC) : DebugAddrWriter(BC) {}
+  DebugAddrWriterDwarf5(BinaryContext *BC, uint8_t AddressByteSize,
+                        std::optional<uint64_t> AddrOffsetSectionBase)
+      : DebugAddrWriter(BC, AddressByteSize),
+        AddrOffsetSectionBase(AddrOffsetSectionBase) {}
 
   /// Write out entries in to .debug_addr section for CUs.
-  virtual void update(DIEBuilder &DIEBlder, DWARFUnit &CUs) override;
+  virtual std::optional<uint64_t> finalize(const size_t BufferSize) override;
+
+  /// Updates address base with the given Offset.
+  virtual void updateAddrBase(DIEBuilder &DIEBlder, DWARFUnit &CU,
+                              const uint64_t Offset) override;
 
 protected:
   /// Given DWARFUnit \p Unit returns either DWO ID or it's offset within
@@ -435,6 +459,10 @@ class DebugAddrWriterDwarf5 : public DebugAddrWriter {
     }
     return Unit.getOffset();
   }
+
+private:
+  std::optional<uint64_t> AddrOffsetSectionBase = std::nullopt;
+  static constexpr uint32_t HeaderSize = 8;
 };
 
 /// This class is NOT thread safe.
@@ -583,12 +611,10 @@ class DebugLoclistWriter : public DebugLocWriter {
 public:
   ~DebugLoclistWriter() {}
   DebugLoclistWriter() = delete;
-  DebugLoclistWriter(DWARFUnit &Unit, uint8_t DV, bool SD)
-      : DebugLocWriter(DV, LocWriterKind::DebugLoclistWriter), CU(Unit),
-        IsSplitDwarf(SD) {
-    assert(DebugLoclistWriter::AddrWriter &&
-           "Please use SetAddressWriter to initialize "
-           "DebugAddrWriter before instantiation.");
+  DebugLoclistWriter(DWARFUnit &Unit, uint8_t DV, bool SD,
+                     DebugAddrWriter &AddrW)
+      : DebugLocWriter(DV, LocWriterKind::DebugLoclistWriter),
+        AddrWriter(AddrW), CU(Unit), IsSplitDwarf(SD) {
     if (DwarfVersion >= 5) {
       LocBodyBuffer = std::make_unique<DebugBufferVector>();
       LocBodyStream = std::make_unique<raw_svector_ostream>(*LocBodyBuffer);
@@ -600,8 +626,6 @@ class DebugLoclistWriter : public DebugLocWriter {
     }
   }
 
-  static void setAddressWriter(DebugAddrWriter *AddrW) { AddrWriter = AddrW; }
-
   /// Stores location lists internally to be written out during finalize phase.
   virtual void addList(DIEBuilder &DIEBldr, DIE &Die, DIEValue &AttrInfo,
                        DebugLocationsVector &LocList) override;
@@ -630,7 +654,7 @@ class DebugLoclistWriter : public DebugLocWriter {
   /// Writes out locations in to a local buffer and applies debug info patches.
   void finalizeDWARF5(DIEBuilder &DIEBldr, DIE &Die);
 
-  static DebugAddrWriter *AddrWriter;
+  DebugAddrWriter &AddrWriter;
   DWARFUnit &CU;
   bool IsSplitDwarf{false};
   // Used for DWARF5 to store location lists before being finalized.

bolt/include/bolt/Rewrite/DWARFRewriter.h

@@ -66,10 +66,6 @@ class DWARFRewriter {
   /// .debug_aranges DWARF section.
   std::unique_ptr<DebugARangesSectionWriter> ARangesSectionWriter;
 
-  /// Stores and serializes information that will be put into the
-  /// .debug_addr DWARF section.
-  std::unique_ptr<DebugAddrWriter> AddrWriter;
-
   /// Stores and serializes information that will be put in to the
   /// .debug_addr DWARF section.
   /// Does not do de-duplication.
@@ -93,6 +89,10 @@ class DWARFRewriter {
   std::unordered_map<uint64_t, std::unique_ptr<DebugRangesSectionWriter>>
       LegacyRangesWritersByCU;
 
+  /// Stores address writer for each CU.
+  std::unordered_map<uint64_t, std::unique_ptr<DebugAddrWriter>>
+      AddressWritersByCU;
+
   std::mutex LocListDebugInfoPatchesMutex;
 
   /// Dwo id specific its RangesBase.
@@ -115,6 +115,7 @@ class DWARFRewriter {
   void updateUnitDebugInfo(DWARFUnit &Unit, DIEBuilder &DIEBldr,
                            DebugLocWriter &DebugLocWriter,
                            DebugRangesSectionWriter &RangesSectionWriter,
+                           DebugAddrWriter &AddressWriter,
                            std::optional<uint64_t> RangesBase = std::nullopt);
 
   /// Patches the binary for an object's address ranges to be updated.
@@ -141,13 +142,15 @@ class DWARFRewriter {
   /// Process and write out CUs that are passsed in.
   void finalizeCompileUnits(DIEBuilder &DIEBlder, DIEStreamer &Streamer,
                             CUOffsetMap &CUMap,
-                            const std::list<DWARFUnit *> &CUs);
+                            const std::list<DWARFUnit *> &CUs,
+                            DebugAddrWriter &FinalAddrWriter);
 
   /// Finalize debug sections in the main binary.
   void finalizeDebugSections(DIEBuilder &DIEBlder,
                              DWARF5AcceleratorTable &DebugNamesTable,
                              DIEStreamer &Streamer, raw_svector_ostream &ObjOS,
-                             CUOffsetMap &CUMap);
+                             CUOffsetMap &CUMap,
+                             DebugAddrWriter &FinalAddrWriter);
 
   /// Patches the binary for DWARF address ranges (e.g. in functions and lexical
   /// blocks) to be updated.

bolt/lib/Core/DebugData.cpp

@@ -184,8 +184,6 @@ void DebugRangesSectionWriter::appendToRangeBuffer(
   *RangesStream << CUBuffer;
 }
 
-DebugAddrWriter *DebugRangeListsSectionWriter::AddrWriter = nullptr;
-
 uint64_t DebugRangeListsSectionWriter::addRanges(
     DebugAddressRangesVector &&Ranges,
     std::map<DebugAddressRangesVector, uint64_t> &CachedRanges) {
@@ -390,7 +388,9 @@ void DebugARangesSectionWriter::writeARangesSection(
   }
 }
 
-DebugAddrWriter::DebugAddrWriter(BinaryContext *BC) : BC(BC) {
+DebugAddrWriter::DebugAddrWriter(BinaryContext *BC,
+                                 const uint8_t AddressByteSize)
+    : BC(BC), AddressByteSize(AddressByteSize) {
   Buffer = std::make_unique<AddressSectionBuffer>();
   AddressStream = std::make_unique<raw_svector_ostream>(*Buffer);
 }
@@ -405,11 +405,6 @@ void DebugAddrWriter::AddressForDWOCU::dump() {
 }
 uint32_t DebugAddrWriter::getIndexFromAddress(uint64_t Address, DWARFUnit &CU) {
   std::lock_guard<std::mutex> Lock(WriterMutex);
-  const uint64_t CUID = getCUID(CU);
-  if (!AddressMaps.count(CUID))
-    AddressMaps[CUID] = AddressForDWOCU();
-
-  AddressForDWOCU &Map = AddressMaps[CUID];
   auto Entry = Map.find(Address);
   if (Entry == Map.end()) {
     auto Index = Map.getNextIndex();
@@ -449,29 +444,23 @@ static void updateAddressBase(DIEBuilder &DIEBlder, DebugAddrWriter &AddrWriter,
   }
 }
 
-void DebugAddrWriter::update(DIEBuilder &DIEBlder, DWARFUnit &CU) {
-  // Handling the case where debug information is a mix of Debug fission and
-  // monolithic.
-  if (!CU.getDWOId())
-    return;
-  const uint64_t CUID = getCUID(CU);
-  auto AM = AddressMaps.find(CUID);
-  // Adding to map even if it did not contribute to .debug_addr.
-  // The Skeleton CU might still have DW_AT_GNU_addr_base.
-  uint64_t Offset = Buffer->size();
-  // If does not exist this CUs DWO section didn't contribute to .debug_addr.
-  if (AM == AddressMaps.end())
-    return;
-  std::vector<IndexAddressPair> SortedMap(AM->second.indexToAddressBegin(),
-                                          AM->second.indexToAdddessEnd());
+void DebugAddrWriter::updateAddrBase(DIEBuilder &DIEBlder, DWARFUnit &CU,
+                                     const uint64_t Offset) {
+  updateAddressBase(DIEBlder, *this, CU, Offset);
+}
+
+std::optional<uint64_t> DebugAddrWriter::finalize(const size_t BufferSize) {
+  if (Map.begin() == Map.end())
+    return std::nullopt;
+  std::vector<IndexAddressPair> SortedMap(Map.indexToAddressBegin(),
+                                          Map.indexToAdddessEnd());
   // Sorting address in increasing order of indices.
   llvm::sort(SortedMap, llvm::less_first());
 
-  uint8_t AddrSize = CU.getAddressByteSize();
   uint32_t Counter = 0;
   auto WriteAddress = [&](uint64_t Address) -> void {
     ++Counter;
-    switch (AddrSize) {
+    switch (AddressByteSize) {
     default:
       assert(false && "Address Size is invalid.");
       break;
@@ -490,10 +479,19 @@ void DebugAddrWriter::update(DIEBuilder &DIEBlder, DWARFUnit &CU) {
       WriteAddress(0);
     WriteAddress(Val.second);
   }
-  updateAddressBase(DIEBlder, *this, CU, Offset);
+  return std::nullopt;
+}
+
+void DebugAddrWriterDwarf5::updateAddrBase(DIEBuilder &DIEBlder, DWARFUnit &CU,
+                                           const uint64_t Offset) {
+  /// Header for DWARF5 has size 8, so we add it to the offset.
+  updateAddressBase(DIEBlder, *this, CU, Offset + HeaderSize);
 }
 
-void DebugAddrWriterDwarf5::update(DIEBuilder &DIEBlder, DWARFUnit &CU) {
+DenseMap<uint64_t, uint64_t> DebugAddrWriter::UnmodifiedAddressOffsets;
+
+std::optional<uint64_t>
+DebugAddrWriterDwarf5::finalize(const size_t BufferSize) {
   // Need to layout all sections within .debug_addr
   // Within each section sort Address by index.
   const endianness Endian = BC->DwCtx->isLittleEndian()
@@ -504,55 +502,44 @@ void DebugAddrWriterDwarf5::update(DIEBuilder &DIEBlder, DWARFUnit &CU) {
                               Endian == llvm::endianness::little, 0);
   DWARFDebugAddrTable AddrTable;
   DIDumpOptions DumpOpts;
-  constexpr uint32_t HeaderSize = 8;
-  const uint64_t CUID = getCUID(CU);
-  const uint8_t AddrSize = CU.getAddressByteSize();
-  auto AMIter = AddressMaps.find(CUID);
   // A case where CU has entry in .debug_addr, but we don't modify addresses
   // for it.
-  if (AMIter == AddressMaps.end()) {
-    AMIter = AddressMaps.insert({CUID, AddressForDWOCU()}).first;
-    std::optional<uint64_t> BaseOffset = CU.getAddrOffsetSectionBase();
-    if (!BaseOffset)
-      return;
+  if (Map.begin() == Map.end()) {
+    if (!AddrOffsetSectionBase)
+      return std::nullopt;
     // Address base offset is to the first entry.
     // The size of header is 8 bytes.
-    uint64_t Offset = *BaseOffset - HeaderSize;
+    uint64_t Offset = *AddrOffsetSectionBase - HeaderSize;
     auto Iter = UnmodifiedAddressOffsets.find(Offset);
-    if (Iter != UnmodifiedAddressOffsets.end()) {
-      updateAddressBase(DIEBlder, *this, CU, Iter->getSecond());
-      return;
-    }
-    UnmodifiedAddressOffsets[Offset] = Buffer->size() + HeaderSize;
-    if (Error Err = AddrTable.extract(AddrData, &Offset, 5, AddrSize,
+    if (Iter != UnmodifiedAddressOffsets.end())
+      return Iter->second;
+    UnmodifiedAddressOffsets[Offset] = BufferSize;
+    if (Error Err = AddrTable.extract(AddrData, &Offset, 5, AddressByteSize,
                                       DumpOpts.WarningHandler)) {
       DumpOpts.RecoverableErrorHandler(std::move(Err));
-      return;
+      return std::nullopt;
     }
-
     uint32_t Index = 0;
     for (uint64_t Addr : AddrTable.getAddressEntries())
-      AMIter->second.insert(Addr, Index++);
+      Map.insert(Addr, Index++);
   }
 
-  updateAddressBase(DIEBlder, *this, CU, Buffer->size() + HeaderSize);
-
-  std::vector<IndexAddressPair> SortedMap(AMIter->second.indexToAddressBegin(),
-                                          AMIter->second.indexToAdddessEnd());
+  std::vector<IndexAddressPair> SortedMap(Map.indexToAddressBegin(),
+                                          Map.indexToAdddessEnd());
   // Sorting address in increasing order of indices.
   llvm::sort(SortedMap, llvm::less_first());
   // Writing out Header
-  const uint32_t Length = SortedMap.size() * AddrSize + 4;
+  const uint32_t Length = SortedMap.size() * AddressByteSize + 4;
   support::endian::write(*AddressStream, Length, Endian);
   support::endian::write(*AddressStream, static_cast<uint16_t>(5), Endian);
-  support::endian::write(*AddressStream, static_cast<uint8_t>(AddrSize),
+  support::endian::write(*AddressStream, static_cast<uint8_t>(AddressByteSize),
                          Endian);
   support::endian::write(*AddressStream, static_cast<uint8_t>(0), Endian);
 
   uint32_t Counter = 0;
   auto writeAddress = [&](uint64_t Address) -> void {
     ++Counter;
-    switch (AddrSize) {
+    switch (AddressByteSize) {
     default:
       llvm_unreachable("Address Size is invalid.");
       break;
@@ -571,6 +558,7 @@ void DebugAddrWriterDwarf5::update(DIEBuilder &DIEBlder, DWARFUnit &CU) {
       writeAddress(0);
     writeAddress(Val.second);
   }
+  return std::nullopt;
 }
 
 void DebugLocWriter::init() {
@@ -723,11 +711,11 @@ void DebugLoclistWriter::addList(DIEBuilder &DIEBldr, DIE &Die,
                                  DIEValue &AttrInfo,
                                  DebugLocationsVector &LocList) {
   if (DwarfVersion < 5)
-    writeLegacyLocList(AttrInfo, LocList, DIEBldr, Die, *AddrWriter, *LocBuffer,
+    writeLegacyLocList(AttrInfo, LocList, DIEBldr, Die, AddrWriter, *LocBuffer,
                        CU, *LocStream);
   else
     writeDWARF5LocList(NumberOfEntries, AttrInfo, LocList, Die, DIEBldr,
-                       *AddrWriter, *LocBodyBuffer, RelativeLocListOffsets, CU,
+                       AddrWriter, *LocBodyBuffer, RelativeLocListOffsets, CU,
                        *LocBodyStream);
 }
 
@@ -789,8 +777,6 @@ void DebugLoclistWriter::finalize(DIEBuilder &DIEBldr, DIE &Die) {
     finalizeDWARF5(DIEBldr, Die);
 }
 
-DebugAddrWriter *DebugLoclistWriter::AddrWriter = nullptr;
-
 static std::string encodeLE(size_t ByteSize, uint64_t NewValue) {
   std::string LE64(ByteSize, 0);
   for (size_t I = 0; I < ByteSize; ++I) {