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

Sayhaan Siddiqui via llvm-commits llvm-commits at lists.llvm.org
Mon Jul 8 15:55:16 PDT 2024


https://github.com/sayhaan updated https://github.com/llvm/llvm-project/pull/98094

>From 447b2193d4747508564c5d9cf2da189603d95400 Mon Sep 17 00:00:00 2001
From: Sayhaan Siddiqui <sayhaan at meta.com>
Date: Mon, 17 Jun 2024 10:16:44 -0700
Subject: [PATCH 1/2] [BOLT][DWARF] Refactor legacy ranges writers

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Differential Revision: https://phabricator.intern.facebook.com/D58679290
---
 bolt/include/bolt/Core/DebugData.h            |  12 ++
 bolt/include/bolt/Rewrite/DWARFRewriter.h     |   4 +
 bolt/lib/Core/DebugData.cpp                   |   6 +
 bolt/lib/Rewrite/DWARFRewriter.cpp            | 132 +++++++++++-------
 bolt/test/X86/debug-fission-single-convert.s  |   6 +-
 bolt/test/X86/dwarf4-df-dualcu.test           |  40 +++---
 .../X86/dwarf4-df-input-lowpc-ranges-cus.test |  80 ++++++-----
 .../X86/dwarf4-df-input-lowpc-ranges.test     |  37 ++---
 8 files changed, 190 insertions(+), 127 deletions(-)

