[llvm] [BOLT] Refactor interface for creating instruction patches. NFCI (PR #129404)

[Maksim Panchenko via llvm-commits]

https://github.com/maksfb created https://github.com/llvm/llvm-project/pull/129404

Add BinaryContext::createInstructionPatch() interface for patching parts of the original binary with new instruction sequences. Refactor PatchEntries pass to use the new interface.

>From f6f8c7326f4eb4be79201ad99560c165c463fa83 Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks at fb.com>
Date: Fri, 28 Feb 2025 21:16:49 -0800
Subject: [PATCH] [BOLT] Refactor interface for creating instruction patches.
 NFCI

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         | 28 +++++++++++++++++++++++++
 bolt/lib/Passes/PatchEntries.cpp        | 19 +++++++----------
 4 files changed, 46 insertions(+), 14 deletions(-)

diff --git a/bolt/include/bolt/Core/BinaryContext.h b/bolt/include/bolt/Core/BinaryContext.h
index 8bec1db70e25a..485979f1a55a1 100644
--- a/bolt/include/bolt/Core/BinaryContext.h
+++ b/bolt/include/bolt/Core/BinaryContext.h
@@ -537,6 +537,17 @@ class BinaryContext {
   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 fa6b5811a4c3b..04ec9165c2ff2 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 f9fc536f3569a..7ec16a64299dc 100644
--- a/bolt/lib/Core/BinaryContext.cpp
+++ b/bolt/lib/Core/BinaryContext.cpp
@@ -2400,6 +2400,34 @@ 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 981d1b70af907..a647dc16820ee 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;