[BOLT][DWARF] Output DWO files as they are being processed

[ayermolo]

Summary

Changed how we handle writing out .dwo and .dwp files. We now write out DWO
sections sooner and destroy DIEBuilder. This should decrease memory footprint.

Ran on clang-17 build in debug mode with split-dwarf.
before
8:07.49 real, 664.62 user, 69.00 sys, 0 amem, 41601612 mmem
8:07.06 real, 669.60 user, 68.75 sys, 0 amem, 41822588 mmem
8:00.36 real, 664.14 user, 66.36 sys, 0 amem, 41561548 mmem

after
8:21.85 real, 682.23 user, 69.64 sys, 0 amem, 39379880 mmem
8:04.58 real, 671.62 user, 66.50 sys, 0 amem, 39735800 mmem
8:10.02 real, 680.67 user, 67.24 sys, 0 amem, 39662888 mmem

Depends on D151906

[ayermolo]

rebase

[ayermolo]

rebase

[ayermolo]

rebase

[ayermolo]

rebase

[ayermolo]

rebase

[ayermolo]

rebase

[maksfb]

bolt/lib/Rewrite/DWARFRewriter.cpp

306

Do we ever invoke this function with a null NameToIndexMap?

361

1523

Assert that RangeListsWriter is non-null, or use a reference.

1604

1605

Do we need to assert that the section is non-null?

1615

1996

nit: remove curly braces.

[ayermolo]

bolt/lib/Rewrite/DWARFRewriter.cpp

306

Not anymore, it looks like.

1605

I guess in theory there might not be a str section somehow? Added just in case.

[ayermolo]

addressed comments

[ayermolo]

rebase

[maksfb]

LGTM

[ayermolo]

rebase

bolt/include/bolt/Rewrite/DWARFRewriter.h

 //===- bolt/Rewrite/DWARFRewriter.h -----------------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef BOLT_REWRITE_DWARF_REWRITER_H
 #define BOLT_REWRITE_DWARF_REWRITER_H
 
 #include "bolt/Core/DIEBuilder.h"
 #include "bolt/Core/DebugData.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/CodeGen/DIE.h"
+#include "llvm/DWP/DWP.h"
 #include "llvm/MC/MCAsmLayout.h"
+#include "llvm/MC/MCContext.h"
+#include "llvm/Support/ToolOutputFile.h"
 #include <cstdint>
 #include <memory>
 #include <mutex>
 #include <optional>
 #include <set>
 #include <unordered_map>
 #include <unordered_set>
 #include <vector>
struct GDBIndexTUEntry {
     uint64_t TypeDIERelativeOffset;
   };
   /// Contains information for CU or TU so we can output correct {cu, tu}-index.
   struct UnitMeta {
     uint64_t Offset;
     uint64_t Length;
     uint64_t TUHash;
   };
+  using UnitMetaVectorType = std::vector<UnitMeta>;
 
 private:
   BinaryContext &BC;
 
   std::mutex DWARFRewriterMutex;
 
   /// Stores and serializes information that will be put into the
   /// .debug_ranges DWARF section.
private:
   using RangeListsDWOWriers =
       std::unordered_map<uint64_t,
                          std::unique_ptr<DebugRangeListsSectionWriter>>;
   /// Store Rangelists writer for each DWO CU.
   RangeListsDWOWriers RangeListsWritersByCU;
 
   std::mutex LocListDebugInfoPatchesMutex;
 
-  /// Dwo id specific its .debug_info.dwo section content.
-  std::unordered_map<uint64_t, std::string> DwoDebugInfoMap;
-
-  /// Dwo id specific its .debug_abbrev.dwo section content.
-  std::unordered_map<uint64_t, std::string> DwoDebugAbbrevMap;
-
-  /// Dwo id specific its .debug_types.dwo section content.
-  std::unordered_map<uint64_t, std::string> DwoDebugTypeMap;
-
   /// Dwo id specific its RangesBase.
   std::unordered_map<uint64_t, uint64_t> DwoRangesBase;
 
   std::unordered_map<DWARFUnit *, uint64_t> LineTablePatchMap;
   std::unordered_map<DWARFUnit *, uint64_t> TypeUnitRelocMap;
 
   /// Entries for GDB Index Types CU List
   using GDBIndexTUEntryType = std::vector<GDBIndexTUEntry>;
   GDBIndexTUEntryType GDBIndexTUEntryVector;
 
-  using UnitMetaVectorType = std::vector<UnitMeta>;
-  using TUnitMetaDwoMapType = std::unordered_map<uint64_t, UnitMetaVectorType>;
-  using CUnitMetaDwoMapType = std::unordered_map<uint64_t, UnitMeta>;
-  CUnitMetaDwoMapType CUnitMetaDwoMap;
-  TUnitMetaDwoMapType TUnitMetaDwoMap;
-
   /// DWARFLegacy is all DWARF versions before DWARF 5.
   enum class DWARFVersion { DWARFLegacy, DWARF5 };
 
   /// Update debug info for all DIEs in \p Unit.
   void updateUnitDebugInfo(DWARFUnit &Unit, DIEBuilder &DIEBldr,
                            DebugLocWriter &DebugLocWriter,
                            DebugRangesSectionWriter &RangesSectionWriter,
                            std::optional<uint64_t> RangesBase = std::nullopt);
private:
 
   /// Patches the binary for DWARF address ranges (e.g. in functions and lexical
   /// blocks) to be updated.
   void updateDebugAddressRanges();
 
   /// Rewrite .gdb_index section if present.
   void updateGdbIndexSection(CUOffsetMap &CUMap, uint32_t NumCUs);
 
