[BOLT] Refactor SectionPatchers map to a Patcher in BinarySection

Summary:
Refactor SectionPatches to avoid the use of extra map and a cast
from StringRef to std::string.

cherry-picked from FBD26756560

(cherry picked from FBD27490641)

bolt/src/BinarySection.h

@@ -11,10 +11,11 @@
 #ifndef LLVM_TOOLS_LLVM_BOLT_BINARY_SECTION_H
 #define LLVM_TOOLS_LLVM_BOLT_BINARY_SECTION_H
 
+#include "DebugData.h"
 #include "Relocation.h"
 #include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/Triple.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/Triple.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/MC/MCSymbol.h"
 #include "llvm/Object/ELFObjectFile.h"
@@ -22,8 +23,9 @@
 #include "llvm/Object/ObjectFile.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/raw_ostream.h"
-#include <set>
 #include <map>
+#include <memory>
+#include <set>
 
 namespace llvm {
 
@@ -69,6 +71,8 @@ class BinarySection {
       : Offset(Offset), Bytes(Bytes.begin(), Bytes.end()) {}
   };
   std::vector<BinaryPatch> Patches;
+  /// Patcher used to apply simple changes to sections of the input binary.
+  std::unique_ptr<BinaryPatcher> Patcher;
 
   // Output info
   bool IsFinalized{false};         // Has this section had output information
@@ -386,6 +390,14 @@ public:
     Patches.emplace_back(BinaryPatch(Offset, Bytes));
   }
 