diff --git a/bolt/include/bolt/Core/DebugData.h b/bolt/include/bolt/Core/DebugData.h
index 585bafa088849..144433ac78a37 100644
--- a/bolt/include/bolt/Core/DebugData.h
+++ b/bolt/include/bolt/Core/DebugData.h
@@ -210,6 +210,15 @@ class DebugRangesSectionWriter {
   static bool classof(const DebugRangesSectionWriter *Writer) {
     return Writer->getKind() == RangesWriterKind::DebugRangesWriter;
   }
+  
+  /// Append a range to the main buffer.
+  void appendToRangeBuffer(const DebugBufferVector &CUBuffer);
+
+  /// Sets Unit DIE to be updated for CU.
+  void setDie(DIE *Die) { this->Die = Die; }
+
+  /// Returns Unit DIE to be updated for CU.
+  DIE *getDie() const { return Die; }
 
   /// Writes out range lists for a current CU being processed.
   void virtual finalizeSection(){};
@@ -232,6 +241,9 @@ class DebugRangesSectionWriter {
   static constexpr uint64_t EmptyRangesOffset{0};
 
 private:
+  /// Stores Unit DIE to be updated for CU.
+  DIE *Die{0};
+
   RangesWriterKind Kind;
 };
 
diff --git a/bolt/include/bolt/Rewrite/DWARFRewriter.h b/bolt/include/bolt/Rewrite/DWARFRewriter.h
index 8dec32de9008e..1776f9efd137f 100644
--- a/bolt/include/bolt/Rewrite/DWARFRewriter.h
+++ b/bolt/include/bolt/Rewrite/DWARFRewriter.h
@@ -89,6 +89,10 @@ class DWARFRewriter {
   /// Store Rangelists writer for each DWO CU.
   RangeListsDWOWriers RangeListsWritersByCU;
 
+  /// Stores ranges writer for each DWO CU.
+  std::unordered_map<uint64_t, std::unique_ptr<DebugRangesSectionWriter>>
+      LegacyRangesWritersByCU;
+
   std::mutex LocListDebugInfoPatchesMutex;
 
   /// Dwo id specific its RangesBase.
diff --git a/bolt/lib/Core/DebugData.cpp b/bolt/lib/Core/DebugData.cpp
index f502a50312470..08d4c45aac791 100644
--- a/bolt/lib/Core/DebugData.cpp
+++ b/bolt/lib/Core/DebugData.cpp
@@ -177,6 +177,12 @@ uint64_t DebugRangesSectionWriter::getSectionOffset() {
   return SectionOffset;
 }
 
+void DebugRangesSectionWriter::appendToRangeBuffer(
+    const DebugBufferVector &CUBuffer) {
+  *RangesStream << CUBuffer;
+  SectionOffset = RangesBuffer->size();
+}
+
 DebugAddrWriter *DebugRangeListsSectionWriter::AddrWriter = nullptr;
 
 uint64_t DebugRangeListsSectionWriter::addRanges(
diff --git a/bolt/lib/Rewrite/DWARFRewriter.cpp b/bolt/lib/Rewrite/DWARFRewriter.cpp
index 8814ebbd10aa5..be7bab9359839 100644
--- a/bolt/lib/Rewrite/DWARFRewriter.cpp
+++ b/bolt/lib/Rewrite/DWARFRewriter.cpp
@@ -646,6 +646,15 @@ void DWARFRewriter::updateDebugInfo() {
 
     } else {
       LocListWritersByCU[CUIndex] = std::make_unique<DebugLocWriter>();
+      if (std::optional<uint64_t> DWOId = CU.getDWOId()) {
+        assert(LegacyRangesWritersByCU.count(*DWOId) == 0 &&
+               "LegacyRangeLists writer for DWO unit already exists.");
+        auto LegacyRangesSectionWriterByCU =
+            std::make_unique<DebugRangesSectionWriter>();
+        LegacyRangesSectionWriterByCU->initSection(CU);
+        LegacyRangesWritersByCU[*DWOId] =
+            std::move(LegacyRangesSectionWriterByCU);
+      }
     }
     return LocListWritersByCU[CUIndex++].get();
   };
@@ -692,6 +701,7 @@ void DWARFRewriter::updateDebugInfo() {
       if (Unit->getVersion() >= 5) {
         TempRangesSectionWriter = RangeListsWritersByCU[*DWOId].get();
       } else {
+        TempRangesSectionWriter = LegacyRangesWritersByCU[*DWOId].get();
         RangesBase = RangesSectionWriter->getSectionOffset();
         setDwoRangesBase(*DWOId, *RangesBase);
       }
@@ -1270,9 +1280,17 @@ void DWARFRewriter::updateDWARFObjectAddressRanges(
     }
 
     if (RangesBaseInfo) {
-      DIEBldr.replaceValue(&Die, RangesBaseInfo.getAttribute(),
-                           RangesBaseInfo.getForm(),
-                           DIEInteger(static_cast<uint32_t>(*RangesBase)));
+      if (RangesBaseInfo.getAttribute() == dwarf::DW_AT_GNU_ranges_base) {
+        auto RangesWriterIterator =
+            LegacyRangesWritersByCU.find(*Unit.getDWOId());
+        assert(RangesWriterIterator != LegacyRangesWritersByCU.end() &&
+               "RangesWriter does not exist for DWOId");
+        RangesWriterIterator->second->setDie(&Die);
+      } else {
+        DIEBldr.replaceValue(&Die, RangesBaseInfo.getAttribute(),
+                             RangesBaseInfo.getForm(),
+                             DIEInteger(static_cast<uint32_t>(*RangesBase)));
+      }
       RangesBase = std::nullopt;
     }
   }
@@ -1290,20 +1308,12 @@ void DWARFRewriter::updateDWARFObjectAddressRanges(
         RangesAttrInfo.getForm() == dwarf::DW_FORM_sec_offset)
       NeedConverted = true;
 
-    uint64_t CurRangeBase = 0;
-    if (Unit.isDWOUnit()) {
-      if (std::optional<uint64_t> DWOId = Unit.getDWOId())
-        CurRangeBase = getDwoRangesBase(*DWOId);
-      else
-        errs() << "BOLT-WARNING: [internal-dwarf-error]: DWOId is not found "
-                  "for DWO Unit.";
-    }
     if (NeedConverted || RangesAttrInfo.getForm() == dwarf::DW_FORM_rnglistx)
       DIEBldr.replaceValue(&Die, dwarf::DW_AT_ranges, dwarf::DW_FORM_rnglistx,
                            DIEInteger(DebugRangesOffset));
     else
       DIEBldr.replaceValue(&Die, dwarf::DW_AT_ranges, RangesAttrInfo.getForm(),
-                           DIEInteger(DebugRangesOffset - CurRangeBase));
+                           DIEInteger(DebugRangesOffset));
 
     if (!RangesBase) {
       if (LowPCAttrInfo &&
@@ -1320,15 +1330,21 @@ void DWARFRewriter::updateDWARFObjectAddressRanges(
 
     // If we are at this point we are in the CU/Skeleton CU, and
     // DW_AT_GNU_ranges_base or DW_AT_rnglists_base doesn't exist.
-    if (Unit.getVersion() <= 4)
+    if (Unit.getVersion() <= 4) {
       DIEBldr.addValue(&Die, dwarf::DW_AT_GNU_ranges_base, dwarf::DW_FORM_data4,
-                       DIEInteger(*RangesBase));
-    else if (Unit.getVersion() == 5)
+                       DIEInteger(INT_MAX));
+      auto RangesWriterIterator =
+          LegacyRangesWritersByCU.find(*Unit.getDWOId());
+      assert(RangesWriterIterator != LegacyRangesWritersByCU.end() &&
+             "RangesWriter does not exist for DWOId");
+      RangesWriterIterator->second->setDie(&Die);
+    } else if (Unit.getVersion() == 5) {
       DIEBldr.addValue(&Die, dwarf::DW_AT_rnglists_base,
                        dwarf::DW_FORM_sec_offset, DIEInteger(*RangesBase));
-    else
+    } else {
       DIEBldr.addValue(&Die, dwarf::DW_AT_rnglists_base,
                        dwarf::DW_FORM_sec_offset, DIEInteger(*RangesBase));
+    }
     return;
   }
 
@@ -1605,6 +1621,30 @@ void DWARFRewriter::finalizeCompileUnits(DIEBuilder &DIEBlder,
                                          DIEStreamer &Streamer,
                                          CUOffsetMap &CUMap,
                                          const std::list<DWARFUnit *> &CUs) {
+  for (DWARFUnit *CU : CUs) {
+    if (CU->getVersion() != 4)
+      continue;
+    std::optional<uint64_t> DWOId = CU->getDWOId();
+    if (!DWOId)
+      continue;
+    auto RangesWriterIterator = LegacyRangesWritersByCU.find(*DWOId);
+    assert(RangesWriterIterator != LegacyRangesWritersByCU.end() &&
+           "RangesWriter does not exist for DWOId");
+    std::unique_ptr<DebugRangesSectionWriter> &LegacyRangesWriter =
+        RangesWriterIterator->second;
+    std::optional<DIE *> Die = LegacyRangesWriter->getDie();
+    if (!Die || !Die.value())
+      continue;
+    DIEValue DvalGNUBase =
+        Die.value()->findAttribute(dwarf::DW_AT_GNU_ranges_base);
+    assert(DvalGNUBase && "GNU_ranges_base attribute does not exist for DWOId");
+    DIEBlder.replaceValue(
+        Die.value(), dwarf::DW_AT_GNU_ranges_base, DvalGNUBase.getForm(),
+        DIEInteger(LegacyRangesSectionWriter->getSectionOffset()));
+    std::unique_ptr<DebugBufferVector> RangesWritersContents =
+        LegacyRangesWriter->releaseBuffer();
+    LegacyRangesSectionWriter->appendToRangeBuffer(*RangesWritersContents);
+  }
   DIEBlder.generateAbbrevs();
   DIEBlder.finish();
   // generate debug_info and CUMap
@@ -2263,7 +2303,6 @@ void DWARFRewriter::convertToRangesPatchDebugInfo(
     DWARFUnit &Unit, DIEBuilder &DIEBldr, DIE &Die,
     uint64_t RangesSectionOffset, DIEValue &LowPCAttrInfo,
     DIEValue &HighPCAttrInfo, std::optional<uint64_t> RangesBase) {
-  uint32_t BaseOffset = 0;
   dwarf::Form LowForm = LowPCAttrInfo.getForm();
   dwarf::Attribute RangeBaseAttribute = dwarf::DW_AT_GNU_ranges_base;
   dwarf::Form RangesForm = dwarf::DW_FORM_sec_offset;
@@ -2278,45 +2317,40 @@ void DWARFRewriter::convertToRangesPatchDebugInfo(
                    Die.getTag() == dwarf::DW_TAG_skeleton_unit;
   if (!IsUnitDie)
     DIEBldr.deleteValue(&Die, LowPCAttrInfo.getAttribute());
-  // In DWARF4 for DW_AT_low_pc in binary DW_FORM_addr is used. In the DWO
-  // section DW_FORM_GNU_addr_index is used. So for if we are converting
-  // DW_AT_low_pc/DW_AT_high_pc and see DW_FORM_GNU_addr_index. We are
-  // converting in DWO section, and DW_AT_ranges [DW_FORM_sec_offset] is
-  // relative to DW_AT_GNU_ranges_base.
-  if (LowForm == dwarf::DW_FORM_GNU_addr_index) {
-    // Ranges are relative to DW_AT_GNU_ranges_base.
-    uint64_t CurRangeBase = 0;
-    if (std::optional<uint64_t> DWOId = Unit.getDWOId()) {
-      CurRangeBase = getDwoRangesBase(*DWOId);
-    }
-    BaseOffset = CurRangeBase;
-  } else {
-    // In DWARF 5 we can have DW_AT_low_pc either as DW_FORM_addr, or
-    // DW_FORM_addrx. Former is when DW_AT_rnglists_base is present. Latter is
-    // when it's absent.
-    if (IsUnitDie) {
-      if (LowForm == dwarf::DW_FORM_addrx) {
-        const uint32_t Index = AddrWriter->getIndexFromAddress(0, Unit);
-        DIEBldr.replaceValue(&Die, LowPCAttrInfo.getAttribute(),
-                             LowPCAttrInfo.getForm(), DIEInteger(Index));
-      } else {
-        DIEBldr.replaceValue(&Die, LowPCAttrInfo.getAttribute(),
-                             LowPCAttrInfo.getForm(), DIEInteger(0));
-      }
+
+  // In DWARF 5 we can have DW_AT_low_pc either as DW_FORM_addr, or
+  // DW_FORM_addrx. Former is when DW_AT_rnglists_base is present. Latter is
+  // when it's absent.
+  if (IsUnitDie) {
+    if (LowForm == dwarf::DW_FORM_addrx) {
+      const uint32_t Index = AddrWriter->getIndexFromAddress(0, Unit);
+      DIEBldr.replaceValue(&Die, LowPCAttrInfo.getAttribute(),
+                           LowPCAttrInfo.getForm(), DIEInteger(Index));
+    } else {
+      DIEBldr.replaceValue(&Die, LowPCAttrInfo.getAttribute(),
+                           LowPCAttrInfo.getForm(), DIEInteger(0));
     }
-    // Original CU didn't have DW_AT_*_base. We converted it's children (or
-    // dwo), so need to insert it into CU.
-    if (RangesBase)
+  }
+  // Original CU didn't have DW_AT_*_base. We converted it's children (or
+  // dwo), so need to insert it into CU.
+  if (RangesBase) {
+    if (Unit.getVersion() >= 5) {
       DIEBldr.addValue(&Die, RangeBaseAttribute, dwarf::DW_FORM_sec_offset,
                        DIEInteger(*RangesBase));
+    } else {
+      DIEBldr.addValue(&Die, RangeBaseAttribute, dwarf::DW_FORM_sec_offset,
+                       DIEInteger(INT_MAX));
+      auto RangesWriterIterator =
+          LegacyRangesWritersByCU.find(*Unit.getDWOId());
+      assert(RangesWriterIterator != LegacyRangesWritersByCU.end() &&
+             "RangesWriter does not exist for DWOId");
+      RangesWriterIterator->second->setDie(&Die);
+    }
   }
 
-  uint64_t RangeAttrVal = RangesSectionOffset - BaseOffset;
-  if (Unit.getVersion() >= 5)
-    RangeAttrVal = RangesSectionOffset;
   // HighPC was conveted into DW_AT_ranges.
   // For DWARF5 we only access ranges through index.
 
   DIEBldr.replaceValue(&Die, HighPCAttrInfo.getAttribute(), dwarf::DW_AT_ranges,
-                       RangesForm, DIEInteger(RangeAttrVal));
+                       RangesForm, DIEInteger(RangesSectionOffset));
 }
diff --git a/bolt/test/X86/debug-fission-single-convert.s b/bolt/test/X86/debug-fission-single-convert.s
index 28fcb6686e0a2..4cd881740b2f8 100644
--- a/bolt/test/X86/debug-fission-single-convert.s
+++ b/bolt/test/X86/debug-fission-single-convert.s
@@ -31,11 +31,11 @@
 # CHECK-DWO-DWO: 00000010
 # CHECK-DWO-DWO: 00000050
 # CHECK-DWO-DWO: DW_TAG_subprogram
-# CHECK-DWO-DWO-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000000
+# CHECK-DWO-DWO-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000010
 # CHECK-DWO-DWO: DW_TAG_subprogram
-# CHECK-DWO-DWO-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000020
+# CHECK-DWO-DWO-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000030
 # CHECK-DWO-DWO: DW_TAG_subprogram
-# CHECK-DWO-DWO-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000040
+# CHECK-DWO-DWO-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000050
 
 # CHECK-ADDR-SEC: .debug_addr contents:
 # CHECK-ADDR-SEC: 0x00000000: Addrs: [
diff --git a/bolt/test/X86/dwarf4-df-dualcu.test b/bolt/test/X86/dwarf4-df-dualcu.test
index b690623b70d83..fb328eb1872e0 100644
--- a/bolt/test/X86/dwarf4-df-dualcu.test
+++ b/bolt/test/X86/dwarf4-df-dualcu.test
@@ -37,36 +37,38 @@
 
 ; BOLT: .debug_ranges
 ; BOLT-NEXT: 00000000 <End of list>
-; BOLT-NEXT: 00000010 [[#%.16x,ADDR:]] [[#%.16x,ADDRB:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR1:]] [[#%.16x,ADDRB1:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR2:]] [[#%.16x,ADDRB2:]]
 ; BOLT-NEXT: 00000010 <End of list>
-; BOLT-NEXT: 00000030 [[#%.16x,ADDR1:]] [[#%.16x,ADDR1B:]]
-; BOLT-NEXT: 00000030 <End of list>
-; BOLT-NEXT: 00000050 [[#%.16x,ADDR2:]] [[#%.16x,ADDR2B:]]
-; BOLT-NEXT: 00000050 [[#%.16x,ADDR3:]] [[#%.16x,ADDR3B:]]
+; BOLT-NEXT: 00000040 <End of list>
+; BOLT-NEXT: 00000050 [[#%.16x,ADDR1:]] [[#%.16x,ADDRB1:]]
 ; BOLT-NEXT: 00000050 <End of list>
-; BOLT-NEXT: 00000080 [[#%.16x,ADDR4:]] [[#%.16x,ADDR4B:]]
-; BOLT-NEXT: 00000080 <End of list>
-; BOLT-NEXT: 000000a0 [[#%.16x,ADDR5:]] [[#%.16x,ADDR5B:]]
-; BOLT-NEXT: 000000a0 <End of list>
+; BOLT-NEXT: 00000070 [[#%.16x,ADDR2:]] [[#%.16x,ADDRB2:]]
+; BOLT-NEXT: 00000070 <End of list>
+; BOLT-NEXT: 00000090 [[#%.16x,ADDR3:]] [[#%.16x,ADDRB3:]]
+; BOLT-NEXT: 00000090 <End of list>
+; BOLT-NEXT: 000000b0 <End of list>
+; BOLT-NEXT: 000000c0 [[#%.16x,ADDR3:]] [[#%.16x,ADDRB3:]]
+; BOLT-NEXT: 000000c0 <End of list>
 
 ; BOLT: DW_TAG_compile_unit
 ; BOLT: DW_AT_GNU_dwo_name [DW_FORM_strp] ( .debug_str[0x00000016] = "main.dwo.dwo")
 ; BOLT-NEXT: DW_AT_GNU_dwo_id
 ; BOLT-NEXT: DW_AT_low_pc [DW_FORM_addr] (0x0000000000000000)
-; BOLT-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000050
-; BOLT-NEXT: [0x[[#ADDR2]], 0x[[#ADDR2B]])
-; BOLT-NEXT: [0x[[#ADDR3]], 0x[[#ADDR3B]]))
+; BOLT-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000010
+; BOLT-NEXT: [0x[[#ADDR1]], 0x[[#ADDRB1]])
+; BOLT-NEXT: [0x[[#ADDR2]], 0x[[#ADDRB2]]))
 ; BOLT-NEXT: DW_AT_GNU_addr_base [DW_FORM_sec_offset]  (0x00000000)
-; BOLT-NEXT: DW_AT_GNU_ranges_base [DW_FORM_sec_offset]  (0x00000010)
+; BOLT-NEXT: DW_AT_GNU_ranges_base [DW_FORM_sec_offset]  (0x00000040)
 ; BOLT-NEXT: Compile
 ; BOLT: DW_TAG_compile_unit
 ; BOLT: DW_AT_GNU_dwo_name [DW_FORM_strp] ( .debug_str[0x00000023] = "helper.dwo.dwo")
 ; BOLT-NEXT: DW_AT_GNU_dwo_id
 ; BOLT-NEXT: DW_AT_low_pc [DW_FORM_addr] (0x0000000000000000)
-; BOLT-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x000000a0
-; BOLT-NEXT: [0x[[#ADDR5]], 0x[[#ADDR5B]])
+; BOLT-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000090
+; BOLT-NEXT: [0x[[#ADDR3]], 0x[[#ADDRB3]])
 ; BOLT-NEXT: DW_AT_GNU_addr_base [DW_FORM_sec_offset]  (0x00000010)
-; BOLT-NEXT: DW_AT_GNU_ranges_base [DW_FORM_sec_offset]  (0x00000080)
+; BOLT-NEXT: DW_AT_GNU_ranges_base [DW_FORM_sec_offset]  (0x000000b0)
 
 ; PRE-BOLT-DWO-MAIN: version = 0x0004
 ; PRE-BOLT-DWO-MAIN: DW_TAG_compile_unit
@@ -113,13 +115,13 @@
 ; BOLT-DWO-MAIN-NEXT: DW_AT_decl_line
 ; BOLT-DWO-MAIN-NEXT: DW_AT_location [DW_FORM_exprloc]	(DW_OP_GNU_addr_index 0x1)
 ; BOLT-DWO-MAIN: DW_TAG_subprogram [4]
-; BOLT-DWO-MAIN-NEXT: DW_AT_ranges [DW_FORM_sec_offset]	(0x00000000
+; BOLT-DWO-MAIN-NEXT: DW_AT_ranges [DW_FORM_sec_offset]	(0x00000010
 ; BOLT-DWO-MAIN-NEXT: )
 ; BOLT-DWO-MAIN-NEXT: DW_AT_frame_base
 ; BOLT-DWO-MAIN-NEXT: DW_AT_linkage_name [DW_FORM_GNU_str_index]	(indexed (00000003) string = "_Z3usePiS_")
 ; BOLT-DWO-MAIN-NEXT: DW_AT_name [DW_FORM_GNU_str_index]	(indexed (00000004) string = "use")
 ; BOLT-DWO-MAIN: DW_TAG_subprogram [6]
-; BOLT-DWO-MAIN-NEXT: DW_AT_ranges [DW_FORM_sec_offset]	(0x00000020
+; BOLT-DWO-MAIN-NEXT: DW_AT_ranges [DW_FORM_sec_offset]	(0x00000030
 ; BOLT-DWO-MAIN-NEXT: )
 ; BOLT-DWO-MAIN-NEXT: DW_AT_frame_base [DW_FORM_exprloc]	(DW_OP_reg6 RBP)
 ; BOLT-DWO-MAIN-NEXT: DW_AT_name [DW_FORM_GNU_str_index]	(indexed (00000005) string = "main")
@@ -160,4 +162,4 @@
 ; BOLT-DWO-HELPER-NEXT: DW_AT_decl_line
 ; BOLT-DWO-HELPER-NEXT: DW_AT_location [DW_FORM_exprloc]	(DW_OP_GNU_addr_index 0x1)
 ; BOLT-DWO-HELPER: DW_TAG_subprogram [4]
-; BOLT-DWO-HELPER-NEXT: DW_AT_ranges [DW_FORM_sec_offset]	(0x00000000
+; BOLT-DWO-HELPER-NEXT: DW_AT_ranges [DW_FORM_sec_offset]	(0x00000010
diff --git a/bolt/test/X86/dwarf4-df-input-lowpc-ranges-cus.test b/bolt/test/X86/dwarf4-df-input-lowpc-ranges-cus.test
index c9abd02bbb7d9..cf9357d5f3c59 100644
--- a/bolt/test/X86/dwarf4-df-input-lowpc-ranges-cus.test
+++ b/bolt/test/X86/dwarf4-df-input-lowpc-ranges-cus.test
@@ -17,45 +17,47 @@
 
 ; BOLT: .debug_ranges
 ; BOLT-NEXT: 00000000 <End of list>
-; BOLT-NEXT: 00000010
-; BOLT-NEXT: 00000010
-; BOLT-NEXT: 00000010
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR1:]] [[#%.16x,ADDRB1:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR2:]] [[#%.16x,ADDRB2:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR3:]] [[#%.16x,ADDRB3:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR4:]] [[#%.16x,ADDRB4:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR5:]] [[#%.16x,ADDRB5:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR6:]] [[#%.16x,ADDRB6:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR7:]] [[#%.16x,ADDRB7:]]
 ; BOLT-NEXT: 00000010 <End of list>
-; BOLT-NEXT: 00000050
-; BOLT-NEXT: 00000050
-; BOLT-NEXT: 00000050
-; BOLT-NEXT: 00000050 <End of list>
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR1:]] [[#%.16x,ADDRB1:]]
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR2:]] [[#%.16x,ADDRB2:]]
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR3:]] [[#%.16x,ADDRB3:]]
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR4:]] [[#%.16x,ADDRB4:]]
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR5:]] [[#%.16x,ADDRB5:]]
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR6:]] [[#%.16x,ADDRB6:]]
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR7:]] [[#%.16x,ADDRB7:]]
 ; BOLT-NEXT: 00000090 <End of list>
-; BOLT-NEXT: 00000110
-; BOLT-NEXT: 00000110
-; BOLT-NEXT: 00000110
-; BOLT-NEXT: 00000110 <End of list>
-; BOLT-NEXT: 00000150
-; BOLT-NEXT: 00000150
-; BOLT-NEXT: 00000150
-; BOLT-NEXT: 00000150 <End of list>
-; BOLT-NEXT: 00000190 [[#%.16x,ADDR8:]] [[#%.16x,ADDRB8:]]
-; BOLT-NEXT: 00000190 [[#%.16x,ADDR9:]] [[#%.16x,ADDRB9:]]
-; BOLT-NEXT: 00000190 [[#%.16x,ADDR10:]] [[#%.16x,ADDRB10:]]
-; BOLT-NEXT: 00000190 [[#%.16x,ADDR11:]] [[#%.16x,ADDRB11:]]
-; BOLT-NEXT: 00000190 [[#%.16x,ADDR12:]] [[#%.16x,ADDRB12:]]
-; BOLT-NEXT: 00000190 [[#%.16x,ADDR13:]] [[#%.16x,ADDRB13:]]
-; BOLT-NEXT: 00000190 [[#%.16x,ADDR14:]] [[#%.16x,ADDRB14:]]
-; BOLT-NEXT: 00000190 <End of list>
+; BOLT-NEXT: 000000a0 [[#%.16x,ADDR1:]] [[#%.16x,ADDRB1:]]
+; BOLT-NEXT: 000000a0 [[#%.16x,ADDR2:]] [[#%.16x,ADDRB2:]]
+; BOLT-NEXT: 000000a0 [[#%.16x,ADDR3:]] [[#%.16x,ADDRB3:]]
+; BOLT-NEXT: 000000a0 <End of list>
+; BOLT-NEXT: 000000e0 [[#%.16x,ADDR5:]] [[#%.16x,ADDRB5:]]
+; BOLT-NEXT: 000000e0 [[#%.16x,ADDR6:]] [[#%.16x,ADDRB6:]]
+; BOLT-NEXT: 000000e0 [[#%.16x,ADDR7:]] [[#%.16x,ADDRB7:]]
+; BOLT-NEXT: 000000e0 <End of list>
+; BOLT-NEXT: 00000120 [[#%.16x,ADDR8:]] [[#%.16x,ADDRB8:]]
+; BOLT-NEXT: 00000120 [[#%.16x,ADDR9:]] [[#%.16x,ADDRB9:]]
+; BOLT-NEXT: 00000120 [[#%.16x,ADDR10:]] [[#%.16x,ADDRB10:]]
+; BOLT-NEXT: 00000120 [[#%.16x,ADDR11:]] [[#%.16x,ADDRB11:]]
+; BOLT-NEXT: 00000120 [[#%.16x,ADDR12:]] [[#%.16x,ADDRB12:]]
+; BOLT-NEXT: 00000120 [[#%.16x,ADDR13:]] [[#%.16x,ADDRB13:]]
+; BOLT-NEXT: 00000120 [[#%.16x,ADDR14:]] [[#%.16x,ADDRB14:]]
+; BOLT-NEXT: 00000120 <End of list>
+; BOLT-NEXT: 000001a0 <End of list>
+; BOLT-NEXT: 000001b0 [[#%.16x,ADDR8:]] [[#%.16x,ADDRB8:]]
+; BOLT-NEXT: 000001b0 [[#%.16x,ADDR9:]] [[#%.16x,ADDRB9:]]
+; BOLT-NEXT: 000001b0 [[#%.16x,ADDR10:]] [[#%.16x,ADDRB10:]]
+; BOLT-NEXT: 000001b0 <End of list>
+; BOLT-NEXT: 000001f0 [[#%.16x,ADDR12:]] [[#%.16x,ADDRB12:]]
+; BOLT-NEXT: 000001f0 [[#%.16x,ADDR13:]] [[#%.16x,ADDRB13:]]
+; BOLT-NEXT: 000001f0 [[#%.16x,ADDR14:]] [[#%.16x,ADDRB14:]]
+; BOLT-NEXT: 000001f0 <End of list>
 
 ; BOLT: DW_TAG_compile_unit
 ; BOLT: DW_AT_GNU_dwo_name [DW_FORM_strp] ( .debug_str[0x{{[0-9a-fA-F]+}}] = "main.dwo.dwo")
 ; BOLT-NEXT: DW_AT_GNU_dwo_id
-; BOLT-NEXT: DW_AT_GNU_ranges_base [DW_FORM_sec_offset]  (0x00000010)
+; BOLT-NEXT: DW_AT_GNU_ranges_base [DW_FORM_sec_offset]  (0x00000090)
 ; BOLT-NEXT: DW_AT_low_pc [DW_FORM_addr] (0x0000000000000000)
-; BOLT-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000090
+; BOLT-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000010
 ; BOLT-NEXT: [0x[[#ADDR1]], 0x[[#ADDRB1]])
 ; BOLT-NEXT: [0x[[#ADDR2]], 0x[[#ADDRB2]])
 ; BOLT-NEXT: [0x[[#ADDR3]], 0x[[#ADDRB3]])
@@ -64,13 +66,14 @@
 ; BOLT-NEXT: [0x[[#ADDR6]], 0x[[#ADDRB6]])
 ; BOLT-NEXT: [0x[[#ADDR7]], 0x[[#ADDRB7]])
 ; BOLT-NEXT: DW_AT_GNU_addr_base [DW_FORM_sec_offset]  (0x00000000)
+; BOLT-NEXT: Compile Unit
 
 ; BOLT: DW_TAG_compile_unit
 ; BOLT: DW_AT_GNU_dwo_name [DW_FORM_strp] ( .debug_str[0x{{[0-9a-fA-F]+}}] = "mainOther.dwo.dwo")
 ; BOLT-NEXT: DW_AT_GNU_dwo_id
-; BOLT-NEXT: DW_AT_GNU_ranges_base [DW_FORM_sec_offset]  (0x00000110)
+; BOLT-NEXT: DW_AT_GNU_ranges_base [DW_FORM_sec_offset]  (0x000001a0)
 ; BOLT-NEXT: DW_AT_low_pc [DW_FORM_addr] (0x0000000000000000)
-; BOLT-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000190
+; BOLT-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000120
 ; BOLT-NEXT: [0x[[#ADDR8]], 0x[[#ADDRB8]])
 ; BOLT-NEXT: [0x[[#ADDR9]], 0x[[#ADDRB9]])
 ; BOLT-NEXT: [0x[[#ADDR10]], 0x[[#ADDRB10]])
@@ -79,19 +82,20 @@
 ; BOLT-NEXT: [0x[[#ADDR13]], 0x[[#ADDRB13]])
 ; BOLT-NEXT: [0x[[#ADDR14]], 0x[[#ADDRB14]])
 ; BOLT-NEXT: DW_AT_GNU_addr_base [DW_FORM_sec_offset]  (0x00000018)
+; BOLT: {{^$}}
 
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
-; BOLT-DWO-MAIN-NEXT:   DW_AT_ranges [DW_FORM_sec_offset] (0x00000000
+; BOLT-DWO-MAIN-NEXT:   DW_AT_ranges [DW_FORM_sec_offset] (0x00000010
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
-; BOLT-DWO-MAIN-NEXT:   DW_AT_ranges [DW_FORM_sec_offset] (0x00000040
+; BOLT-DWO-MAIN-NEXT:   DW_AT_ranges [DW_FORM_sec_offset] (0x00000050
 
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
-; BOLT-DWO-MAIN-NEXT:   DW_AT_ranges [DW_FORM_sec_offset] (0x00000000
+; BOLT-DWO-MAIN-NEXT:   DW_AT_ranges [DW_FORM_sec_offset] (0x00000010
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
-; BOLT-DWO-MAIN-NEXT:   DW_AT_ranges [DW_FORM_sec_offset] (0x00000040
+; BOLT-DWO-MAIN-NEXT:   DW_AT_ranges [DW_FORM_sec_offset] (0x00000050
diff --git a/bolt/test/X86/dwarf4-df-input-lowpc-ranges.test b/bolt/test/X86/dwarf4-df-input-lowpc-ranges.test
index 276bea4ba0c1c..ab4353a282475 100644
--- a/bolt/test/X86/dwarf4-df-input-lowpc-ranges.test
+++ b/bolt/test/X86/dwarf4-df-input-lowpc-ranges.test
@@ -15,29 +15,30 @@
 
 ; BOLT: .debug_ranges
 ; BOLT-NEXT: 00000000 <End of list>
-; BOLT-NEXT: 00000010
-; BOLT-NEXT: 00000010
-; BOLT-NEXT: 00000010
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR1:]] [[#%.16x,ADDRB1:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR2:]] [[#%.16x,ADDRB2:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR3:]] [[#%.16x,ADDRB3:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR4:]] [[#%.16x,ADDRB4:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR5:]] [[#%.16x,ADDRB5:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR6:]] [[#%.16x,ADDRB6:]]
+; BOLT-NEXT: 00000010 [[#%.16x,ADDR7:]] [[#%.16x,ADDRB7:]]
 ; BOLT-NEXT: 00000010 <End of list>
-; BOLT-NEXT: 00000050
-; BOLT-NEXT: 00000050
-; BOLT-NEXT: 00000050
-; BOLT-NEXT: 00000050 <End of list>
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR1:]] [[#%.16x,ADDRB1:]]
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR2:]] [[#%.16x,ADDRB2:]]
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR3:]] [[#%.16x,ADDRB3:]]
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR4:]] [[#%.16x,ADDRB4:]]
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR5:]] [[#%.16x,ADDRB5:]]
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR6:]] [[#%.16x,ADDRB6:]]
-; BOLT-NEXT: 00000090 [[#%.16x,ADDR7:]] [[#%.16x,ADDRB7:]]
 ; BOLT-NEXT: 00000090 <End of list>
+; BOLT-NEXT: 000000a0 [[#%.16x,ADDR1:]] [[#%.16x,ADDRB1:]]
+; BOLT-NEXT: 000000a0 [[#%.16x,ADDR2:]] [[#%.16x,ADDRB2:]]
+; BOLT-NEXT: 000000a0 [[#%.16x,ADDR3:]] [[#%.16x,ADDRB3:]]
+; BOLT-NEXT: 000000a0 <End of list>
+; BOLT-NEXT: 000000e0 [[#%.16x,ADDR5:]] [[#%.16x,ADDRB5:]]
+; BOLT-NEXT: 000000e0 [[#%.16x,ADDR6:]] [[#%.16x,ADDRB6:]]
+; BOLT-NEXT: 000000e0 [[#%.16x,ADDR7:]] [[#%.16x,ADDRB7:]]
+; BOLT-NEXT: 000000e0 <End of list>
 
 ; BOLT: DW_TAG_compile_unit
 ; BOLT: DW_AT_GNU_dwo_name [DW_FORM_strp] ( .debug_str[0x{{[0-9a-fA-F]+}}] = "main.dwo.dwo")
 ; BOLT-NEXT: DW_AT_GNU_dwo_id
-; BOLT-NEXT: DW_AT_GNU_ranges_base [DW_FORM_sec_offset]  (0x00000010)
+; BOLT-NEXT: DW_AT_GNU_ranges_base [DW_FORM_sec_offset]  (0x00000090)
 ; BOLT-NEXT: DW_AT_low_pc [DW_FORM_addr] (0x0000000000000000)
-; BOLT-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000090
+; BOLT-NEXT: DW_AT_ranges [DW_FORM_sec_offset] (0x00000010
 ; BOLT-NEXT: [0x[[#ADDR1]], 0x[[#ADDRB1]])
 ; BOLT-NEXT: [0x[[#ADDR2]], 0x[[#ADDRB2]])
 ; BOLT-NEXT: [0x[[#ADDR3]], 0x[[#ADDRB3]])
@@ -48,9 +49,9 @@
 ; BOLT-NEXT: DW_AT_GNU_addr_base [DW_FORM_sec_offset]  (0x00000000)
 
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
-; BOLT-DWO-MAIN-NEXT:   DW_AT_ranges [DW_FORM_sec_offset] (0x00000000
+; BOLT-DWO-MAIN-NEXT:   DW_AT_ranges [DW_FORM_sec_offset] (0x00000010
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
 ; BOLT-DWO-MAIN:        DW_TAG_subprogram
-; BOLT-DWO-MAIN-NEXT:   DW_AT_ranges [DW_FORM_sec_offset] (0x00000040
+; BOLT-DWO-MAIN-NEXT:   DW_AT_ranges [DW_FORM_sec_offset] (0x00000050

>From 2ed77d25ff926d09b5cef752de7fd7eae12213a1 Mon Sep 17 00:00:00 2001
From: Sayhaan Siddiqui <sayhaan at meta.com>
Date: Mon, 8 Jul 2024 10:26:50 -0700
Subject: [PATCH 2/2] Refactor address writers

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Differential Revision: https://phabricator.intern.facebook.com/D59493903
---
 bolt/include/bolt/Core/DebugData.h        | 53 ++++++++++++-----
 bolt/include/bolt/Rewrite/DWARFRewriter.h |  5 ++
 bolt/lib/Core/DebugData.cpp               | 71 +++++++++++------------
 bolt/lib/Rewrite/DWARFRewriter.cpp        | 63 ++++++++++++--------
 4 files changed, 118 insertions(+), 74 deletions(-)

diff --git a/bolt/include/bolt/Core/DebugData.h b/bolt/include/bolt/Core/DebugData.h
index 144433ac78a37..05217c699f0dd 100644
--- a/bolt/include/bolt/Core/DebugData.h
+++ b/bolt/include/bolt/Core/DebugData.h
@@ -344,14 +344,34 @@ class DebugAddrWriter {
   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 void update();
 
   /// Return buffer with all the entries in .debug_addr already writen out using
   /// update(...).
   virtual AddressSectionBuffer &finalize() { return *Buffer; }
 
   /// Returns False if .debug_addr section was created..
-  bool isInitialized() const { return !AddressMaps.empty(); }
+  bool isInitialized() { return Map.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 buffer for the given
+  /// CU.
+  void appendToAddressBuffer(const AddressSectionBuffer &Buffer) {
+    *AddressStream << Buffer;
+  }
+
+  /// Sets AddressByteSize for the CU.
+  void setAddressByteSize(uint8_t AddressByteSize) {
+    this->AddressByteSize = AddressByteSize;
+  }
+
+  /// Sets AddrOffsetSectionBase for the CU.
+  void setAddrOffsetSectionBase(std::optional<uint64_t> AddrOffsetSectionBase) {
+    this->AddrOffsetSectionBase = AddrOffsetSectionBase;
+  }
 
 protected:
   class AddressForDWOCU {
@@ -396,6 +416,8 @@ class DebugAddrWriter {
 
     void dump();
 
+    bool empty = false;
+
   private:
     AddressToIndexMap AddressToIndex;
     IndexToAddressMap IndexToAddress;
@@ -408,14 +430,16 @@ 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;
+  std::optional<uint64_t> AddrOffsetSectionBase;
   /// 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 {
@@ -424,7 +448,10 @@ class DebugAddrWriterDwarf5 : public DebugAddrWriter {
   DebugAddrWriterDwarf5(BinaryContext *BC) : DebugAddrWriter(BC) {}
 
   /// Write out entries in to .debug_addr section for CUs.
-  virtual void update(DIEBuilder &DIEBlder, DWARFUnit &CUs) override;
+  virtual void update() override;
+
+  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
@@ -584,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);
@@ -601,7 +626,7 @@ class DebugLoclistWriter : public DebugLocWriter {
     }
   }
 
-  static void setAddressWriter(DebugAddrWriter *AddrW) { AddrWriter = AddrW; }
+  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,
@@ -631,7 +656,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.
diff --git a/bolt/include/bolt/Rewrite/DWARFRewriter.h b/bolt/include/bolt/Rewrite/DWARFRewriter.h
index 1776f9efd137f..a24ef2d4aeff8 100644
--- a/bolt/include/bolt/Rewrite/DWARFRewriter.h
+++ b/bolt/include/bolt/Rewrite/DWARFRewriter.h
@@ -93,6 +93,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 +119,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.
diff --git a/bolt/lib/Core/DebugData.cpp b/bolt/lib/Core/DebugData.cpp
index 08d4c45aac791..52df5d5e55707 100644
--- a/bolt/lib/Core/DebugData.cpp
+++ b/bolt/lib/Core/DebugData.cpp
@@ -393,6 +393,7 @@ void DebugARangesSectionWriter::writeARangesSection(
 DebugAddrWriter::DebugAddrWriter(BinaryContext *BC) : BC(BC) {
   Buffer = std::make_unique<AddressSectionBuffer>();
   AddressStream = std::make_unique<raw_svector_ostream>(*Buffer);
+  Map = AddressForDWOCU();
 }
 
 void DebugAddrWriter::AddressForDWOCU::dump() {
@@ -405,11 +406,8 @@ 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];
+  if (Map.begin() == Map.end())
+    Map.empty = true;
   auto Entry = Map.find(Address);
   if (Entry == Map.end()) {
     auto Index = Map.getNextIndex();
@@ -449,25 +447,20 @@ 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())
+void DebugAddrWriter::updateAddrBase(DIEBuilder &DIEBlder, DWARFUnit &CU,
+                                     const uint64_t Offset) {
+  updateAddressBase(DIEBlder, *this, CU, Offset);
+}
+
+void DebugAddrWriter::update() {
+  if (Map.indexToAddressBegin() == Map.indexToAdddessEnd())
     return;
-  std::vector<IndexAddressPair> SortedMap(AM->second.indexToAddressBegin(),
-                                          AM->second.indexToAdddessEnd());
+  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();
+  uint8_t AddrSize = AddressByteSize;
   uint32_t Counter = 0;
   auto WriteAddress = [&](uint64_t Address) -> void {
     ++Counter;
@@ -490,10 +483,23 @@ void DebugAddrWriter::update(DIEBuilder &DIEBlder, DWARFUnit &CU) {
       WriteAddress(0);
     WriteAddress(Val.second);
   }
-  updateAddressBase(DIEBlder, *this, CU, Offset);
 }
 
-void DebugAddrWriterDwarf5::update(DIEBuilder &DIEBlder, DWARFUnit &CU) {
+void DebugAddrWriterDwarf5::updateAddrBase(DIEBuilder &DIEBlder, DWARFUnit &CU,
+                                           const uint64_t Offset) {
+  /// Doesn't update address base if the CU doesn't access .debug_addr.
+  if (Map.indexToAddressBegin() == Map.indexToAdddessEnd()) {
+    std::optional<uint64_t> BaseOffset = CU.getAddrOffsetSectionBase();
+    if (!BaseOffset)
+      return;
+  }
+  /// Header for DWARF5 has size 8, so we add it to the offset.
+  updateAddressBase(DIEBlder, *this, CU, Offset + 8);
+}
+
+DenseMap<uint64_t, uint64_t> DebugAddrWriter::UnmodifiedAddressOffsets;
+
+void DebugAddrWriterDwarf5::update() {
   // Need to layout all sections within .debug_addr
   // Within each section sort Address by index.
   const endianness Endian = BC->DwCtx->isLittleEndian()
@@ -505,14 +511,11 @@ void DebugAddrWriterDwarf5::update(DIEBuilder &DIEBlder, DWARFUnit &CU) {
   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);
+  const uint8_t AddrSize = AddressByteSize;
   // 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 (Map.indexToAddressBegin() == Map.indexToAdddessEnd()) {
+    std::optional<uint64_t> BaseOffset = AddrOffsetSectionBase;
     if (!BaseOffset)
       return;
     // Address base offset is to the first entry.
@@ -520,7 +523,6 @@ void DebugAddrWriterDwarf5::update(DIEBuilder &DIEBlder, DWARFUnit &CU) {
     uint64_t Offset = *BaseOffset - HeaderSize;
     auto Iter = UnmodifiedAddressOffsets.find(Offset);
     if (Iter != UnmodifiedAddressOffsets.end()) {
-      updateAddressBase(DIEBlder, *this, CU, Iter->getSecond());
       return;
     }
     UnmodifiedAddressOffsets[Offset] = Buffer->size() + HeaderSize;
@@ -529,16 +531,13 @@ void DebugAddrWriterDwarf5::update(DIEBuilder &DIEBlder, DWARFUnit &CU) {
       DumpOpts.RecoverableErrorHandler(std::move(Err));
       return;
     }
-
     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
@@ -789,8 +788,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) {
diff --git a/bolt/lib/Rewrite/DWARFRewriter.cpp b/bolt/lib/Rewrite/DWARFRewriter.cpp
index be7bab9359839..cec98c1d7304b 100644
--- a/bolt/lib/Rewrite/DWARFRewriter.cpp
+++ b/bolt/lib/Rewrite/DWARFRewriter.cpp
@@ -615,7 +615,6 @@ void DWARFRewriter::updateDebugInfo() {
   if (BC.isDWARF5Used()) {
     AddrWriter = std::make_unique<DebugAddrWriterDwarf5>(&BC);
     RangeListsSectionWriter = std::make_unique<DebugRangeListsSectionWriter>();
-    DebugRangeListsSectionWriter::setAddressWriter(AddrWriter.get());
   } else {
     AddrWriter = std::make_unique<DebugAddrWriter>(&BC);
   }
@@ -623,8 +622,6 @@ void DWARFRewriter::updateDebugInfo() {
   if (BC.isDWARFLegacyUsed())
     LegacyRangesSectionWriter = std::make_unique<DebugRangesSectionWriter>();
 
-  DebugLoclistWriter::setAddressWriter(AddrWriter.get());
-
   uint32_t CUIndex = 0;
   std::mutex AccessMutex;
   // Needs to be invoked in the same order as CUs are processed.
@@ -632,8 +629,16 @@ void DWARFRewriter::updateDebugInfo() {
     std::lock_guard<std::mutex> Lock(AccessMutex);
     const uint16_t DwarfVersion = CU.getVersion();
     if (DwarfVersion >= 5) {
-      LocListWritersByCU[CUIndex] =
-          std::make_unique<DebugLoclistWriter>(CU, DwarfVersion, false);
+      auto AddrW = std::make_unique<DebugAddrWriterDwarf5>(&BC);
+      AddressWritersByCU[CU.getOffset()] = std::move(AddrW);
+      DebugRangeListsSectionWriter::setAddressWriter(
+          AddressWritersByCU[CU.getOffset()].get());
+      AddressWritersByCU[CU.getOffset()]->setAddressByteSize(
+          CU.getAddressByteSize());
+      AddressWritersByCU[CU.getOffset()]->setAddrOffsetSectionBase(
+          CU.getAddrOffsetSectionBase());
+      LocListWritersByCU[CUIndex] = std::make_unique<DebugLoclistWriter>(
+          CU, DwarfVersion, false, AddressWritersByCU[CU.getOffset()].get());
 
       if (std::optional<uint64_t> DWOId = CU.getDWOId()) {
         assert(RangeListsWritersByCU.count(*DWOId) == 0 &&
@@ -645,6 +650,10 @@ void DWARFRewriter::updateDebugInfo() {
       }
 
     } else {
+      auto AddrW = std::make_unique<DebugAddrWriter>(&BC);
+      AddressWritersByCU[CU.getOffset()] = std::move(AddrW);
+      AddressWritersByCU[CU.getOffset()]->setAddressByteSize(
+          CU.getAddressByteSize());
       LocListWritersByCU[CUIndex] = std::make_unique<DebugLocWriter>();
       if (std::optional<uint64_t> DWOId = CU.getDWOId()) {
         assert(LegacyRangesWritersByCU.count(*DWOId) == 0 &&
@@ -676,6 +685,7 @@ void DWARFRewriter::updateDebugInfo() {
     DebugRangesSectionWriter *RangesSectionWriter =
         Unit->getVersion() >= 5 ? RangeListsSectionWriter.get()
                                 : LegacyRangesSectionWriter.get();
+    auto &AddressWriter = AddressWritersByCU[Unit->getOffset()];
     // Skipping CUs that failed to load.
     if (SplitCU) {
       DIEBuilder DWODIEBuilder(BC, &(*SplitCU)->getContext(), DebugNamesTable,
@@ -696,7 +706,8 @@ void DWARFRewriter::updateDebugInfo() {
       DWODIEBuilder.updateDWONameCompDirForTypes(DWOStrOffstsWriter,
                                                  DWOStrWriter, **SplitCU,
                                                  DwarfOutputPath, DWOName);
-      DebugLoclistWriter DebugLocDWoWriter(*Unit, Unit->getVersion(), true);
+      DebugLoclistWriter DebugLocDWoWriter(*Unit, Unit->getVersion(), true,
+                                           AddressWriter.get());
       DebugRangesSectionWriter *TempRangesSectionWriter = RangesSectionWriter;
       if (Unit->getVersion() >= 5) {
         TempRangesSectionWriter = RangeListsWritersByCU[*DWOId].get();
@@ -707,7 +718,7 @@ void DWARFRewriter::updateDebugInfo() {
       }
 
       updateUnitDebugInfo(*(*SplitCU), DWODIEBuilder, DebugLocDWoWriter,
-                          *TempRangesSectionWriter);
+                          *TempRangesSectionWriter, *AddressWriter.get());
       DebugLocDWoWriter.finalize(DWODIEBuilder,
                                  *DWODIEBuilder.getUnitDIEbyUnit(**SplitCU));
       if (Unit->getVersion() >= 5)
@@ -725,11 +736,10 @@ void DWARFRewriter::updateDebugInfo() {
     }
 
     updateUnitDebugInfo(*Unit, *DIEBlder, *DebugLocWriter, *RangesSectionWriter,
-                        RangesBase);
+                        *AddressWriter.get(), RangesBase);
     DebugLocWriter->finalize(*DIEBlder, *DIEBlder->getUnitDIEbyUnit(*Unit));
     if (Unit->getVersion() >= 5)
       RangesSectionWriter->finalizeSection();
-    AddrWriter->update(*DIEBlder, *Unit);
   };
 
   DIEBuilder DIEBlder(BC, BC.DwCtx.get(), DebugNamesTable);
@@ -777,7 +787,7 @@ void DWARFRewriter::updateDebugInfo() {
 void DWARFRewriter::updateUnitDebugInfo(
     DWARFUnit &Unit, DIEBuilder &DIEBldr, DebugLocWriter &DebugLocWriter,
     DebugRangesSectionWriter &RangesSectionWriter,
-    std::optional<uint64_t> RangesBase) {
+    DebugAddrWriter &AddressWriter, std::optional<uint64_t> RangesBase) {
   // Cache debug ranges so that the offset for identical ranges could be reused.
   std::map<DebugAddressRangesVector, uint64_t> CachedRanges;
 
@@ -811,7 +821,7 @@ void DWARFRewriter::updateUnitDebugInfo(
 
     if (FormLowPC == dwarf::DW_FORM_addrx ||
         FormLowPC == dwarf::DW_FORM_GNU_addr_index)
-      LowPC = AddrWriter->getIndexFromAddress(LowPC, Unit);
+      LowPC = AddressWriter.getIndexFromAddress(LowPC, Unit);
 
     if (LowPCVal)
       DIEBldr.replaceValue(Die, AttrLowPC, FormLowPC, DIEInteger(LowPC));
@@ -975,7 +985,7 @@ void DWARFRewriter::updateUnitDebugInfo(
 
         if (AttrVal.getForm() == dwarf::DW_FORM_addrx) {
           const uint32_t Index =
-              AddrWriter->getIndexFromAddress(UpdatedAddress, Unit);
+              AddressWriter.getIndexFromAddress(UpdatedAddress, Unit);
           DIEBldr.replaceValue(Die, AttrVal.getAttribute(), AttrVal.getForm(),
                                DIEInteger(Index));
         } else if (AttrVal.getForm() == dwarf::DW_FORM_addr) {
@@ -1192,7 +1202,7 @@ void DWARFRewriter::updateUnitDebugInfo(
                 assert(EntryAddress && "Address is not found.");
                 assert(Index <= std::numeric_limits<uint32_t>::max() &&
                        "Invalid Operand Index.");
-                const uint32_t AddrIndex = AddrWriter->getIndexFromAddress(
+                const uint32_t AddrIndex = AddressWriter.getIndexFromAddress(
                     EntryAddress->Address, Unit);
                 // update Index into .debug_address section for DW_AT_location.
                 // The Size field is not stored in IR, we need to minus 1 in
@@ -1244,7 +1254,7 @@ void DWARFRewriter::updateUnitDebugInfo(
           std::lock_guard<std::mutex> Lock(DWARFRewriterMutex);
           if (Form == dwarf::DW_FORM_addrx ||
               Form == dwarf::DW_FORM_GNU_addr_index) {
-            const uint32_t Index = AddrWriter->getIndexFromAddress(
+            const uint32_t Index = AddressWriter.getIndexFromAddress(
                 NewAddress ? NewAddress : Address, Unit);
             DIEBldr.replaceValue(Die, LowPCAttrInfo.getAttribute(),
                                  LowPCAttrInfo.getForm(), DIEInteger(Index));
@@ -1561,14 +1571,10 @@ void DWARFRewriter::finalizeDebugSections(
           LocationListSectionContents->size());
   }
 
-  // AddrWriter should be finalized after debug_loc since more addresses can be
-  // added there.
-  if (AddrWriter->isInitialized()) {
-    AddressSectionBuffer AddressSectionContents = AddrWriter->finalize();
-    BC.registerOrUpdateNoteSection(".debug_addr",
-                                   copyByteArray(AddressSectionContents),
-                                   AddressSectionContents.size());
-  }
+  AddressSectionBuffer AddressSectionContents = AddrWriter->finalize();
+  BC.registerOrUpdateNoteSection(".debug_addr",
+                                 copyByteArray(AddressSectionContents),
+                                 AddressSectionContents.size());
 
   Streamer.emitAbbrevs(DIEBlder.getAbbrevs(), BC.DwCtx->getMaxVersion());
   Streamer.finish();
@@ -1622,6 +1628,16 @@ void DWARFRewriter::finalizeCompileUnits(DIEBuilder &DIEBlder,
                                          CUOffsetMap &CUMap,
                                          const std::list<DWARFUnit *> &CUs) {
   for (DWARFUnit *CU : CUs) {
+    auto AddressWriterIterator = AddressWritersByCU.find(CU->getOffset());
+    assert(AddressWriterIterator != AddressWritersByCU.end() &&
+           "AddressWriter does not exist for CU");
+    auto &AddressWriter = AddressWriterIterator->second;
+    AddressWriter->update();
+    AddressWriter->updateAddrBase(DIEBlder, *CU, AddrWriter->finalize().size());
+    if (AddressWriter->isInitialized()) {
+      AddressSectionBuffer AddressSectionContents = AddressWriter->finalize();
+      AddrWriter->appendToAddressBuffer(AddressSectionContents);
+    }
     if (CU->getVersion() != 4)
       continue;
     std::optional<uint64_t> DWOId = CU->getDWOId();
@@ -2323,7 +2339,8 @@ void DWARFRewriter::convertToRangesPatchDebugInfo(
   // when it's absent.
   if (IsUnitDie) {
     if (LowForm == dwarf::DW_FORM_addrx) {
-      const uint32_t Index = AddrWriter->getIndexFromAddress(0, Unit);
+      const uint32_t Index =
+          AddressWritersByCU[Unit.getOffset()]->getIndexFromAddress(0, Unit);
       DIEBldr.replaceValue(&Die, LowPCAttrInfo.getAttribute(),
                            LowPCAttrInfo.getForm(), DIEInteger(Index));
     } else {



More information about the llvm-commits mailing list