-  /// Output .dwo files.
-  void writeDWOFiles(std::unordered_map<uint64_t, std::string> &DWOIdToName);
-
-  /// Output .dwp files.
-  void writeDWP(std::unordered_map<uint64_t, std::string> &DWOIdToName);
-
   /// DWARFDie contains a pointer to a DIE and hence gets invalidated once the
   /// embedded DIE is destroyed. This wrapper class stores a DIE internally and
   /// could be cast to a DWARFDie that is valid even after the initial DIE is
   /// destroyed.
   struct DWARFDieWrapper {
     DWARFUnit *Unit;
     DWARFDebugInfoEntry DIE;
 
public:
   DWARFRewriter(BinaryContext &BC) : BC(BC) {}
 
   /// Main function for updating the DWARF debug info.
   void updateDebugInfo();
 
   /// Update stmt_list for CUs based on the new .debug_line \p Layout.
   void updateLineTableOffsets(const MCAsmLayout &Layout);
 
-  /// Given a \p DWOId, return its DebugLocWriter if it exists.
-  DebugLocWriter *getDebugLocWriter(uint64_t DWOId) {
-    auto Iter = LocListWritersByCU.find(DWOId);
-    return Iter == LocListWritersByCU.end() ? nullptr
-                                            : LocListWritersByCU[DWOId].get();
-  }
-
-  StringRef getDwoDebugInfoStr(uint64_t DWOId) {
-    return DwoDebugInfoMap[DWOId];
-  }
-
-  StringRef getDwoDebugAbbrevStr(uint64_t DWOId) {
-    return DwoDebugAbbrevMap[DWOId];
-  }
-
-  StringRef getDwoDebugTypeStr(uint64_t DWOId) {
-    return DwoDebugTypeMap[DWOId];
-  }
-
   uint64_t getDwoRangesBase(uint64_t DWOId) { return DwoRangesBase[DWOId]; }
 
-  void setDwoDebugInfoStr(uint64_t DWOId, StringRef Str) {
-    DwoDebugInfoMap[DWOId] = Str.str();
-  }
-
-  void setDwoDebugAbbrevStr(uint64_t DWOId, StringRef Str) {
-    DwoDebugAbbrevMap[DWOId] = Str.str();
-  }
-
-  void setDwoDebugTypeStr(uint64_t DWOId, StringRef Str) {
-    DwoDebugTypeMap[DWOId] = Str.str();
-  }
-
   void setDwoRangesBase(uint64_t DWOId, uint64_t RangesBase) {
     DwoRangesBase[DWOId] = RangesBase;
   }
 
   /// Adds an GDBIndexTUEntry if .gdb_index seciton exists.
   void addGDBTypeUnitEntry(const GDBIndexTUEntry &&Entry);
 
   /// Returns all entries needed for Types CU list
   const GDBIndexTUEntryType &getGDBIndexTUEntryVector() const {
     return GDBIndexTUEntryVector;
   }
 
-  /// Stores meta data for each CU per DWO ID. It's used to create cu-index for
-  /// DWARF5.
-  void addCUnitMetaEntry(const uint64_t DWOId, const UnitMeta &Entry) {
-    auto RetVal = CUnitMetaDwoMap.insert({DWOId, Entry});
-    if (!RetVal.second)
-      errs() << "BOLT-WARNING: [internal-dwarf-error]: Trying to set CU meta "
-                "data twice for DWOID: "
-             << Twine::utohexstr(DWOId) << ".\n";
-  }
+  using OverriddenSectionsMap = std::unordered_map<DWARFSectionKind, StringRef>;
+  /// Output .dwo files.
+  void writeDWOFiles(DWARFUnit &, const OverriddenSectionsMap &,
+                     const std::string &, DebugLocWriter &);
+  using KnownSectionsEntry = std::pair<MCSection *, DWARFSectionKind>;
+  struct DWPState {
+    std::unique_ptr<ToolOutputFile> Out;
+    std::unique_ptr<BinaryContext> TmpBC;
+    std::unique_ptr<MCStreamer> Streamer;
+    std::unique_ptr<DWPStringPool> Strings;
+    const MCObjectFileInfo *MCOFI = nullptr;
+    const DWARFUnitIndex *CUIndex = nullptr;
+    std::deque<SmallString<32>> UncompressedSections;
+    MapVector<uint64_t, UnitIndexEntry> IndexEntries;
+    MapVector<uint64_t, UnitIndexEntry> TypeIndexEntries;
+    StringMap<KnownSectionsEntry> KnownSections;
+    uint32_t ContributionOffsets[8] = {};
+    uint32_t IndexVersion = 2;
+    uint64_t DebugInfoSize = 0;
+    uint16_t Version = 0;
+    bool IsDWP = false;
+  };
+  /// Init .dwp file
+  void initDWPState(DWPState &);
 
-  /// Stores meta data for each TU per DWO ID. It's used to create cu-index for
-  /// DWARF5.
-  void addTUnitMetaEntry(const uint64_t DWOId, const UnitMeta &Entry) {
-    TUnitMetaDwoMap[DWOId].emplace_back(Entry);
-  }
+  /// Write out .dwp File
+  void finalizeDWP(DWPState &);
 
-  /// Returns Meta data for TUs in offset increasing order.
-  UnitMetaVectorType &getTUnitMetaEntries(const uint64_t DWOId) {
-    return TUnitMetaDwoMap[DWOId];
-  }
-  /// Returns Meta data for TUs in offset increasing order.
-  const UnitMeta &getCUnitMetaEntry(const uint64_t DWOId) {
-    return CUnitMetaDwoMap[DWOId];
-  }
+  /// add content of dwo to .dwp file.
+  void updateDWP(DWARFUnit &, const OverriddenSectionsMap &, const UnitMeta &,
+                 UnitMetaVectorType &, DWPState &, DebugLocWriter &);
 };
 
 } // namespace bolt
 } // namespace llvm
 
 #endif

bolt/lib/Rewrite/DWARFRewriter.cpp

 #include "bolt/Utils/Utils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/BinaryFormat/Dwarf.h"
 #include "llvm/CodeGen/AsmPrinter.h"
 #include "llvm/CodeGen/DIE.h"
 #include "llvm/DWARFLinker/DWARFStreamer.h"
-#include "llvm/DWP/DWP.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
 #include "llvm/DebugInfo/DWARF/DWARFDebugLoc.h"
 #include "llvm/DebugInfo/DWARF/DWARFExpression.h"
 #include "llvm/DebugInfo/DWARF/DWARFFormValue.h"
 #include "llvm/DebugInfo/DWARF/DWARFTypeUnit.h"
 #include "llvm/DebugInfo/DWARF/DWARFUnit.h"
 #include "llvm/MC/MCAsmBackend.h"
 #include "llvm/MC/MCAsmLayout.h"
-#include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/MC/MCStreamer.h"
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/Endian.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/LEB128.h"
 #include "llvm/Support/ThreadPool.h"
-#include "llvm/Support/ToolOutputFile.h"
 #include "llvm/Support/raw_ostream.h"
 #include <algorithm>
 #include <cstdint>
 #include <functional>
 #include <iterator>
 #include <memory>
 #include <optional>
 #include <string>
