[BOLT][NFCI] Migrate pseudo probes to MetadataRewriter interface

Use new MetdataRewriter interface to update pseudo probes and move
ProbeDecoder out of BinaryContext into new PseudoProbeRewriter class.

Depends on D154021

Reviewed By: Amir

Differential Revision: https://reviews.llvm.org/D154022
Differential Revision: https://reviews.llvm.org/D154023

Author: maksfb

bolt/include/bolt/Core/BinaryContext.h

@@ -29,7 +29,6 @@
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
-#include "llvm/MC/MCPseudoProbe.h"
 #include "llvm/MC/MCSectionELF.h"
 #include "llvm/MC/MCSectionMachO.h"
 #include "llvm/MC/MCStreamer.h"
@@ -681,9 +680,6 @@ class BinaryContext {
   /// and are referenced from BinaryFunction.
   std::list<std::pair<BinaryFunction *, uint64_t>> InterproceduralReferences;
 
-  /// PseudoProbe decoder
-  MCPseudoProbeDecoder ProbeDecoder;
-
   /// DWARF encoding. Available encoding types defined in BinaryFormat/Dwarf.h
   /// enum Constants, e.g. DW_EH_PE_omit.
   unsigned LSDAEncoding = dwarf::DW_EH_PE_omit;

bolt/include/bolt/Core/BinaryFunction.h

@@ -423,21 +423,6 @@ class BinaryFunction {
     return BB->getIndex();
   }
 
-  /// Return basic block that originally contained offset \p Offset
-  /// from the function start.
-  BinaryBasicBlock *getBasicBlockContainingOffset(uint64_t Offset);
-
-  const BinaryBasicBlock *getBasicBlockContainingOffset(uint64_t Offset) const {
-    return const_cast<BinaryFunction *>(this)->getBasicBlockContainingOffset(
-        Offset);
-  }
-
-  /// Return basic block that started at offset \p Offset.
-  BinaryBasicBlock *getBasicBlockAtOffset(uint64_t Offset) {
-    BinaryBasicBlock *BB = getBasicBlockContainingOffset(Offset);
-    return BB && BB->getOffset() == Offset ? BB : nullptr;
-  }
-
   /// Release memory taken by the list.
   template <typename T> BinaryFunction &clearList(T &List) {
     T TempList;
@@ -900,6 +885,21 @@ class BinaryFunction {
     return LabelToBB.lookup(Label);
   }
 
+  /// Return basic block that originally contained offset \p Offset
+  /// from the function start.
+  BinaryBasicBlock *getBasicBlockContainingOffset(uint64_t Offset);
+
+  const BinaryBasicBlock *getBasicBlockContainingOffset(uint64_t Offset) const {
+    return const_cast<BinaryFunction *>(this)->getBasicBlockContainingOffset(
+        Offset);
+  }
+
+  /// Return basic block that started at offset \p Offset.
+  BinaryBasicBlock *getBasicBlockAtOffset(uint64_t Offset) {
+    BinaryBasicBlock *BB = getBasicBlockContainingOffset(Offset);
+    return BB && BB->getOffset() == Offset ? BB : nullptr;
+  }
+
   /// Retrieve the landing pad BB associated with invoke instruction \p Invoke
   /// that is in \p BB. Return nullptr if none exists
   BinaryBasicBlock *getLandingPadBBFor(const BinaryBasicBlock &BB,

bolt/include/bolt/Rewrite/MetadataRewriters.h

@@ -19,6 +19,8 @@ class BinaryContext;
 
 // The list of rewriter build functions.
 
+std::unique_ptr<MetadataRewriter> createPseudoProbeRewriter(BinaryContext &);
+
 std::unique_ptr<MetadataRewriter> createSDTRewriter(BinaryContext &);
 
 } // namespace bolt

bolt/include/bolt/Rewrite/RewriteInstance.h

@@ -201,12 +201,6 @@ class RewriteInstance {
   /// Update LKMarkers' locations for the output binary.
   void updateLKMarkers();
 
-  /// Update address of MCDecodedPseudoProbe.
-  void updatePseudoProbes();
-
-  /// Encode MCDecodedPseudoProbe.
-  void encodePseudoProbes();
-
   /// Return the list of code sections in the output order.
   std::vector<BinarySection *> getCodeSections();
 
@@ -395,10 +389,6 @@ class RewriteInstance {
   /// of appending contents to it.
   bool willOverwriteSection(StringRef SectionName);
 
-  /// Parse .pseudo_probe_desc section and .pseudo_probe section
-  /// Setup Pseudo probe decoder
-  void parsePseudoProbe();
-
 public:
   /// Standard ELF sections we overwrite.
   static constexpr const char *SectionsToOverwrite[] = {
@@ -580,15 +570,6 @@ class RewriteInstance {
   /// .note.gnu.build-id section.
   ErrorOr<BinarySection &> BuildIDSection{std::errc::bad_address};
 
-  /// .pseudo_probe_desc section.
-  /// Contains information about pseudo probe description, like its related
-  /// function
-  ErrorOr<BinarySection &> PseudoProbeDescSection{std::errc::bad_address};
-
-  /// .pseudo_probe section.
-  /// Contains information about pseudo probe details, like its address
-  ErrorOr<BinarySection &> PseudoProbeSection{std::errc::bad_address};
-
   /// Helper for accessing sections by name.
   BinarySection *getSection(const Twine &Name) {
     ErrorOr<BinarySection &> ErrOrSection = BC->getUniqueSectionByName(Name);

bolt/lib/Rewrite/CMakeLists.txt

@@ -16,6 +16,7 @@ add_llvm_library(LLVMBOLTRewrite
   JITLinkLinker.cpp
   MachORewriteInstance.cpp
   MetadataManager.cpp
+  PseudoProbeRewriter.cpp
   RewriteInstance.cpp
   SDTRewriter.cpp