From 5a11912ece2731eb9c50f80fdfd75bd1dfc2ebc8 Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks@fb.com>
Date: Sat, 1 Mar 2025 19:20:17 -0800
Subject: [PATCH] [BOLT] Refactor interface for creating instruction patches.
 NFCI (#129404)

Add BinaryContext::createInstructionPatch() interface for patching parts
of the original binary with new instruction sequences. Refactor
PatchEntries pass to use the new interface.
---
 bolt/include/bolt/Core/BinaryContext.h  | 11 +++++++++++
 bolt/include/bolt/Passes/PatchEntries.h |  2 --
 bolt/lib/Core/BinaryContext.cpp         | 26 +++++++++++++++++++++++++
 bolt/lib/Passes/PatchEntries.cpp        | 19 +++++++-----------
 4 files changed, 44 insertions(+), 14 deletions(-)

diff --git a/bolt/include/bolt/Core/BinaryContext.h b/bolt/include/bolt/Core/BinaryContext.h
index 8bec1db70e25..485979f1a55a 100644
--- a/bolt/include/bolt/Core/BinaryContext.h
+++ b/bolt/include/bolt/Core/BinaryContext.h
@@ -537,6 +537,17 @@ public:
   BinaryFunction *createInjectedBinaryFunction(const std::string &Name,
                                                bool IsSimple = true);
 
+  /// Patch the original binary contents at address \p Address with a sequence
+  /// of instructions from the \p Instructions list. The callee is responsible
+  /// for checking that the sequence doesn't cross any function or section
+  /// boundaries.
+  ///
+  /// Optional \p Name can be assigned to the patch. The name will be emitted to
+  /// the symbol table at \p Address.
+  BinaryFunction *createInstructionPatch(uint64_t Address,
+                                         InstructionListType &Instructions,
+                                         const Twine &Name = "");
+
   std::vector<BinaryFunction *> &getInjectedBinaryFunctions() {
     return InjectedBinaryFunctions;
   }
diff --git a/bolt/include/bolt/Passes/PatchEntries.h b/bolt/include/bolt/Passes/PatchEntries.h
index fa6b5811a4c3..04ec9165c2ff 100644
--- a/bolt/include/bolt/Passes/PatchEntries.h
+++ b/bolt/include/bolt/Passes/PatchEntries.h
@@ -26,8 +26,6 @@ class PatchEntries : public BinaryFunctionPass {
   struct Patch {
     const MCSymbol *Symbol;
     uint64_t Address;
-    uint64_t FileOffset;
-    BinarySection *Section;
   };
 
 public:
diff --git a/bolt/lib/Core/BinaryContext.cpp b/bolt/lib/Core/BinaryContext.cpp
index f9fc536f3569..8fa1e367c685 100644
--- a/bolt/lib/Core/BinaryContext.cpp
+++ b/bolt/lib/Core/BinaryContext.cpp
@@ -2400,6 +2400,32 @@ BinaryContext::createInjectedBinaryFunction(const std::string &Name,
   return BF;
 }
 
+BinaryFunction *BinaryContext::createInstructionPatch(
+    uint64_t Address, InstructionListType &Instructions, const Twine &Name) {
+  ErrorOr<BinarySection &> Section = getSectionForAddress(Address);
+  assert(Section && "cannot get section for patching");
+  assert(Section->hasSectionRef() && Section->isText() &&
+         "can only patch input file code sections");
+
+  const uint64_t FileOffset =
+      Section->getInputFileOffset() + Address - Section->getAddress();
+
+  std::string PatchName = Name.str();
+  if (PatchName.empty()) {
+    // Assign unique name to the patch.
+    static uint64_t N = 0;
+    PatchName = "__BP_" + std::to_string(N++);
+  }
+
+  BinaryFunction *PBF = createInjectedBinaryFunction(PatchName);
+  PBF->setOutputAddress(Address);
+  PBF->setFileOffset(FileOffset);
+  PBF->setOriginSection(&Section.get());
+  PBF->addBasicBlock()->addInstructions(Instructions);
+
+  return PBF;
+}
+
 std::pair<size_t, size_t>
 BinaryContext::calculateEmittedSize(BinaryFunction &BF, bool FixBranches) {
   // Adjust branch instruction to match the current layout.
diff --git a/bolt/lib/Passes/PatchEntries.cpp b/bolt/lib/Passes/PatchEntries.cpp
index 981d1b70af90..a647dc16820e 100644
--- a/bolt/lib/Passes/PatchEntries.cpp
+++ b/bolt/lib/Passes/PatchEntries.cpp
@@ -83,9 +83,8 @@ Error PatchEntries::runOnFunctions(BinaryContext &BC) {
         return false;
       }
 
-      PendingPatches.emplace_back(Patch{Symbol, Function.getAddress() + Offset,
-                                        Function.getFileOffset() + Offset,
-                                        Function.getOriginSection()});
+      PendingPatches.emplace_back(
+          Patch{Symbol, Function.getAddress() + Offset});
       NextValidByte = Offset + PatchSize;
       if (NextValidByte > Function.getMaxSize()) {
         if (opts::Verbosity >= 1)
@@ -118,16 +117,12 @@ Error PatchEntries::runOnFunctions(BinaryContext &BC) {
     }
 
     for (Patch &Patch : PendingPatches) {
-      BinaryFunction *PatchFunction = BC.createInjectedBinaryFunction(
+      // Add instruction patch to the binary.
+      InstructionListType Instructions;
+      BC.MIB->createLongTailCall(Instructions, Patch.Symbol, BC.Ctx.get());
+      BinaryFunction *PatchFunction = BC.createInstructionPatch(
+          Patch.Address, Instructions,
           NameResolver::append(Patch.Symbol->getName(), ".org.0"));
-      // Force the function to be emitted at the given address.
-      PatchFunction->setOutputAddress(Patch.Address);
-      PatchFunction->setFileOffset(Patch.FileOffset);
-      PatchFunction->setOriginSection(Patch.Section);
-
-      InstructionListType Seq;
-      BC.MIB->createLongTailCall(Seq, Patch.Symbol, BC.Ctx.get());
-      PatchFunction->addBasicBlock()->addInstructions(Seq);
 
       // Verify the size requirements.
       uint64_t HotSize, ColdSize;