if (std::optional<object::SectionedAddress> SA =
     return SA->Address;
   return std::nullopt;
 }
 
 /// Returns DWO Name to be used. Handles case where user specifies output DWO
 /// directory, and there are duplicate names. Assumes DWO ID is unique.
 static std::string
 getDWOName(llvm::DWARFUnit &CU,
-           std::unordered_map<std::string, uint32_t> *NameToIndexMap,
-           std::unordered_map<uint64_t, std::string> &DWOIdToName) {
+           std::unordered_map<std::string, uint32_t> *NameToIndexMap) {
   std::optional<uint64_t> DWOId = CU.getDWOId();
   assert(DWOId && "DWO ID not found.");
   (void)DWOId;
-  auto NameIter = DWOIdToName.find(*DWOId);
-  if (NameIter != DWOIdToName.end())
-    return NameIter->second;
 
   std::string DWOName = dwarf::toString(
       CU.getUnitDIE().find({dwarf::DW_AT_dwo_name, dwarf::DW_AT_GNU_dwo_name}),
       "");
   assert(!DWOName.empty() &&
          "DW_AT_dwo_name/DW_AT_GNU_dwo_name does not exists.");
   if (NameToIndexMap && !opts::DwarfOutputPath.empty()) {
     auto Iter = NameToIndexMap->find(DWOName);
     if (Iter == NameToIndexMap->end())
       Iter = NameToIndexMap->insert({DWOName, 0}).first;
     DWOName.append(std::to_string(Iter->second));
     ++Iter->second;
   }
   DWOName.append(".dwo");
-  DWOIdToName[*DWOId] = DWOName;
   return DWOName;
 }
 
 static std::unique_ptr<DIEStreamer>
 createDIEStreamer(const Triple &TheTriple, raw_pwrite_stream &OutFile,
                   StringRef Swift5ReflectionSegmentName, DIEBuilder &DIEBldr,
                   DWARFRewriter &Rewriter) {
 
static DWARFRewriter::UnitMeta emitUnit(DIEBuilder &DIEBldr,
   const DIEBuilder::UnitInfo &U = DIEBldr.getUnitInfoByDwarfUnit(Unit);
   Streamer->emitUnit(Unit, *UnitDIE);
   uint64_t TypeHash = 0;
   if (DWARFTypeUnit *DTU = dyn_cast_or_null<DWARFTypeUnit>(&Unit))
     TypeHash = DTU->getTypeHash();
   return {U.UnitOffset, U.UnitLength, TypeHash};
 }
 
-static void emitDWOBuilder(DIEBuilder &DWODIEBuilder, DWARFRewriter &Rewriter,
-                           const DWARFUnit &SplitCU, const uint64_t DWOId) {
-  // populate debug_info and debug_abbrev for current dwo into StringRef.
+static void emitDWOBuilder(const std::string &DWOName,
+                           DIEBuilder &DWODIEBuilder, DWARFRewriter &Rewriter,
+                           const DWARFUnit &SplitCU, DWARFUnit &CU,
+                           DWARFRewriter::DWPState &State,
+                           DebugLocWriter &LocWriter) {
+  // populate debug_info and debug_abbrev for current dwo into stringRef.
   DWODIEBuilder.generateAbbrevs();
   DWODIEBuilder.finish();
 
   SmallVector<char, 20> OutBuffer;
   std::shared_ptr<raw_svector_ostream> ObjOS =
       std::make_shared<raw_svector_ostream>(OutBuffer);
   const object::ObjectFile *File = SplitCU.getContext().getDWARFObj().getFile();
   auto TheTriple = std::make_unique<Triple>(File->makeTriple());
   std::unique_ptr<DIEStreamer> Streamer = createDIEStreamer(
       *TheTriple, *ObjOS, "DwoStreamerInitAug2", DWODIEBuilder, Rewriter);
+  DWARFRewriter::UnitMetaVectorType TUMetaVector;
+  DWARFRewriter::UnitMeta CUMI = {0, 0, 0};
   if (SplitCU.getContext().getMaxDWOVersion() >= 5) {
     // TODO: Handle DWP as input. Right now it will iterate over all of CUs and
     // TUs
     for (std::unique_ptr<llvm::DWARFUnit> &CU :
          SplitCU.getContext().dwo_info_section_units()) {
       DWARFRewriter::UnitMeta MI = emitUnit(DWODIEBuilder, Streamer, *CU.get());
       if (CU->isTypeUnit())
-        Rewriter.addTUnitMetaEntry(DWOId, MI);
+        TUMetaVector.emplace_back(MI);
       else
-        Rewriter.addCUnitMetaEntry(DWOId, MI);
+        CUMI = MI;
     }
   } else {
     for (std::unique_ptr<llvm::DWARFUnit> &CU :
          SplitCU.getContext().dwo_compile_units())
       emitUnit(DWODIEBuilder, Streamer, *CU.get());
 
     // emit debug_types sections for dwarf4
     for (DWARFUnit *CU : DWODIEBuilder.getDWARF4TUVector()) {
       DWARFRewriter::UnitMeta MI = emitUnit(DWODIEBuilder, Streamer, *CU);
-      Rewriter.addTUnitMetaEntry(DWOId, MI);
+      TUMetaVector.emplace_back(MI);
     }
   }
 
   Streamer->emitAbbrevs(DWODIEBuilder.getAbbrevs(),
                         SplitCU.getContext().getMaxVersion());
   Streamer->finish();
 
   std::unique_ptr<MemoryBuffer> ObjectMemBuffer =
       MemoryBuffer::getMemBuffer(ObjOS->str(), "in-memory object file", false);
   std::unique_ptr<object::ObjectFile> Obj = cantFail(
       object::ObjectFile::createObjectFile(ObjectMemBuffer->getMemBufferRef()),
       "error creating in-memory object");
 
+  DWARFRewriter::OverriddenSectionsMap OverriddenSections;
   for (const SectionRef &Secs : Obj->sections()) {
     StringRef Contents = cantFail(Secs.getContents());
     StringRef Name = cantFail(Secs.getName());
-    if (Name.equals(".debug_abbrev"))
-      Rewriter.setDwoDebugAbbrevStr(DWOId, Contents);
-    else if (Name.equals(".debug_info"))
-      Rewriter.setDwoDebugInfoStr(DWOId, Contents);
-    else if (Name.equals(".debug_types"))
-      Rewriter.setDwoDebugTypeStr(DWOId, Contents);
+    DWARFSectionKind Kind =
+        StringSwitch<DWARFSectionKind>(Name)
+            .Case(".debug_abbrev", DWARFSectionKind::DW_SECT_ABBREV)
+            .Case(".debug_info", DWARFSectionKind::DW_SECT_INFO)
+            .Case(".debug_types", DWARFSectionKind::DW_SECT_EXT_TYPES)
+            .Default(DWARFSectionKind::DW_SECT_EXT_unknown);
+    if (Kind == DWARFSectionKind::DW_SECT_EXT_unknown)
+      continue;
+    OverriddenSections[Kind] = Contents;
   }
+  if (opts::WriteDWP)
+    Rewriter.updateDWP(CU, OverriddenSections, CUMI, TUMetaVector, State,
+                       LocWriter);
+  else
+    Rewriter.writeDWOFiles(CU, OverriddenSections, DWOName, LocWriter);
 }
 
 void DWARFRewriter::addStringHelper(DIEBuilder &DIEBldr, DIE &Die,
                                     const DWARFUnit &Unit,
                                     DIEValue &DIEAttrInfo, StringRef Str) {
   uint32_t NewOffset = StrWriter->addString(Str);
   if (Unit.getVersion() == 5) {
     StrOffstsWriter->updateAddressMap(DIEAttrInfo.getDIEInteger().getValue(),
if (DwarfVersion >= 5) {
             std::make_unique<DebugRangeListsSectionWriter>();
         RangeListsSectionWriter->initSection(*CU.get());
         RangeListsWritersByCU[*DWOId] = std::move(RangeListsSectionWriter);
       }
 
     } else {
       LocListWritersByCU[CUIndex] = std::make_unique<DebugLocWriter>();
     }
-
-    if (std::optional<uint64_t> DWOId = CU->getDWOId()) {
-      assert(LocListWritersByCU.count(*DWOId) == 0 &&
-             "LocList writer for DWO unit already exists.");
-      // Work around some bug in llvm-15. If I pass in directly lld reports
-      // undefined symbol.
-      LocListWritersByCU[*DWOId] =
-          std::make_unique<DebugLoclistWriter>(*CU.get(), DwarfVersion, true);
-    }
     ++CUIndex;
   }
 
   // Unordered maps to handle name collision if output DWO directory is
   // specified.
   std::unordered_map<std::string, uint32_t> NameToIndexMap;
-  std::unordered_map<uint64_t, std::string> DWOIdToName;
   std::mutex AccessMutex;
 
   auto updateDWONameCompDir = [&](DWARFUnit &Unit, DIEBuilder &DIEBldr,
-                                  DIE &UnitDIE) -> void {
+                                  DIE &UnitDIE) -> std::string {
     DIEValue DWONameAttrInfo = UnitDIE.findAttribute(dwarf::DW_AT_dwo_name);
     if (!DWONameAttrInfo)
       DWONameAttrInfo = UnitDIE.findAttribute(dwarf::DW_AT_GNU_dwo_name);
     assert(DWONameAttrInfo && "DW_AT_dwo_name is not in Skeleton CU.");
     std::string ObjectName;
 
     {
       std::lock_guard<std::mutex> Lock(AccessMutex);
-      ObjectName = getDWOName(Unit, &NameToIndexMap, DWOIdToName);
+      ObjectName = getDWOName(Unit, &NameToIndexMap);
     }
     addStringHelper(DIEBldr, UnitDIE, Unit, DWONameAttrInfo,
                     ObjectName.c_str());
 
     DIEValue CompDirAttrInfo = UnitDIE.findAttribute(dwarf::DW_AT_comp_dir);
     assert(CompDirAttrInfo && "DW_AT_comp_dir is not in Skeleton CU.");
 
     if (!opts::DwarfOutputPath.empty()) {
       addStringHelper(DIEBldr, UnitDIE, Unit, CompDirAttrInfo,
                       opts::DwarfOutputPath.c_str());
     }
+    return ObjectName;
   };
 
+  DWPState State;
+  if (opts::WriteDWP)
+    initDWPState(State);
   auto processUnitDIE = [&](size_t CUIndex, DWARFUnit *Unit,
                             DIEBuilder *DIEBlder) {
     // Check if the unit is a skeleton and we need special updates for it and
     // its matching split/DWO CU.
     std::optional<DWARFUnit *> SplitCU;
     std::optional<uint64_t> RangesBase;
     std::optional<uint64_t> DWOId = Unit->getDWOId();
     StrOffstsWriter->initialize(Unit->getStringOffsetSection(),
                                 Unit->getStringOffsetsTableContribution());
     if (DWOId)
       SplitCU = BC.getDWOCU(*DWOId);
 
     DebugLocWriter *DebugLocWriter = nullptr;
     DebugRangesSectionWriter *RangesSectionWriter =
         Unit->getVersion() >= 5 ? RangeListsSectionWriter.get()
                                 : LegacyRangesSectionWriter.get();
     // Skipping CUs that failed to load.
     if (SplitCU) {
       DIEBuilder DWODIEBuilder(&(*SplitCU)->getContext(), true);
-      updateDWONameCompDir(*Unit, *DIEBlder,
-                           *DIEBlder->getUnitDIEbyUnit(*Unit));
+      std::string DWOName = updateDWONameCompDir(
+          *Unit, *DIEBlder, *DIEBlder->getUnitDIEbyUnit(*Unit));
 
-      {
-        std::lock_guard<std::mutex> Lock(AccessMutex);
-        DebugLocWriter = LocListWritersByCU[*DWOId].get();
-      }
+      DebugLoclistWriter DebugLocDWoWriter(*Unit, Unit->getVersion(), true);
       DebugRangesSectionWriter *TempRangesSectionWriter = RangesSectionWriter;
       if (Unit->getVersion() >= 5) {
         TempRangesSectionWriter = RangeListsWritersByCU[*DWOId].get();
       } else {
         RangesBase = RangesSectionWriter->getSectionOffset();
         // For DWARF5 there is now .debug_rnglists.dwo, so don't need to
         // update rnglists base.
         if (RangesBase) {
           DwoRangesBase[*DWOId] = *RangesBase;
           setDwoRangesBase(*DWOId, *RangesBase);
         }
       }
 
-      updateUnitDebugInfo(*(*SplitCU), DWODIEBuilder, *DebugLocWriter,
+      updateUnitDebugInfo(*(*SplitCU), DWODIEBuilder, DebugLocDWoWriter,
                           *TempRangesSectionWriter);
-      DebugLocWriter->finalize(DWODIEBuilder,
+      DebugLocDWoWriter.finalize(DWODIEBuilder,
                                  *DWODIEBuilder.getUnitDIEbyUnit(**SplitCU));
       if (Unit->getVersion() >= 5)
         TempRangesSectionWriter->finalizeSection();
 
-      emitDWOBuilder(DWODIEBuilder, *this, **SplitCU, *DWOId);
+      emitDWOBuilder(DWOName, DWODIEBuilder, *this, **SplitCU, *Unit, State,
+                     DebugLocDWoWriter);
     }
 
     {
       std::lock_guard<std::mutex> Lock(AccessMutex);
       auto LocListWriterIter = LocListWritersByCU.find(CUIndex);
       if (LocListWriterIter != LocListWritersByCU.end())
         DebugLocWriter = LocListWriterIter->second.get();
     }
if (opts::NoThreads || opts::DeterministicDebugInfo) {
     ThreadPool &ThreadPool = ParallelUtilities::getThreadPool();
     for (std::unique_ptr<DWARFUnit> &CU : BC.DwCtx->compile_units()) {
       ThreadPool.async(processUnitDIE, CUIndex, CU.get(), &DIEBlder);
       CUIndex++;
     }
     ThreadPool.wait();
   }
 
-  CUOffsetMap OffsetMap = finalizeDebugSections(DIEBlder);
-
   if (opts::WriteDWP)
-    writeDWP(DWOIdToName);
-  else
-    writeDWOFiles(DWOIdToName);
+    finalizeDWP(State);
 
+  CUOffsetMap OffsetMap = finalizeDebugSections(DIEBlder);
   updateGdbIndexSection(OffsetMap, CUIndex);
 }
 
 void DWARFRewriter::updateUnitDebugInfo(
     DWARFUnit &Unit, DIEBuilder &DIEBldr, DebugLocWriter &DebugLocWriter,
     DebugRangesSectionWriter &RangesSectionWriter,
     std::optional<uint64_t> RangesBase) {
   // Cache debug ranges so that the offset for identical ranges could be reused.
CUOffsetMap DWARFRewriter::finalizeDebugSections(DIEBuilder &DIEBlder) {
   }
   return CUMap;
 }
 
 // Creates all the data structures necessary for creating MCStreamer.
 // They are passed by reference because they need to be kept around.
 // Also creates known debug sections. These are sections handled by
 // handleDebugDataPatching.
-using KnownSectionsEntry = std::pair<MCSection *, DWARFSectionKind>;
 namespace {
 
 std::unique_ptr<BinaryContext>
 createDwarfOnlyBC(const object::ObjectFile &File) {
   return cantFail(BinaryContext::createBinaryContext(
       &File, false,
       DWARFContext::create(File, DWARFContext::ProcessDebugRelocations::Ignore,
                            nullptr, "", WithColor::defaultErrorHandler,
                            WithColor::defaultWarningHandler)));
 }
 
-StringMap<KnownSectionsEntry>
+StringMap<DWARFRewriter::KnownSectionsEntry>
 createKnownSectionsMap(const MCObjectFileInfo &MCOFI) {
-  StringMap<KnownSectionsEntry> KnownSectionsTemp = {
+  StringMap<DWARFRewriter::KnownSectionsEntry> KnownSectionsTemp = {
       {"debug_info.dwo", {MCOFI.getDwarfInfoDWOSection(), DW_SECT_INFO}},
       {"debug_types.dwo", {MCOFI.getDwarfTypesDWOSection(), DW_SECT_EXT_TYPES}},
       {"debug_str_offsets.dwo",
        {MCOFI.getDwarfStrOffDWOSection(), DW_SECT_STR_OFFSETS}},
       {"debug_str.dwo", {MCOFI.getDwarfStrDWOSection(), DW_SECT_EXT_unknown}},
       {"debug_loc.dwo", {MCOFI.getDwarfLocDWOSection(), DW_SECT_EXT_LOC}},
       {"debug_abbrev.dwo", {MCOFI.getDwarfAbbrevDWOSection(), DW_SECT_ABBREV}},
       {"debug_line.dwo", {MCOFI.getDwarfLineDWOSection(), DW_SECT_LINE}},
StringRef getSectionName(const SectionRef &Section) {
   Name = Name.substr(Name.find_first_not_of("._"));
   return Name;
 }
 
 // Exctracts an appropriate slice if input is DWP.
 // Applies patches or overwrites the section.
 std::optional<StringRef> updateDebugData(
     DWARFContext &DWCtx, StringRef SectionName, StringRef SectionContents,
-    const StringMap<KnownSectionsEntry> &KnownSections, MCStreamer &Streamer,
-    DWARFRewriter &Writer, const DWARFUnitIndex::Entry *CUDWOEntry,
-    uint64_t DWOId, std::unique_ptr<DebugBufferVector> &OutputBuffer,
-    DebugRangeListsSectionWriter *RangeListsWriter) {
+    const StringMap<DWARFRewriter::KnownSectionsEntry> &KnownSections,
+    MCStreamer &Streamer, DWARFRewriter &Writer,
+    const DWARFUnitIndex::Entry *CUDWOEntry, uint64_t DWOId,
+    std::unique_ptr<DebugBufferVector> &OutputBuffer,
+    DebugRangeListsSectionWriter *RangeListsWriter, DebugLocWriter &LocWriter,
+    const llvm::bolt::DWARFRewriter::OverriddenSectionsMap &OverridenSections) {
 
   using DWOSectionContribution =
       const DWARFUnitIndex::Entry::SectionContribution;
   auto getSliceData = [&](const DWARFUnitIndex::Entry *DWOEntry,
                           StringRef OutData, DWARFSectionKind Sec,
                           uint64_t &DWPOffset) -> StringRef {
     if (DWOEntry) {
       DWOSectionContribution *DWOContrubution = DWOEntry->getContribution(Sec);
       DWPOffset = DWOContrubution->getOffset();
       OutData = OutData.substr(DWPOffset, DWOContrubution->getLength());
     }
     return OutData;
   };
 
   auto SectionIter = KnownSections.find(SectionName);
   if (SectionIter == KnownSections.end())
     return std::nullopt;
   Streamer.switchSection(SectionIter->second.first);
   uint64_t DWPOffset = 0;
 
+  auto getOverridenSection =
+      [&](DWARFSectionKind Kind) -> std::optional<StringRef> {
+    auto Iter = OverridenSections.find(Kind);
+    if (Iter == OverridenSections.end()) {
+      errs()
+          << "BOLT-WARNING: [internal-dwarf-error]: Could not find overriden "
+             "section for: "
+          << Twine::utohexstr(DWOId) << ".\n";
+      return std::nullopt;
+    }
+    return Iter->second;
+  };
   switch (SectionIter->second.second) {
   default: {
     if (!SectionName.equals("debug_str.dwo"))
       errs() << "BOLT-WARNING: unsupported debug section: " << SectionName
              << "\n";
     return SectionContents;
   }
   case DWARFSectionKind::DW_SECT_INFO: {
-    return Writer.getDwoDebugInfoStr(DWOId);
+    return getOverridenSection(DWARFSectionKind::DW_SECT_INFO);
   }
   case DWARFSectionKind::DW_SECT_EXT_TYPES: {
-    return Writer.getDwoDebugTypeStr(DWOId);
+    return getOverridenSection(DWARFSectionKind::DW_SECT_EXT_TYPES);
   }
   case DWARFSectionKind::DW_SECT_STR_OFFSETS: {
     return getSliceData(CUDWOEntry, SectionContents,
                         DWARFSectionKind::DW_SECT_STR_OFFSETS, DWPOffset);
   }
   case DWARFSectionKind::DW_SECT_ABBREV: {
-    return Writer.getDwoDebugAbbrevStr(DWOId);
+    return getOverridenSection(DWARFSectionKind::DW_SECT_ABBREV);
   }
   case DWARFSectionKind::DW_SECT_EXT_LOC:
   case DWARFSectionKind::DW_SECT_LOCLISTS: {
-    DebugLocWriter *LocWriter = Writer.getDebugLocWriter(DWOId);
-    OutputBuffer = LocWriter->getBuffer();
+    OutputBuffer = LocWriter.getBuffer();
     // Creating explicit StringRef here, otherwise
     // with impicit conversion it will take null byte as end of
     // string.
     return StringRef(reinterpret_cast<const char *>(OutputBuffer->data()),
                      OutputBuffer->size());
   }
   case DWARFSectionKind::DW_SECT_LINE: {
     return getSliceData(CUDWOEntry, SectionContents,
                         DWARFSectionKind::DW_SECT_LINE, DWPOffset);
   }
   case DWARFSectionKind::DW_SECT_RNGLISTS: {
     OutputBuffer = RangeListsWriter->releaseBuffer();
     return StringRef(reinterpret_cast<const char *>(OutputBuffer->data()),
                      OutputBuffer->size());
   }
   }
 }
 
 } // namespace
 
-void DWARFRewriter::writeDWP(
-    std::unordered_map<uint64_t, std::string> &DWOIdToName) {
+void DWARFRewriter::initDWPState(DWPState &State) {
   SmallString<0> OutputNameStr;
   StringRef OutputName;
   if (opts::DwarfOutputPath.empty()) {
     OutputName =
         Twine(opts::OutputFilename).concat(".dwp").toStringRef(OutputNameStr);
   } else {
     StringRef ExeFileName = llvm::sys::path::filename(opts::OutputFilename);
     OutputName = Twine(opts::DwarfOutputPath)
                      .concat("/")
                      .concat(ExeFileName)
                      .concat(".dwp")
                      .toStringRef(OutputNameStr);
     errs() << "BOLT-WARNING: dwarf-output-path is in effect and .dwp file will "
               "possibly be written to another location that is not the same as "
               "the executable\n";
   }
   std::error_code EC;
-  std::unique_ptr<ToolOutputFile> Out =
+  State.Out =
       std::make_unique<ToolOutputFile>(OutputName, EC, sys::fs::OF_None);
-
   const object::ObjectFile *File = BC.DwCtx->getDWARFObj().getFile();
-  std::unique_ptr<BinaryContext> TmpBC = createDwarfOnlyBC(*File);
-  std::unique_ptr<MCStreamer> Streamer = TmpBC->createStreamer(Out->os());
-  const MCObjectFileInfo &MCOFI = *Streamer->getContext().getObjectFileInfo();
-  StringMap<KnownSectionsEntry> KnownSections = createKnownSectionsMap(MCOFI);
-  MCSection *const StrSection = MCOFI.getDwarfStrDWOSection();
-  MCSection *const StrOffsetSection = MCOFI.getDwarfStrOffDWOSection();
+  State.TmpBC = createDwarfOnlyBC(*File);
+  State.Streamer = State.TmpBC->createStreamer(State.Out->os());
+  State.MCOFI = State.Streamer->getContext().getObjectFileInfo();
+  State.KnownSections = createKnownSectionsMap(*State.MCOFI);
+  MCSection *const StrSection = State.MCOFI->getDwarfStrDWOSection();
 
   // Data Structures for DWP book keeping
   // Size of array corresponds to the number of sections supported by DWO format
   // in DWARF4/5.
-  uint32_t ContributionOffsets[8] = {};
-  std::deque<SmallString<32>> UncompressedSections;
-  DWPStringPool Strings(*Streamer, StrSection);
-  MapVector<uint64_t, UnitIndexEntry> IndexEntries;
-  MapVector<uint64_t, UnitIndexEntry> TypeIndexEntries;
-  uint16_t Version = 0;
-  uint32_t IndexVersion = 2;
+
+  State.Strings = std::make_unique<DWPStringPool>(*State.Streamer, StrSection);
 
   // Setup DWP code once.
   DWARFContext *DWOCtx = BC.getDWOContext();
-  const DWARFUnitIndex *CUIndex = nullptr;
-  bool IsDWP = false;
+
   if (DWOCtx) {
-    CUIndex = &DWOCtx->getCUIndex();
-    IsDWP = !CUIndex->getRows().empty();
+    State.CUIndex = &DWOCtx->getCUIndex();
+    State.IsDWP = !State.CUIndex->getRows().empty();
   }
-  uint64_t DebugInfoSize = 0;
-  for (const std::unique_ptr<DWARFUnit> &CU : BC.DwCtx->compile_units()) {
-    std::optional<uint64_t> DWOId = CU->getDWOId();
-    if (!DWOId)
-      continue;
+}
+
+void DWARFRewriter::finalizeDWP(DWPState &State) {
+  if (State.Version < 5) {
+    // Lie about there being no info contributions so the TU index only includes
+    // the type unit contribution for DWARF < 5. In DWARFv5 the TU index has a
+    // contribution to the info section, so we do not want to lie about it.
+    State.ContributionOffsets[0] = 0;
+  }
+  writeIndex(*State.Streamer.get(), State.MCOFI->getDwarfTUIndexSection(),
+             State.ContributionOffsets, State.TypeIndexEntries,
+             State.IndexVersion);
 
-    // Skipping CUs that we failed to load.
-    std::optional<DWARFUnit *> DWOCU = BC.getDWOCU(*DWOId);
-    if (!DWOCU)
-      continue;
+  if (State.Version < 5) {
+    // Lie about the type contribution for DWARF < 5. In DWARFv5 the type
+    // section does not exist, so no need to do anything about this.
+    State.ContributionOffsets[getContributionIndex(DW_SECT_EXT_TYPES, 2)] = 0;
+    // Unlie about the info contribution
+    State.ContributionOffsets[0] = 1;
+  }
+  writeIndex(*State.Streamer.get(), State.MCOFI->getDwarfCUIndexSection(),
+             State.ContributionOffsets, State.IndexEntries, State.IndexVersion);
+
+  State.Streamer->finish();
+  State.Out->keep();
+}
+
+void DWARFRewriter::updateDWP(DWARFUnit &CU,
+                              const OverriddenSectionsMap &OverridenSections,
+                              const DWARFRewriter::UnitMeta &CUMI,
+                              DWARFRewriter::UnitMetaVectorType &TUMetaVector,
+                              DWPState &State, DebugLocWriter &LocWriter) {
+  uint64_t DWOId = *CU.getDWOId();
+  MCSection *const StrOffsetSection = State.MCOFI->getDwarfStrOffDWOSection();
+  std::optional<DWARFUnit *> DWOCU = BC.getDWOCU(DWOId);
+    return;
 
-    if (Version == 0) {
-      Version = CU->getVersion();
-      IndexVersion = Version < 5 ? 2 : 5;
-    } else if (Version != CU->getVersion()) {
-      errs() << "BOLT-ERROR: Incompatible DWARF compile unit versions.\n";
-      exit(1);
-    }
+  if (State.Version == 0) {
+    State.Version = CU.getVersion();
+    State.IndexVersion = State.Version < 5 ? 2 : 5;
+  } else if (State.Version != CU.getVersion()) {
 
-    UnitIndexEntry CurEntry = {};
-    CurEntry.DWOName =
-        dwarf::toString(CU->getUnitDIE().find(
-                            {dwarf::DW_AT_dwo_name, dwarf::DW_AT_GNU_dwo_name}),
-                        "");
-    const char *Name = CU->getUnitDIE().getShortName();
-    if (Name)
-      CurEntry.Name = Name;
-    StringRef CurStrSection;
-    StringRef CurStrOffsetSection;
+  CurEntry.DWOName = dwarf::toString(
+      CU.getUnitDIE().find({dwarf::DW_AT_dwo_name, dwarf::DW_AT_GNU_dwo_name}),
+  const char *Name = CU.getUnitDIE().getShortName();
 
-    // This maps each section contained in this file to its length.
-    // This information is later on used to calculate the contributions,
-    // i.e. offset and length, of each compile/type unit to a section.
-    std::vector<std::pair<DWARFSectionKind, uint32_t>> SectionLength;
 
-    const DWARFUnitIndex::Entry *CUDWOEntry = nullptr;
-    if (IsDWP)
-      CUDWOEntry = CUIndex->getFromHash(*DWOId);
+  if (State.IsDWP)
+    CUDWOEntry = State.CUIndex->getFromHash(DWOId);
 
-    bool StrSectionWrittenOut = false;
-    const object::ObjectFile *DWOFile =
-        (*DWOCU)->getContext().getDWARFObj().getFile();
 
-    DebugRangeListsSectionWriter *RangeListssWriter = nullptr;
-    if (CU->getVersion() == 5) {
-      assert(RangeListsWritersByCU.count(*DWOId) != 0 &&
-             "No RangeListsWriter for DWO ID.");
-      RangeListssWriter = RangeListsWritersByCU[*DWOId].get();
-    }
-    auto AddType = [&](unsigned int Index, uint32_t IndexVersion,
-                       uint64_t Offset, uint64_t Length,
-                       uint64_t Hash) -> void {
-      UnitIndexEntry TUEntry = CurEntry;
-      if (IndexVersion < 5)
-        TUEntry.Contributions[0] = {};
-      TUEntry.Contributions[Index].setOffset(Offset);
-      TUEntry.Contributions[Index].setLength(Length);
-      ContributionOffsets[Index] += TUEntry.Contributions[Index].getLength32();
-      TypeIndexEntries.insert(std::make_pair(Hash, TUEntry));
-    };
-    for (const SectionRef &Section : DWOFile->sections()) {
-      std::unique_ptr<DebugBufferVector> OutputData;
-      StringRef SectionName = getSectionName(Section);
-      Expected<StringRef> ContentsExp = Section.getContents();
-      assert(ContentsExp && "Invalid contents.");
-      std::optional<StringRef> TOutData = updateDebugData(
-          (*DWOCU)->getContext(), SectionName, *ContentsExp, KnownSections,
-          *Streamer, *this, CUDWOEntry, *DWOId, OutputData, RangeListssWriter);
-      if (!TOutData)
-        continue;
+  if (CU.getVersion() == 5) {
+    assert(RangeListsWritersByCU.count(DWOId) != 0 &&
+    RangeListssWriter = RangeListsWritersByCU[DWOId].get();
+  auto AddType = [&](unsigned int Index, uint32_t IndexVersion, uint64_t Offset,
+                     uint64_t Length, uint64_t Hash) -> void {
+    State.ContributionOffsets[Index] +=
+        TUEntry.Contributions[Index].getLength32();
+    State.TypeIndexEntries.insert(std::make_pair(Hash, TUEntry));
+        (*DWOCU)->getContext(), SectionName, *ContentsExp, State.KnownSections,
+        *State.Streamer, *this, CUDWOEntry, DWOId, OutputData,
+        RangeListssWriter, LocWriter, OverridenSections);
 
-      StringRef OutData = *TOutData;
-      if (SectionName == "debug_types.dwo") {
-        Streamer->emitBytes(OutData);
-        continue;
-      }
+      State.Streamer->emitBytes(OutData);
 
-      if (SectionName.equals("debug_str.dwo")) {
-        CurStrSection = OutData;
-      } else {
-        // Since handleDebugDataPatching returned true, we already know this is
-        // a known section.
-        auto SectionIter = KnownSections.find(SectionName);
-        if (SectionIter->second.second == DWARFSectionKind::DW_SECT_STR_OFFSETS)
-          CurStrOffsetSection = OutData;
-        else
-          Streamer->emitBytes(OutData);
-        unsigned int Index =
-            getContributionIndex(SectionIter->second.second, IndexVersion);
-        uint64_t Offset = ContributionOffsets[Index];
-        uint64_t Length = OutData.size();
-        if (CU->getVersion() >= 5 &&
-            SectionIter->second.second == DWARFSectionKind::DW_SECT_INFO) {
-          for (UnitMeta &MI : getTUnitMetaEntries(*DWOId))
-            MI.Offset += DebugInfoSize;
+      auto SectionIter = State.KnownSections.find(SectionName);
+        State.Streamer->emitBytes(OutData);
+          getContributionIndex(SectionIter->second.second, State.IndexVersion);
+      uint64_t Offset = State.ContributionOffsets[Index];
+      if (CU.getVersion() >= 5 &&
+        for (UnitMeta &MI : TUMetaVector)
+          MI.Offset += State.DebugInfoSize;
 
-          const UnitMeta &MI = getCUnitMetaEntry(*DWOId);
-          Offset = DebugInfoSize + MI.Offset;
-          Length = MI.Length;
-          DebugInfoSize += OutData.size();
-        }
-        CurEntry.Contributions[Index].setOffset(Offset);
-        CurEntry.Contributions[Index].setLength(Length);
-        ContributionOffsets[Index] +=
-            CurEntry.Contributions[Index].getLength32();
-      }
+        Offset = State.DebugInfoSize + CUMI.Offset;
+        Length = CUMI.Length;
+        State.DebugInfoSize += OutData.size();
+      State.ContributionOffsets[Index] +=
 
-      // Strings are combined in to a new string section, and de-duplicated
-      // based on hash.
-      if (!StrSectionWrittenOut && !CurStrOffsetSection.empty() &&
-          !CurStrSection.empty()) {
-        writeStringsAndOffsets(*Streamer.get(), Strings, StrOffsetSection,
-                               CurStrSection, CurStrOffsetSection,
-                               CU->getVersion());
-        StrSectionWrittenOut = true;
-      }
-    }
-    CompileUnitIdentifiers CUI{*DWOId, CurEntry.Name.c_str(),
-                               CurEntry.DWOName.c_str()};
-    auto P = IndexEntries.insert(std::make_pair(CUI.Signature, CurEntry));
-    if (!P.second) {
-      Error Err = buildDuplicateError(*P.first, CUI, "");
-      errs() << "BOLT-ERROR: " << toString(std::move(Err)) << "\n";
-      return;
-    }
+      writeStringsAndOffsets(*State.Streamer.get(), *State.Strings.get(),
+                             StrOffsetSection, CurStrSection,
+                             CurStrOffsetSection, CU.getVersion());
+  CompileUnitIdentifiers CUI{DWOId, CurEntry.Name.c_str(),
+  auto P = State.IndexEntries.insert(std::make_pair(CUI.Signature, CurEntry));
 
-    // Handling TU
-    const unsigned Index = getContributionIndex(
-        IndexVersion < 5 ? DW_SECT_EXT_TYPES : DW_SECT_INFO, IndexVersion);
-    for (UnitMeta &MI : getTUnitMetaEntries(*DWOId))
-      AddType(Index, IndexVersion, MI.Offset, MI.Length, MI.TUHash);
-  }
-
-  if (Version < 5) {
-    // Lie about there being no info contributions so the TU index only includes
-    // the type unit contribution for DWARF < 5. In DWARFv5 the TU index has a
-    // contribution to the info section, so we do not want to lie about it.
-    ContributionOffsets[0] = 0;
-  }
-  writeIndex(*Streamer.get(), MCOFI.getDwarfTUIndexSection(),
-             ContributionOffsets, TypeIndexEntries, IndexVersion);
-
-  if (Version < 5) {
-    // Lie about the type contribution for DWARF < 5. In DWARFv5 the type
-    // section does not exist, so no need to do anything about this.
-    ContributionOffsets[getContributionIndex(DW_SECT_EXT_TYPES, 2)] = 0;
-    // Unlie about the info contribution
-    ContributionOffsets[0] = 1;
-  }
-  writeIndex(*Streamer.get(), MCOFI.getDwarfCUIndexSection(),
-             ContributionOffsets, IndexEntries, IndexVersion);
-
-  Streamer->finish();
-  Out->keep();
+      State.IndexVersion < 5 ? DW_SECT_EXT_TYPES : DW_SECT_INFO,
+      State.IndexVersion);
+  for (UnitMeta &MI : TUMetaVector)
+    AddType(Index, State.IndexVersion, MI.Offset, MI.Length, MI.TUHash);
 }
 
 void DWARFRewriter::writeDWOFiles(
-    std::unordered_map<uint64_t, std::string> &DWOIdToName) {
+    DWARFUnit &CU, const OverriddenSectionsMap &OverridenSections,
+    const std::string &DWOName, DebugLocWriter &LocWriter) {
   // Setup DWP code once.
   DWARFContext *DWOCtx = BC.getDWOContext();
+  const uint64_t DWOId = *CU.getDWOId();
   const DWARFUnitIndex *CUIndex = nullptr;
   bool IsDWP = false;
   if (DWOCtx) {
     CUIndex = &DWOCtx->getCUIndex();
     IsDWP = !CUIndex->getRows().empty();
   }
 
-  for (const std::unique_ptr<DWARFUnit> &CU : BC.DwCtx->compile_units()) {
-    std::optional<uint64_t> DWOId = CU->getDWOId();
-    if (!DWOId)
-      continue;
-
-    // Skipping CUs that we failed to load.
-    std::optional<DWARFUnit *> DWOCU = BC.getDWOCU(*DWOId);
-    if (!DWOCU)
-      continue;
+  std::optional<DWARFUnit *> DWOCU = BC.getDWOCU(DWOId);
+  if (!DWOCU) {
+    errs() << "BOLT-WARNING: [internal-dwarf-error]: CU for DWO_ID "
+           << Twine::utohexstr(DWOId) << " is not found.\n";
+    return;
+  }
 
-    std::string CompDir = opts::DwarfOutputPath.empty()
-                              ? CU->getCompilationDir()
-                              : opts::DwarfOutputPath.c_str();
-    std::string ObjectName = getDWOName(*CU.get(), nullptr, DWOIdToName);
-    auto FullPath = CompDir.append("/").append(ObjectName);
+                            ? CU.getCompilationDir()
+  auto FullPath = CompDir.append("/").append(DWOName);
 
-    std::error_code EC;
-    std::unique_ptr<ToolOutputFile> TempOut =
-        std::make_unique<ToolOutputFile>(FullPath, EC, sys::fs::OF_None);
 
-    const DWARFUnitIndex::Entry *CUDWOEntry = nullptr;
-    if (IsDWP)
-      CUDWOEntry = CUIndex->getFromHash(*DWOId);
+    CUDWOEntry = CUIndex->getFromHash(DWOId);
 
-    const object::ObjectFile *File =
-        (*DWOCU)->getContext().getDWARFObj().getFile();
-    std::unique_ptr<BinaryContext> TmpBC = createDwarfOnlyBC(*File);
-    std::unique_ptr<MCStreamer> Streamer = TmpBC->createStreamer(TempOut->os());
-    const MCObjectFileInfo &MCOFI = *Streamer->getContext().getObjectFileInfo();
-    StringMap<KnownSectionsEntry> KnownSections = createKnownSectionsMap(MCOFI);
 
-    DebugRangeListsSectionWriter *RangeListssWriter = nullptr;
-    if (CU->getVersion() == 5) {
-      assert(RangeListsWritersByCU.count(*DWOId) != 0 &&
-             "No RangeListsWriter for DWO ID.");
-      RangeListssWriter = RangeListsWritersByCU[*DWOId].get();
+  if (CU.getVersion() == 5) {
+    assert(RangeListsWritersByCU.count(DWOId) != 0 &&
+    RangeListssWriter = RangeListsWritersByCU[DWOId].get();
 
-      // Handling .debug_rnglists.dwo seperatly. The original .o/.dwo might not
-      // have .debug_rnglists so won't be part of the loop below.
-      if (!RangeListssWriter->empty()) {
-        std::unique_ptr<DebugBufferVector> OutputData;
-        if (std::optional<StringRef> OutData =
-                updateDebugData((*DWOCU)->getContext(), "debug_rnglists.dwo",
-                                "", KnownSections, *Streamer, *this, CUDWOEntry,
-                                *DWOId, OutputData, RangeListssWriter))
-          Streamer->emitBytes(*OutData);
-      }
-    }
+      if (std::optional<StringRef> OutData = updateDebugData(
+              (*DWOCU)->getContext(), "debug_rnglists.dwo", "", KnownSections,
+              *Streamer, *this, CUDWOEntry, DWOId, OutputData,
+              RangeListssWriter, LocWriter, OverridenSections))
 
     for (const SectionRef &Section : File->sections()) {
       std::unique_ptr<DebugBufferVector> OutputData;
       StringRef SectionName = getSectionName(Section);
       if (SectionName == "debug_rnglists.dwo")
         continue;
       Expected<StringRef> ContentsExp = Section.getContents();
       assert(ContentsExp && "Invalid contents.");
-      if (std::optional<StringRef> OutData =
-              updateDebugData((*DWOCU)->getContext(), SectionName, *ContentsExp,
-                              KnownSections, *Streamer, *this, CUDWOEntry,
-                              *DWOId, OutputData, RangeListssWriter))
+      if (std::optional<StringRef> OutData = updateDebugData(
+              (*DWOCU)->getContext(), SectionName, *ContentsExp, KnownSections,
+              *Streamer, *this, CUDWOEntry, DWOId, OutputData,
+              RangeListssWriter, LocWriter, OverridenSections))
         Streamer->emitBytes(*OutData);
     }
     Streamer->finish();
     TempOut->keep();
-  }
-}
 
 void DWARFRewriter::addGDBTypeUnitEntry(const GDBIndexTUEntry &&Entry) {
   std::lock_guard<std::mutex> Lock(DWARFRewriterMutex);
   if (!BC.getGdbIndexSection())
     return;
   GDBIndexTUEntryVector.emplace_back(Entry);
 }
 
if (Version == DWARFVersion::DWARF5 &&
       continue;
 
     // Filter out DWARF5, writing out DWARF4
     if (Version == DWARFVersion::DWARFLegacy &&
         (LocListWriter && LocListWriter->getDwarfVersion() >= 5))
       continue;
 
     // Skipping DWARF4/5 split dwarf.
-    if (LocListWriter && (LocListWriter->getDwarfVersion() <= 4 ||
-                          (LocListWriter->getDwarfVersion() >= 5 &&
-                           LocListWriter->isSplitDwarf()))) {
+    if (LocListWriter && LocListWriter->getDwarfVersion() <= 4) {
       continue;
     }
     std::unique_ptr<DebugBufferVector> CurrCULocationLists =
         LocWriter->getBuffer();
     *LocStream << *CurrCULocationLists;
   }
 
   return LocBuffer;