+  /// Register patcher for this section.
+  void registerPatcher(std::unique_ptr<BinaryPatcher> BPatcher) {
+    Patcher = std::move(BPatcher);
+  }
+
+  /// Returns the patcher
+  BinaryPatcher *getPatcher() { return Patcher.get(); }
+
   /// Lookup the relocation (if any) at the given /p Offset.
   const Relocation *getRelocationAt(uint64_t Offset) const {
     Relocation Key{Offset, 0, 0, 0, 0};

bolt/src/DWARFRewriter.cpp

@@ -11,6 +11,7 @@
 #include "DWARFRewriter.h"
 #include "BinaryContext.h"
 #include "BinaryFunction.h"
+#include "DebugData.h"
 #include "ParallelUtilities.h"
 #include "Utils.h"
 #include "llvm/ADT/STLExtras.h"
@@ -35,6 +36,7 @@
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/Timer.h"
 #include <algorithm>
+#include <llvm/Support/Error.h>
 
 #undef  DEBUG_TYPE
 #define DEBUG_TYPE "bolt"
@@ -66,14 +68,18 @@ DeterministicDebugInfo("deterministic-debuginfo",
 } // namespace opts
 
 void DWARFRewriter::updateDebugInfo() {
-  SectionPatchers[".debug_abbrev"] = std::make_unique<DebugAbbrevPatcher>();
-  SectionPatchers[".debug_info"] = std::make_unique<SimpleBinaryPatcher>();
-
-  DebugInfoPatcher =
-      static_cast<SimpleBinaryPatcher *>(SectionPatchers[".debug_info"].get());
-  AbbrevPatcher =
-      static_cast<DebugAbbrevPatcher *>(SectionPatchers[".debug_abbrev"].get());
-  assert(DebugInfoPatcher && AbbrevPatcher && "Patchers not initialized.");
+  auto DebugAbbrev = BC.getUniqueSectionByName(".debug_abbrev");
+  auto DebugInfo = BC.getUniqueSectionByName(".debug_info");
+  if (DebugAbbrev) {
+    DebugAbbrev->registerPatcher(std::make_unique<DebugAbbrevPatcher>());
+    AbbrevPatcher =
+        static_cast<DebugAbbrevPatcher *>(DebugAbbrev->getPatcher());
+  }
+  if (DebugInfo) {
+    DebugInfo->registerPatcher(std::make_unique<SimpleBinaryPatcher>());
+    DebugInfoPatcher =
+        static_cast<SimpleBinaryPatcher *>(DebugInfo->getPatcher());
+  }
 
   ARangesSectionWriter = std::make_unique<DebugARangesSectionWriter>();
   RangesSectionWriter = std::make_unique<DebugRangesSectionWriter>(&BC);
@@ -634,6 +640,8 @@ void DWARFRewriter::updateGdbIndexSection() {
 
 void
 DWARFRewriter::convertToRanges(const DWARFAbbreviationDeclaration *Abbrev) {
+  if (!AbbrevPatcher)
+    return;
   dwarf::Form HighPCForm = Abbrev->findAttribute(dwarf::DW_AT_high_pc)->Form;
   std::lock_guard<std::mutex> Lock(AbbrevPatcherMutex);
   AbbrevPatcher->addAttributePatch(Abbrev,
@@ -757,6 +765,8 @@ void getRangeAttrData(
 }
 
 void DWARFRewriter::patchLowHigh(DWARFDie DIE, DebugAddressRange Range) {
+  if (!DebugInfoPatcher)
+    return;
   uint64_t LowPCOffset, HighPCOffset;
   DWARFFormValue LowPCFormValue, HighPCFormValue;
   getRangeAttrData(
@@ -772,6 +782,8 @@ void DWARFRewriter::patchLowHigh(DWARFDie DIE, DebugAddressRange Range) {
 
 void DWARFRewriter::convertToRanges(DWARFDie DIE,
                                     uint64_t RangesSectionOffset) {
+  if (!DebugInfoPatcher)
+    return;
   uint64_t LowPCOffset, HighPCOffset;
   DWARFFormValue LowPCFormValue, HighPCFormValue;
   getRangeAttrData(

bolt/src/DWARFRewriter.h

@@ -28,10 +28,6 @@ class DWARFRewriter {
 
   BinaryContext &BC;
 
-  using SectionPatchersType = RewriteInstance::SectionPatchersType;
-
-  SectionPatchersType &SectionPatchers;
-
   SimpleBinaryPatcher *DebugInfoPatcher{nullptr};
 
   std::mutex DebugInfoPatcherMutex;
@@ -141,9 +137,7 @@ class DWARFRewriter {
   void flushPendingRanges();
 
 public:
-  DWARFRewriter(BinaryContext &BC,
-                SectionPatchersType &SectionPatchers)
-    : BC(BC), SectionPatchers(SectionPatchers) {}
+  DWARFRewriter(BinaryContext &BC) : BC(BC) {}
 
   /// Main function for updating the DWARF debug info.
   void updateDebugInfo();

bolt/src/RewriteInstance.cpp

@@ -19,6 +19,7 @@
 #include "DWARFRewriter.h"
 #include "DataAggregator.h"
 #include "DataReader.h"
+#include "DebugData.h"
 #include "Exceptions.h"
 #include "ExecutableFileMemoryManager.h"
 #include "MCPlusBuilder.h"
@@ -33,8 +34,8 @@
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/BinaryFormat/Magic.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
-#include "llvm/ExecutionEngine/RuntimeDyld.h"
 #include "llvm/ExecutionEngine/RTDyldMemoryManager.h"
+#include "llvm/ExecutionEngine/RuntimeDyld.h"
 #include "llvm/MC/MCAsmBackend.h"
 #include "llvm/MC/MCAsmInfo.h"
 #include "llvm/MC/MCAsmLayout.h"
@@ -67,6 +68,7 @@
 #include "llvm/Target/TargetMachine.h"
 #include <algorithm>
 #include <fstream>
+#include <llvm/Support/Error.h>
 #include <stack>
 #include <system_error>
 #include <thread>
@@ -493,7 +495,7 @@ RewriteInstance::RewriteInstance(ELFObjectFileBase *File, const int Argc,
   BAT = std::make_unique<BoltAddressTranslation>(*BC);
 
   if (opts::UpdateDebugSections)
-    DebugInfoRewriter = std::make_unique<DWARFRewriter>(*BC, SectionPatchers);
+    DebugInfoRewriter = std::make_unique<DWARFRewriter>(*BC);
 
   if (opts::Instrument) {
     BC->setRuntimeLibrary(std::make_unique<InstrumentationRuntimeLibrary>());
@@ -3126,9 +3128,12 @@ void RewriteInstance::updateSDTMarkers() {
   NamedRegionTimer T("updateSDTMarkers", "update SDT markers", TimerGroupName,
                      TimerGroupDesc, opts::TimeRewrite);
 
-  SectionPatchers[".note.stapsdt"] = std::make_unique<SimpleBinaryPatcher>();
-  auto *SDTNotePatcher = static_cast<SimpleBinaryPatcher *>(
-      SectionPatchers[".note.stapsdt"].get());
+  if (!SDTSection)
+    return;
+  SDTSection->registerPatcher(std::make_unique<SimpleBinaryPatcher>());
+
+  auto *SDTNotePatcher =
+      static_cast<SimpleBinaryPatcher *>(SDTSection->getPatcher());
   for (auto &SDTInfoKV : BC->SDTMarkers) {
     const auto OriginalAddress = SDTInfoKV.first;
     auto &SDTInfo = SDTInfoKV.second;
@@ -3170,16 +3175,11 @@ void RewriteInstance::updateLKMarkers() {
     for (auto &LKMarkerInfo : LKMarkerInfoKV.second) {
       StringRef SectionName = LKMarkerInfo.SectionName;
       SimpleBinaryPatcher *LKPatcher;
-      if (SectionPatchers.find(std::string(SectionName)) !=
-          SectionPatchers.end()) {
-        LKPatcher = static_cast<SimpleBinaryPatcher *>(
-            SectionPatchers[std::string(SectionName)].get());
-      } else {
-        SectionPatchers[std::string(SectionName)] =
-            std::make_unique<SimpleBinaryPatcher>();
-        LKPatcher = static_cast<SimpleBinaryPatcher *>(
-            SectionPatchers[std::string(SectionName)].get());
-      }
+      auto BSec = BC->getUniqueSectionByName(SectionName);
+      assert(BSec && "missing section info for kernel section");
+      if (!BSec->getPatcher())
+        BSec->registerPatcher(std::make_unique<SimpleBinaryPatcher>());
+      LKPatcher = static_cast<SimpleBinaryPatcher *>(BSec->getPatcher());
       PatchCounts[std::string(SectionName)]++;
       if (LKMarkerInfo.IsPCRelative) {
         LKPatcher->addLE32Patch(LKMarkerInfo.SectionOffset,
@@ -3684,6 +3684,7 @@ void RewriteInstance::rewriteNoteSections() {
 
     StringRef SectionName =
         cantFail(Obj.getSectionName(Section), "cannot get section name");
+    auto BSec = BC->getUniqueSectionByName(SectionName);
 
     if (shouldStrip(Section, SectionName))
       continue;
@@ -3700,10 +3701,8 @@ void RewriteInstance::rewriteNoteSections() {
       Size = Section.sh_size;
       std::string Data =
           std::string(InputFile->getData().substr(Section.sh_offset, Size));
-      auto SectionPatchersIt = SectionPatchers.find(std::string(SectionName));
-      if (SectionPatchersIt != SectionPatchers.end()) {
-        (*SectionPatchersIt->second).patchBinary(Data);
-      }
+      if (BSec && BSec->getPatcher())
+        BSec->getPatcher()->patchBinary(Data);
       OS << Data;
 
       // Add padding as the section extension might rely on the alignment.
@@ -3711,7 +3710,6 @@ void RewriteInstance::rewriteNoteSections() {
     }
 
     // Perform section post-processing.
-    auto BSec = BC->getUniqueSectionByName(SectionName);
     uint8_t *SectionData = nullptr;
     if (BSec && !BSec->isAllocatable()) {
       assert(BSec->getAlignment() <= Section.sh_addralign &&

bolt/src/RewriteInstance.h

@@ -368,9 +368,6 @@ public:
   /// Return true if the section holds linux kernel symbol information.
   static bool isKSymtabSection(StringRef SectionName);
 
-  using SectionPatchersType =
-      std::map<std::string, std::unique_ptr<BinaryPatcher>>;
-
 private:
   /// Get the contents of the LSDA section for this binary.
   ArrayRef<uint8_t> getLSDAData();
@@ -442,10 +439,6 @@ private:
 
   std::unique_ptr<BoltAddressTranslation> BAT;
 
-  /// Patchers used to apply simple changes to sections of the input binary.
-  /// Maps section name -> patcher.
-  SectionPatchersType SectionPatchers;
-
   /// Number of local symbols in newly written symbol table.
   uint64_t NumLocalSymbols{0};