[llvm] [BOLT] Add BinaryFunction::registerBranch(). NFC (PR #83337)

[Maksim Panchenko via llvm-commits]

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

Add an external interface to register a branch in a function that is in disassembled state. Allows to make custom modifications to the disassembler. E.g., a pre-CFG pass can add an instruction and register a branch that will later be used during the CFG construction.

>From 4601a866ae49225c30b8181340f5171e99a1e860 Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks at fb.com>
Date: Wed, 28 Feb 2024 12:57:05 -0800
Subject: [PATCH] [BOLT] Add BinaryFunction::registerBranch(). NFC

Add an external interface to register a branch in a function that is in
disassembled state. Allows to make custom modifications to the
disassembler. E.g., a pre-CFG pass can add an instruction and register
a branch that will later be used during the CFG construction.
---
 bolt/include/bolt/Core/BinaryFunction.h |  8 ++++++++
 bolt/lib/Core/BinaryFunction.cpp        | 10 ++++++++++
 2 files changed, 18 insertions(+)

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index a177178769e456..c170fa6397cc92 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -2056,6 +2056,14 @@ class BinaryFunction {
   /// Returns false if disassembly failed.
   Error disassemble();
 
+  /// An external interface to register a branch while the function is in
+  /// disassembled state. Allows to make custom modifications to the
+  /// disassembler. E.g., a pre-CFG pass can add an instruction and register
+  /// a branch that will later be used during the CFG construction.
+  ///
+  /// Return a label at the branch destination.
+  MCSymbol *registerBranch(uint64_t Src, uint64_t Dst);
+
   Error handlePCRelOperand(MCInst &Instruction, uint64_t Address,
                            uint64_t Size);
 
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 00df42c11e2239..c4dd5ef439fdc0 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -1445,6 +1445,16 @@ Error BinaryFunction::disassemble() {
   return Error::success();
 }
 
+MCSymbol *BinaryFunction::registerBranch(uint64_t Src, uint64_t Dst) {
+  assert(CurrentState == Disassembled &&
+         "Cannot register branch unless function is in disassembled state.");
+  assert(containsAddress(Src) && containsAddress(Dst) &&
+         "Cannot register external branch.");
+  MCSymbol *Target = getOrCreateLocalLabel(Dst);
+  TakenBranches.emplace_back(Src - getAddress(), Dst - getAddress());
+  return Target;
+}
+
 bool BinaryFunction::scanExternalRefs() {
   bool Success = true;
   bool DisassemblyFailed = false;