[llvm] [BOLT][AArch64] Expand cmpbr when reversing would overflow (PR #202998)

Alexandros Lamprineas via llvm-commits llvm-commits at lists.llvm.org
Wed Aug 5 07:44:58 PDT 2026


https://github.com/labrinea updated https://github.com/llvm/llvm-project/pull/202998

>From 754ab44a4e7fa5304fbada343cc8c172c8579f7b Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Thu, 4 Jun 2026 10:57:50 +0100
Subject: [PATCH 1/9] [BOLT][AArch64] Transform cmpbr ~> cmp + br when
 inversion not possible (#185731)

When reordering blocks we may have to invert branches. Sometimes this
isn't possible for compare-and-branch instructions because the immediate
value would overflow/underflow after the adjustment. In such cases I am
splitting the instruction into a compare followed by a branch. For this
to be legal we should be sure that the condition flags are not being
clobbered. Liveness analysis may help here.
---
 bolt/include/bolt/Core/BinaryFunction.h       |   4 +-
 bolt/include/bolt/Core/MCPlusBuilder.h        |  11 +-
 bolt/include/bolt/Passes/LongJmp.h            |   4 +-
 bolt/include/bolt/Utils/CommandLineOpts.h     |   4 +
 bolt/lib/Core/BinaryFunction.cpp              |   7 +-
 bolt/lib/Passes/BinaryPasses.cpp              |  16 +-
 bolt/lib/Passes/LongJmp.cpp                   |  27 +++-
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   |  90 +++++++++--
 bolt/lib/Target/AArch64/CMakeLists.txt        |   6 +-
 bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp  |   6 +-
 bolt/lib/Target/X86/X86MCPlusBuilder.cpp      |   6 +-
 bolt/lib/Utils/CommandLineOpts.cpp            |   6 +
 .../AArch64/compare-and-branch-inversion.S    | 118 ++++++++++-----
 bolt/unittests/Core/CMakeLists.txt            |   1 +
 bolt/unittests/Core/MCPlusBuilder.cpp         | 140 ++++++++++++++++--
 15 files changed, 365 insertions(+), 81 deletions(-)

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 84fbd5661fd0a..51651e41497da 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -65,6 +65,8 @@ class DWARFUnit;
 
 namespace bolt {
 
+class DataflowInfoManager;
+
 using InputOffsetToAddressMapTy = std::unordered_multimap<uint64_t, uint64_t>;
 
 /// Types of macro-fusion alignment corrections.
@@ -2463,7 +2465,7 @@ class BinaryFunction {
   /// while the second successor - false/fall-through branch.
   ///
   /// When we reverse the branch condition, the CFG is updated accordingly.
-  void fixBranches();
+  void fixBranches(DataflowInfoManager *DIM = nullptr);
 
   /// Mark function as finalized. No further optimizations are permitted.
   void setFinalized() { CurrentState = State::CFG_Finalized; }
diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 84b44a9ab5483..c98b56d14523c 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -52,6 +52,7 @@ namespace bolt {
 class BinaryBasicBlock;
 class BinaryContext;
 class BinaryFunction;
+class DataflowInfoManager;
 
 /// Different types of indirect branches encountered during disassembly.
 enum class IndirectBranchType : char {
@@ -475,7 +476,8 @@ class MCPlusBuilder {
   }
 
   /// Check whether this conditional branch can be reversed
-  virtual bool isReversibleBranch(const MCInst &Inst) const {
+  virtual bool isReversibleBranch(const MCInst &Inst,
+                                  DataflowInfoManager *DIM = nullptr) const {
     assert(!isUnsupportedInstruction(Inst) && isConditionalBranch(Inst) &&
            "Instruction is not known conditional branch");
 
@@ -2141,8 +2143,11 @@ class MCPlusBuilder {
   }
 
   /// Reverses the branch condition in Inst and update its taken target to TBB.
-  virtual void reverseBranchCondition(MCInst &Inst, const MCSymbol *TBB,
-                                      MCContext *Ctx) const {
+  /// Assumes that the branch is reversible.
+  virtual void
+  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                         const MCSymbol *TBB, MCContext *Ctx,
+                         DataflowInfoManager *DIM = nullptr) const {
     llvm_unreachable("not implemented");
   }
 
diff --git a/bolt/include/bolt/Passes/LongJmp.h b/bolt/include/bolt/Passes/LongJmp.h
index 4633d30104d43..ccdebf9e1ed0a 100644
--- a/bolt/include/bolt/Passes/LongJmp.h
+++ b/bolt/include/bolt/Passes/LongJmp.h
@@ -14,6 +14,8 @@
 namespace llvm {
 namespace bolt {
 
+class DataflowInfoManager;
+
 /// LongJmp is veneer-insertion pass originally written for AArch64 that
 /// compensates for its short-range branches, typically done during linking. We
 /// pull this pass inside BOLT because here we can do a better job at stub
@@ -74,7 +76,7 @@ class LongJmpPass : public BinaryFunctionPass {
   /// Relax all internal function branches including those between fragments.
   /// Assume that fragments are placed in different sections but are within
   /// 128MB of each other.
-  void relaxLocalBranches(BinaryFunction &BF);
+  void relaxLocalBranches(BinaryFunction &BF, DataflowInfoManager *DIM);
 
   ///                 -- Layout estimation methods --
   /// Try to do layout before running the emitter, by looking at BinaryFunctions
diff --git a/bolt/include/bolt/Utils/CommandLineOpts.h b/bolt/include/bolt/Utils/CommandLineOpts.h
index e11b18d3489cf..88b56217f9512 100644
--- a/bolt/include/bolt/Utils/CommandLineOpts.h
+++ b/bolt/include/bolt/Utils/CommandLineOpts.h
@@ -132,6 +132,10 @@ extern llvm::cl::opt<bool> UpdateDebugSections;
 // dbgs() for output within DEBUG().
 extern llvm::cl::opt<unsigned> Verbosity;
 
+// Option to control whether liveness analysis should be used by
+// FixupBranches and LongJmpPass. Needed for branch inversion on AArch64.
+extern llvm::cl::opt<bool> LivenessAnalysis;
+
 /// Return true if we should process all functions in the binary.
 bool processAllFunctions();
 
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 200e286d8e80e..ddfd163be84ec 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3650,7 +3650,7 @@ bool BinaryFunction::validateCFG() const {
   return true;
 }
 
-void BinaryFunction::fixBranches() {
+void BinaryFunction::fixBranches(DataflowInfoManager *DIM) {
   assert(isSimple() && "Expected function with valid CFG.");
 
   auto &MIB = BC.MIB;
@@ -3709,7 +3709,7 @@ void BinaryFunction::fixBranches() {
 
       // Reverse branch condition and swap successors.
       auto swapSuccessors = [&]() {
-        if (!MIB->isReversibleBranch(*CondBranch)) {
+        if (!MIB->isReversibleBranch(*CondBranch, DIM)) {
           if (opts::Verbosity) {
             BC.outs() << "BOLT-INFO: unable to swap successors in " << *this
                       << '\n';
@@ -3719,7 +3719,8 @@ void BinaryFunction::fixBranches() {
         std::swap(TSuccessor, FSuccessor);
         BB->swapConditionalSuccessors();
         auto L = BC.scopeLock();
-        MIB->reverseBranchCondition(*CondBranch, TSuccessor->getLabel(), Ctx);
+        MIB->reverseBranchCondition(BB, *CondBranch, TSuccessor->getLabel(),
+                                    Ctx, DIM);
         return true;
       };
 
diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index adf2bbae52d11..55b88233dcbde 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -11,8 +11,10 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/BinaryPasses.h"
+#include "bolt/Core/BinaryFunctionCallGraph.h"
 #include "bolt/Core/FunctionLayout.h"
 #include "bolt/Core/ParallelUtilities.h"
+#include "bolt/Passes/DataflowInfoManager.h"
 #include "bolt/Passes/ReorderAlgorithm.h"
 #include "bolt/Passes/ReorderFunctions.h"
 #include "bolt/Utils/CommandLineOpts.h"
@@ -545,12 +547,22 @@ bool ReorderBasicBlocks::modifyFunctionLayout(BinaryFunction &BF,
 }
 
 Error FixupBranches::runOnFunctions(BinaryContext &BC) {
+  std::unique_ptr<BinaryFunctionCallGraph> CG;
+  std::unique_ptr<RegAnalysis> RA;
+  std::unique_ptr<DataflowInfoManager> DIM;
+
+  if (opts::LivenessAnalysis) {
+    CG = std::make_unique<BinaryFunctionCallGraph>(buildCallGraph(BC));
+    RA = std::make_unique<RegAnalysis>(BC, &BC.getBinaryFunctions(), CG.get());
+  }
   for (auto &It : BC.getBinaryFunctions()) {
     BinaryFunction &Function = It.second;
     if (!BC.shouldEmit(Function) || !Function.isSimple())
       continue;
 
-    Function.fixBranches();
+    if (opts::LivenessAnalysis)
+      DIM = std::make_unique<DataflowInfoManager>(Function, RA.get(), nullptr);
+    Function.fixBranches(DIM.get());
   }
   return Error::success();
 }
@@ -961,7 +973,7 @@ uint64_t SimplifyConditionalTailCalls::fixTailCalls(BinaryFunction &BF) {
       uint64_t Count = 0;
       if (CondSucc != BB) {
         // Patch the new target address into the conditional branch.
-        MIB->reverseBranchCondition(*CondBranch, CalleeSymbol, Ctx);
+        MIB->reverseBranchCondition(PredBB, *CondBranch, CalleeSymbol, Ctx);
         // Since we reversed the condition on the branch we need to change
         // the target for the unconditional branch or add a unconditional
         // branch to the old target.  This has to be done manually since
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index b771e6a8b120a..03d5aa4bf3e2c 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -11,7 +11,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/LongJmp.h"
+#include "bolt/Core/BinaryFunctionCallGraph.h"
 #include "bolt/Core/ParallelUtilities.h"
+#include "bolt/Passes/DataflowInfoManager.h"
 #include "bolt/Utils/CommandLineOpts.h"
 #include "llvm/Support/MathExtras.h"
 
@@ -662,7 +664,8 @@ Error LongJmpPass::relax(BinaryFunction &Func, bool &Modified) {
   return Error::success();
 }
 
-void LongJmpPass::relaxLocalBranches(BinaryFunction &BF) {
+void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
+                                     DataflowInfoManager *DIM) {
   BinaryContext &BC = BF.getBinaryContext();
   auto &MIB = BC.MIB;
 
@@ -832,7 +835,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF) {
       // If the other successor is a fall-through, invert the condition code.
       BinaryBasicBlock *NextBB =
           BF->getLayout().getBasicBlockAfter(BB, /*IgnoreSplits*/ false);
-      bool IsReversibleBranch = MIB->isReversibleBranch(Inst);
+      bool IsReversibleBranch = MIB->isReversibleBranch(Inst, DIM);
       bool ShouldReverseBranch = BB->getConditionalSuccessor(false) == NextBB;
 
       // Create a trampoline basic block for the fall-through target of the
@@ -850,7 +853,8 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF) {
       if (ShouldReverseBranch && IsReversibleBranch) {
         BB->swapConditionalSuccessors();
         auto L = BC.scopeLock();
-        MIB->reverseBranchCondition(Inst, NextBB->getLabel(), BC.Ctx.get());
+        MIB->reverseBranchCondition(BB, Inst, NextBB->getLabel(), BC.Ctx.get(),
+                                    DIM);
       } else {
         auto L = BC.scopeLock();
         MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(), BC.Ctx.get());
@@ -935,12 +939,23 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
           opts::SplitStrategy != opts::SplitFunctionsStrategy::CDSplit) &&
          "LongJmp cannot work with functions split in more than two fragments");
 
+  std::unique_ptr<BinaryFunctionCallGraph> CG;
+  std::unique_ptr<RegAnalysis> RA;
+  std::unique_ptr<DataflowInfoManager> DIM;
+
+  if (opts::LivenessAnalysis) {
+    CG = std::make_unique<BinaryFunctionCallGraph>(buildCallGraph(BC));
+    RA = std::make_unique<RegAnalysis>(BC, &BC.getBinaryFunctions(), CG.get());
+  }
+
   if (opts::CompactCodeModel) {
     BC.outs()
         << "BOLT-INFO: relaxing branches for compact code model (<128MB)\n";
 
     ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
-      relaxLocalBranches(BF);
+      if (opts::LivenessAnalysis)
+        DIM = std::make_unique<DataflowInfoManager>(BF, RA.get(), nullptr);
+      relaxLocalBranches(BF, DIM.get());
     };
 
     ParallelUtilities::PredicateTy SkipPredicate =
@@ -965,12 +980,14 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
     tentativeLayout(BC, Sorted);
     updateStubGroups();
     for (BinaryFunction *Func : Sorted) {
+      if (opts::LivenessAnalysis)
+        DIM = std::make_unique<DataflowInfoManager>(*Func, RA.get(), nullptr);
       if (auto E = relax(*Func, Modified))
         return Error(std::move(E));
       // Don't ruin non-simple functions, they can't afford to have the layout
       // changed.
       if (Modified && Func->isSimple())
-        Func->fixBranches();
+        Func->fixBranches(DIM.get());
     }
   } while (Modified);
   BC.outs() << "BOLT-INFO: Inserted " << NumHotStubs
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index b4a9dff9d25b2..ecb7bac9d7796 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -22,6 +22,7 @@
 #include "bolt/Core/BinaryFunction.h"
 #include "bolt/Core/MCInstUtils.h"
 #include "bolt/Core/MCPlusBuilder.h"
+#include "bolt/Passes/DataflowInfoManager.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCDisassembler/MCDisassembler.h"
@@ -2045,6 +2046,25 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     exit(1);
   }
 
+  unsigned getInvertedCC(unsigned Opcode) const {
+    // clang-format off
+    switch (Opcode) {
+    default:
+      llvm_unreachable("Failed to invert condition code");
+      return Opcode;
+    // Compare register with immediate and branch.
+    case AArch64::CBGTWri:  return AArch64CC::LE;
+    case AArch64::CBGTXri:  return AArch64CC::LE;
+    case AArch64::CBLTWri:  return AArch64CC::GE;
+    case AArch64::CBLTXri:  return AArch64CC::GE;
+    case AArch64::CBHIWri:  return AArch64CC::LS;
+    case AArch64::CBHIXri:  return AArch64CC::LS;
+    case AArch64::CBLOWri:  return AArch64CC::HS;
+    case AArch64::CBLOXri:  return AArch64CC::HS;
+    }
+    // clang-format on
+  }
+
   unsigned getInvertedBranchOpcode(unsigned Opcode) const {
     // clang-format off
     switch (Opcode) {
@@ -2171,38 +2191,78 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     }
   }
 
-  bool isReversibleBranch(const MCInst &Inst) const override {
+  bool isReversibleBranch(const MCInst &Inst,
+                          DataflowInfoManager *DIM = nullptr) const override {
     if (isCompAndBranch(Inst)) {
+      bool MayClobberFlags =
+          DIM ? DIM->getLivenessAnalysis().getLiveIn(Inst).test(getFlagsReg())
+              : true;
       unsigned InvertedOpcode = getInvertedBranchOpcode(Inst.getOpcode());
-      if (needsImmDec(InvertedOpcode) && Inst.getOperand(1).getImm() == 0)
+      if (needsImmDec(InvertedOpcode) && Inst.getOperand(1).getImm() == 0 &&
+          MayClobberFlags)
         return false;
-      if (needsImmInc(InvertedOpcode) && Inst.getOperand(1).getImm() == 63)
+      if (needsImmInc(InvertedOpcode) && Inst.getOperand(1).getImm() == 63 &&
+          MayClobberFlags)
         return false;
     }
     return MCPlusBuilder::isReversibleBranch(Inst);
   }
 
-  void reverseBranchCondition(MCInst &Inst, const MCSymbol *TBB,
-                              MCContext *Ctx) const override {
-    if (!isReversibleBranch(Inst)) {
-      errs() << "BOLT-ERROR: Cannot reverse branch " << Inst << "\n";
-      exit(1);
-    }
+  void
+  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                         const MCSymbol *TBB, MCContext *Ctx,
+                         DataflowInfoManager *DIM = nullptr) const override {
+    assert(isReversibleBranch(Inst, DIM) && "Irreversible branch");
 
     if (isTB(Inst) || isCB(Inst) || isCompAndBranch(Inst)) {
+      bool ImmediateOutOfBounds = false;
       unsigned InvertedOpcode = getInvertedBranchOpcode(Inst.getOpcode());
-      Inst.setOpcode(InvertedOpcode);
-      assert(Inst.getOpcode() != 0 && "Invalid branch instruction");
+      assert(InvertedOpcode != 0 && "Invalid branch instruction");
       // The FEAT_CMPBR compare-and-branch instructions cannot encode all
       // the possible condition codes, therefore we either have to adjust
       // the immediate value by +-1, or to swap the register operands
       // when reversing the branch condition.
       if (needsRegSwap(InvertedOpcode))
         std::swap(Inst.getOperand(0), Inst.getOperand(1));
-      else if (needsImmDec(InvertedOpcode))
-        Inst.getOperand(1).setImm(Inst.getOperand(1).getImm() - 1);
-      else if (needsImmInc(InvertedOpcode))
-        Inst.getOperand(1).setImm(Inst.getOperand(1).getImm() + 1);
+      else if (needsImmDec(InvertedOpcode)) {
+        if (Inst.getOperand(1).getImm() == 0)
+          ImmediateOutOfBounds = true;
+        else
+          Inst.getOperand(1).setImm(Inst.getOperand(1).getImm() - 1);
+      } else if (needsImmInc(InvertedOpcode)) {
+        if (Inst.getOperand(1).getImm() == 63)
+          ImmediateOutOfBounds = true;
+        else
+          Inst.getOperand(1).setImm(Inst.getOperand(1).getImm() + 1);
+      }
+      if (ImmediateOutOfBounds) {
+        auto is32BitVariant = [](unsigned Opcode) {
+          switch (Opcode) {
+          default:
+            return false;
+          case AArch64::CBGTWri:
+          case AArch64::CBLTWri:
+          case AArch64::CBHIWri:
+          case AArch64::CBLOWri:
+            return true;
+          }
+        };
+        InstructionListType Code;
+        MCInstBuilder Cmp =
+            is32BitVariant(InvertedOpcode)
+                ? MCInstBuilder(AArch64::SUBSWri).addReg(AArch64::WZR)
+                : MCInstBuilder(AArch64::SUBSXri).addReg(AArch64::XZR);
+        Cmp.addReg(Inst.getOperand(0).getReg())
+            .addImm(Inst.getOperand(1).getImm())
+            .addImm(0);
+        Code.emplace_back(std::move(Cmp));
+        Code.emplace_back(MCInstBuilder(AArch64::Bcc)
+                              .addImm(getInvertedCC(Inst.getOpcode()))
+                              .addExpr(MCSymbolRefExpr::create(TBB, *Ctx)));
+        Parent->replaceInstruction(Parent->findInstruction(&Inst), Code);
+        return;
+      }
+      Inst.setOpcode(InvertedOpcode);
     } else if (Inst.getOpcode() == AArch64::Bcc) {
       Inst.getOperand(0).setImm(AArch64CC::getInvertedCondCode(
           static_cast<AArch64CC::CondCode>(Inst.getOperand(0).getImm())));
diff --git a/bolt/lib/Target/AArch64/CMakeLists.txt b/bolt/lib/Target/AArch64/CMakeLists.txt
index 1e171748aece6..e28ed0bd66ba1 100644
--- a/bolt/lib/Target/AArch64/CMakeLists.txt
+++ b/bolt/lib/Target/AArch64/CMakeLists.txt
@@ -29,7 +29,11 @@ add_llvm_library(LLVMBOLTTargetAArch64
   AArch64CommonTableGen
   )
 
-target_link_libraries(LLVMBOLTTargetAArch64 PRIVATE LLVMBOLTCore LLVMBOLTUtils)
+target_link_libraries(LLVMBOLTTargetAArch64 PRIVATE
+  LLVMBOLTCore
+  LLVMBOLTPasses
+  LLVMBOLTUtils
+  )
 
 include_directories(
   ${LLVM_MAIN_SRC_DIR}/lib/Target/AArch64
diff --git a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
index d1a0572277874..b954a392c4f26 100644
--- a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
+++ b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
@@ -162,8 +162,10 @@ class RISCVMCPlusBuilder : public MCPlusBuilder {
     }
   }
 
-  void reverseBranchCondition(MCInst &Inst, const MCSymbol *TBB,
-                              MCContext *Ctx) const override {
+  void
+  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                         const MCSymbol *TBB, MCContext *Ctx,
+                         DataflowInfoManager *DIM = nullptr) const override {
     auto Opcode = getInvertedBranchOpcode(Inst.getOpcode());
     Inst.setOpcode(Opcode);
     replaceBranchTarget(Inst, TBB, Ctx);
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index 11a297f514530..88dcf994e5dc7 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -2811,8 +2811,10 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     Inst.addOperand(MCOperand::createImm(CC));
   }
 
-  void reverseBranchCondition(MCInst &Inst, const MCSymbol *TBB,
-                              MCContext *Ctx) const override {
+  void
+  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                         const MCSymbol *TBB, MCContext *Ctx,
+                         DataflowInfoManager *DIM = nullptr) const override {
     unsigned InvCC = getInvertedCondCode(getCondCode(Inst));
     assert(InvCC != X86::COND_INVALID && "invalid branch instruction");
     Inst.getOperand(Info->get(Inst.getOpcode()).NumOperands - 1).setImm(InvCC);
diff --git a/bolt/lib/Utils/CommandLineOpts.cpp b/bolt/lib/Utils/CommandLineOpts.cpp
index 20b24c3b4acc5..82fde5788397a 100644
--- a/bolt/lib/Utils/CommandLineOpts.cpp
+++ b/bolt/lib/Utils/CommandLineOpts.cpp
@@ -365,6 +365,12 @@ cl::opt<unsigned>
               cl::init(0), cl::ZeroOrMore, cl::cat(BoltCategory),
               cl::sub(cl::SubCommand::getAll()));
 
+cl::opt<bool> LivenessAnalysis(
+    "liveness-analysis",
+    cl::desc("use liveness analysis in FixupBranches and LongJmpPass"
+             "(needed for branch inversion on AArch64)"),
+    cl::init(false), cl::cat(BoltCategory));
+
 bool processAllFunctions() {
   if (opts::AggregateOnly)
     return false;
diff --git a/bolt/test/AArch64/compare-and-branch-inversion.S b/bolt/test/AArch64/compare-and-branch-inversion.S
index 28167416c31cb..0ea084e904101 100644
--- a/bolt/test/AArch64/compare-and-branch-inversion.S
+++ b/bolt/test/AArch64/compare-and-branch-inversion.S
@@ -1,18 +1,24 @@
 # This test checks that branch inversion works when reordering blocks which
 # contain short range conditional branches. Handles edge cases, like when
 # the immediate value is the upper or lower allowed value in which case the
-# transformation bails.
+# transformation bails. If liveness analysis proves that the condition flags
+# are dead we can replace the branch with cmp + b.cc
 
 # REQUIRES: system-linux, asserts
 
 # RUN: %clang %cflags -march=armv9-a+cmpbr -Wl,-q %s -o %t
 # RUN: link_fdata --no-lbr %s %t %t.fdata
 # RUN: llvm-strip --strip-unneeded %t
+#
 # RUN: llvm-bolt -v=1 %t -o %t.bolt --data %t.fdata --reorder-blocks=ext-tsp --compact-code-model \
-# RUN:   | FileCheck %s --check-prefix=BOLT-INFO
-# RUN: llvm-objdump -d %t.bolt | FileCheck %s
+# RUN:   | FileCheck %s --check-prefix=BOLT-INFO-NO-LIVENESS
+# RUN: llvm-objdump -d %t.bolt | FileCheck %s --check-prefix=COMMON --check-prefix=NO-LIVENESS
+#
+# RUN: llvm-bolt -v=1 %t -o %t.bolt --data %t.fdata --reorder-blocks=ext-tsp --compact-code-model \
+# RUN:   --liveness-analysis | FileCheck %s --check-prefix=BOLT-INFO-LIVENESS
+# RUN: llvm-objdump -d %t.bolt | FileCheck %s --check-prefix=COMMON --check-prefix=LIVENESS
 
-# CHECK: Disassembly of section .text:
+# COMMON: Disassembly of section .text:
 
   .globl immediate_increment
   .type immediate_increment, %function
@@ -29,12 +35,12 @@ immediate_increment:
     mov x0, #2
     ret
 
-# CHECK: <immediate_increment>:
-# CHECK-NEXT:            {{.*}} cblt x0, #0x1, 0x[[ADDR0:[0-9a-f]+]] <{{.*}}>
-# CHECK-NEXT:            {{.*}} mov  x0, #0x2 // =2
-# CHECK-NEXT:            {{.*}} ret
-# CHECK-NEXT: [[ADDR0]]: {{.*}} mov x0, #0x1 // =1
-# CHECK-NEXT:            {{.*}} ret
+# COMMON: <immediate_increment>:
+# COMMON-NEXT:            {{.*}} cblt x0, #0x1, 0x[[ADDR0:[0-9a-f]+]] <{{.*}}>
+# COMMON-NEXT:            {{.*}} mov  x0, #0x2 // =2
+# COMMON-NEXT:            {{.*}} ret
+# COMMON-NEXT: [[ADDR0]]: {{.*}} mov x0, #0x1 // =1
+# COMMON-NEXT:            {{.*}} ret
 
   .globl immediate_decrement
   .type immediate_decrement, %function
@@ -51,12 +57,12 @@ immediate_decrement:
     mov x0, #2
     ret
 
-# CHECK: <immediate_decrement>:
-# CHECK-NEXT:            {{.*}} cbhi x0, #0x0, 0x[[ADDR1:[0-9a-f]+]] <{{.*}}>
-# CHECK-NEXT:            {{.*}} mov  x0, #0x2 // =2
-# CHECK-NEXT:            {{.*}} ret
-# CHECK-NEXT: [[ADDR1]]: {{.*}} mov x0, #0x1 // =1
-# CHECK-NEXT:            {{.*}} ret
+# COMMON: <immediate_decrement>:
+# COMMON-NEXT:            {{.*}} cbhi x0, #0x0, 0x[[ADDR1:[0-9a-f]+]] <{{.*}}>
+# COMMON-NEXT:            {{.*}} mov  x0, #0x2 // =2
+# COMMON-NEXT:            {{.*}} ret
+# COMMON-NEXT: [[ADDR1]]: {{.*}} mov x0, #0x1 // =1
+# COMMON-NEXT:            {{.*}} ret
 
   .globl register_swap
   .type register_swap, %function
@@ -73,37 +79,77 @@ register_swap:
     mov x0, #2
     ret
 
-# CHECK: <register_swap>:
-# CHECK-NEXT:            {{.*}} cbgt x1, x0, 0x[[ADDR2:[0-9a-f]+]] <{{.*}}>
-# CHECK-NEXT:            {{.*}} mov  x0, #0x2 // =2
-# CHECK-NEXT:            {{.*}} ret
-# CHECK-NEXT: [[ADDR2]]: {{.*}} mov x0, #0x1 // =1
-# CHECK-NEXT:            {{.*}} ret
+# COMMON: <register_swap>:
+# COMMON-NEXT:            {{.*}} cbgt x1, x0, 0x[[ADDR2:[0-9a-f]+]] <{{.*}}>
+# COMMON-NEXT:            {{.*}} mov  x0, #0x2 // =2
+# COMMON-NEXT:            {{.*}} ret
+# COMMON-NEXT: [[ADDR2]]: {{.*}} mov x0, #0x1 // =1
+# COMMON-NEXT:            {{.*}} ret
 
-  .globl irreversible
-  .type irreversible, %function
-irreversible:
+  .globl immediate_overflow
+  .type immediate_overflow, %function
+immediate_overflow:
 .entry3:
-# FDATA: 1 irreversible #.entry3# 10
+# FDATA: 1 immediate_overflow #.entry3# 10
     cbgt x0, #63, .exit3
 .cold3:
-# FDATA: 1 irreversible #.cold3# 1
+# FDATA: 1 immediate_overflow #.cold3# 1
     mov x0, #1
     ret
 .exit3:
-# FDATA: 1 irreversible #.exit3# 10
+# FDATA: 1 immediate_overflow #.exit3# 10
+    mov x0, #2
+    ret
+
+# BOLT-INFO-NO-LIVENESS: unable to swap successors in immediate_overflow
+#
+# Without liveness the blocks get reordered, but since the branch is
+# irreversible an additional unconditional branch is emitted.
+# This codegen is suboptimal yet correct.
+#
+# NO-LIVENESS: <immediate_overflow>:
+# NO-LIVENESS-NEXT:            {{.*}} cbgt x0, #0x3f, 0x[[ADDR3:[0-9a-f]+]] <{{.*}}>
+# NO-LIVENESS-NEXT:            {{.*}} b               0x[[ADDR4:[0-9a-f]+]] <{{.*}}>
+# NO-LIVENESS-NEXT: [[ADDR3]]: {{.*}} mov  x0, #0x2 // =2
+# NO-LIVENESS-NEXT:            {{.*}} ret
+# NO-LIVENESS-NEXT: [[ADDR4]]: {{.*}} mov x0, #0x1 // =1
+# NO-LIVENESS-NEXT:            {{.*}} ret
+
+# LIVENESS: <immediate_overflow>:
+# LIVENESS-NEXT:            {{.*}} cmp  x0, #0x3f
+# LIVENESS-NEXT:            {{.*}} b.le 0x[[ADDR5:[0-9a-f]+]] <{{.*}}>
+# LIVENESS-NEXT:            {{.*}} mov  x0, #0x2 // =2
+# LIVENESS-NEXT:            {{.*}} ret
+# LIVENESS-NEXT: [[ADDR5]]: {{.*}} mov x0, #0x1 // =1
+# LIVENESS-NEXT:            {{.*}} ret
+
+  .globl irreversible
+  .type irreversible, %function
+irreversible:
+.entry4:
+# FDATA: 1 irreversible #.entry4# 10
+    cmp x0, #63
+    cbgt x0, #63, .exit4
+.cold4:
+# FDATA: 1 irreversible #.cold4# 1
+    csel x0, x1, x2, le
+    ret
+.exit4:
+# FDATA: 1 irreversible #.exit4# 10
     mov x0, #2
     ret
 
-# BOLT-INFO: unable to swap successors in irreversible
+# BOLT-INFO-NO-LIVENESS: unable to swap successors in irreversible
+# BOLT-INFO-LIVENESS:    unable to swap successors in irreversible
 
-# CHECK: <irreversible>:
-# CHECK-NEXT:            {{.*}} cbgt x0, #0x3f, 0x[[ADDR3:[0-9a-f]+]] <{{.*}}>
-# CHECK-NEXT:            {{.*}} b               0x[[ADDR4:[0-9a-f]+]] <{{.*}}>
-# CHECK-NEXT: [[ADDR3]]: {{.*}} mov  x0, #0x2 // =2
-# CHECK-NEXT:            {{.*}} ret
-# CHECK-NEXT: [[ADDR4]]: {{.*}} mov x0, #0x1 // =1
-# CHECK-NEXT:            {{.*}} ret
+# COMMON: <irreversible>:
+# COMMON-NEXT:            {{.*}} cmp  x0, #0x3f
+# COMMON-NEXT:            {{.*}} cbgt x0, #0x3f, 0x[[ADDR6:[0-9a-f]+]] <{{.*}}>
+# COMMON-NEXT:            {{.*}} b               0x[[ADDR7:[0-9a-f]+]] <{{.*}}>
+# COMMON-NEXT: [[ADDR6]]: {{.*}} mov  x0, #0x2 // =2
+# COMMON-NEXT:            {{.*}} ret
+# COMMON-NEXT: [[ADDR7]]: {{.*}} csel x0, x1, x2, le
+# COMMON-NEXT:            {{.*}} ret
 
 ## Force relocation mode.
 .reloc 0, R_AARCH64_NONE
diff --git a/bolt/unittests/Core/CMakeLists.txt b/bolt/unittests/Core/CMakeLists.txt
index 297dec7449202..b755afc8a66da 100644
--- a/bolt/unittests/Core/CMakeLists.txt
+++ b/bolt/unittests/Core/CMakeLists.txt
@@ -24,6 +24,7 @@ target_link_libraries(CoreTests
   PRIVATE
   LLVMBOLTCore
   LLVMBOLTRewrite
+  LLVMBOLTPasses
   LLVMBOLTProfile
   LLVMBOLTUtils
   )
diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
index e67460fe2a6a6..5ea242d65e533 100644
--- a/bolt/unittests/Core/MCPlusBuilder.cpp
+++ b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -18,6 +18,9 @@
 
 #include "bolt/Core/BinaryBasicBlock.h"
 #include "bolt/Core/BinaryFunction.h"
+#include "bolt/Core/BinaryFunctionCallGraph.h"
+#include "bolt/Passes/BinaryPasses.h"
+#include "bolt/Passes/DataflowInfoManager.h"
 #include "bolt/Rewrite/RewriteInstance.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
@@ -233,8 +236,8 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                            .addExpr(MCSymbolRefExpr::create(
                                TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(NeedsImmInc));
-  BC->MIB->reverseBranchCondition(NeedsImmInc, TargetBB->getLabel(),
-                                  BC->Ctx.get());
+  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, NeedsImmInc,
+                                  TargetBB->getLabel(), BC->Ctx.get());
   ASSERT_EQ(NeedsImmInc.getOpcode(), AArch64::CBLTXri);
   ASSERT_EQ(NeedsImmInc.getOperand(1).getImm(), 1);
 
@@ -247,8 +250,8 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                            .addExpr(MCSymbolRefExpr::create(
                                TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(NeedsImmDec));
-  BC->MIB->reverseBranchCondition(NeedsImmDec, TargetBB->getLabel(),
-                                  BC->Ctx.get());
+  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, NeedsImmDec,
+                                  TargetBB->getLabel(), BC->Ctx.get());
   ASSERT_EQ(NeedsImmDec.getOpcode(), AArch64::CBHIXri);
   ASSERT_EQ(NeedsImmDec.getOperand(1).getImm(), 0);
 
@@ -261,8 +264,8 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                                    .addExpr(MCSymbolRefExpr::create(
                                        TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(CompRegNeedsRegSwap));
-  BC->MIB->reverseBranchCondition(CompRegNeedsRegSwap, TargetBB->getLabel(),
-                                  BC->Ctx.get());
+  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, CompRegNeedsRegSwap,
+                                  TargetBB->getLabel(), BC->Ctx.get());
   ASSERT_EQ(CompRegNeedsRegSwap.getOpcode(), AArch64::CBGTXrr);
   ASSERT_EQ(CompRegNeedsRegSwap.getOperand(0).getReg(), AArch64::X1);
   ASSERT_EQ(CompRegNeedsRegSwap.getOperand(1).getReg(), AArch64::X0);
@@ -276,8 +279,8 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                                     .addExpr(MCSymbolRefExpr::create(
                                         TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(CompByteNeedsRegSwap));
-  BC->MIB->reverseBranchCondition(CompByteNeedsRegSwap, TargetBB->getLabel(),
-                                  BC->Ctx.get());
+  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, CompByteNeedsRegSwap,
+                                  TargetBB->getLabel(), BC->Ctx.get());
   ASSERT_EQ(CompByteNeedsRegSwap.getOpcode(), AArch64::CBBHSWrr);
   ASSERT_EQ(CompByteNeedsRegSwap.getOperand(0).getReg(), AArch64::W1);
   ASSERT_EQ(CompByteNeedsRegSwap.getOperand(1).getReg(), AArch64::W0);
@@ -291,8 +294,8 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                                     .addExpr(MCSymbolRefExpr::create(
                                         TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(CompHalfNeedsRegSwap));
-  BC->MIB->reverseBranchCondition(CompHalfNeedsRegSwap, TargetBB->getLabel(),
-                                  BC->Ctx.get());
+  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, CompHalfNeedsRegSwap,
+                                  TargetBB->getLabel(), BC->Ctx.get());
   ASSERT_EQ(CompHalfNeedsRegSwap.getOpcode(), AArch64::CBHHIWrr);
   ASSERT_EQ(CompHalfNeedsRegSwap.getOperand(0).getReg(), AArch64::W1);
   ASSERT_EQ(CompHalfNeedsRegSwap.getOperand(1).getReg(), AArch64::W0);
@@ -318,6 +321,123 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
   ASSERT_FALSE(BC->MIB->isReversibleBranch(Overflows));
 }
 
+TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Underflows) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  BinaryBasicBlock *EntryBB = BF->addBasicBlock();
+  BinaryBasicBlock *FallThroughBB = BF->addBasicBlock();
+  BinaryBasicBlock *TargetBB = BF->addBasicBlock();
+  BF->addEntryPoint(*EntryBB);
+  EntryBB->addSuccessor(TargetBB);
+  EntryBB->addSuccessor(FallThroughBB);
+
+  // Inversion requires expansion, immediate value underflows.
+  // cblt x0, #0, target ~> cmp x0, #0
+  //                        b.ge target
+  auto I =
+      EntryBB->addInstruction(MCInstBuilder(AArch64::CBLTXri)
+                                  .addReg(AArch64::X0)
+                                  .addImm(0)
+                                  .addExpr(MCSymbolRefExpr::create(
+                                      TargetBB->getLabel(), *BC->Ctx.get())));
+  BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
+  RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
+  DataflowInfoManager DIM(*BF, &RA, nullptr);
+
+  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &DIM));
+  BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
+                                  BC->Ctx.get(), &DIM);
+  I = EntryBB->begin();
+  ASSERT_EQ(I->getOpcode(), AArch64::SUBSXri);
+  ASSERT_EQ(I->getOperand(0).getReg(), AArch64::XZR);
+  ASSERT_EQ(I->getOperand(1).getReg(), AArch64::X0);
+  ASSERT_EQ(I->getOperand(2).getImm(), 0);
+  ASSERT_EQ(I->getOperand(3).getImm(), 0);
+  I++;
+  ASSERT_EQ(I->getOpcode(), AArch64::Bcc);
+  ASSERT_EQ(I->getOperand(0).getImm(), AArch64CC::GE);
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Overflows) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  BinaryBasicBlock *EntryBB = BF->addBasicBlock();
+  BinaryBasicBlock *FallThroughBB = BF->addBasicBlock();
+  BinaryBasicBlock *TargetBB = BF->addBasicBlock();
+  BF->addEntryPoint(*EntryBB);
+  EntryBB->addSuccessor(TargetBB);
+  EntryBB->addSuccessor(FallThroughBB);
+
+  // Inversion requires expansion, immediate value overflows.
+  // cbhi w0, #63, target ~> cmp w0, #63
+  //                         b.ls target
+  auto I =
+      EntryBB->addInstruction(MCInstBuilder(AArch64::CBHIWri)
+                                  .addReg(AArch64::W0)
+                                  .addImm(63)
+                                  .addExpr(MCSymbolRefExpr::create(
+                                      TargetBB->getLabel(), *BC->Ctx.get())));
+  BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
+  RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
+  DataflowInfoManager DIM(*BF, &RA, nullptr);
+
+  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &DIM));
+  BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
+                                  BC->Ctx.get(), &DIM);
+  I = EntryBB->begin();
+  ASSERT_EQ(I->getOpcode(), AArch64::SUBSWri);
+  ASSERT_EQ(I->getOperand(0).getReg(), AArch64::WZR);
+  ASSERT_EQ(I->getOperand(1).getReg(), AArch64::W0);
+  ASSERT_EQ(I->getOperand(2).getImm(), 63);
+  ASSERT_EQ(I->getOperand(3).getImm(), 0);
+  I++;
+  ASSERT_EQ(I->getOpcode(), AArch64::Bcc);
+  ASSERT_EQ(I->getOperand(0).getImm(), AArch64CC::LS);
+}
+
+TEST_P(MCPlusBuilderTester, AArch64_IsReversibleBranch_LiveCondFlags) {
+  if (GetParam() != Triple::aarch64)
+    GTEST_SKIP();
+
+  BinaryFunction *BF = BC->createInjectedBinaryFunction("BF", true);
+  BinaryBasicBlock *EntryBB = BF->addBasicBlock();
+  BinaryBasicBlock *FallThroughBB = BF->addBasicBlock();
+  BinaryBasicBlock *TargetBB = BF->addBasicBlock();
+  BF->addEntryPoint(*EntryBB);
+  EntryBB->addSuccessor(TargetBB);
+  EntryBB->addSuccessor(FallThroughBB);
+
+  // cmp x0, #63
+  EntryBB->addInstruction(MCInstBuilder(AArch64::SUBSXri)
+                              .addReg(AArch64::XZR)
+                              .addReg(AArch64::X0)
+                              .addImm(63)
+                              .addImm(0));
+  // cbgt x0, #63, target
+  auto I =
+      EntryBB->addInstruction(MCInstBuilder(AArch64::CBGTXri)
+                                  .addReg(AArch64::X0)
+                                  .addImm(63)
+                                  .addExpr(MCSymbolRefExpr::create(
+                                      TargetBB->getLabel(), *BC->Ctx.get())));
+  // csel x0, x1, x2, le
+  FallThroughBB->addInstruction(MCInstBuilder(AArch64::CSELXr)
+                                    .addReg(AArch64::X0)
+                                    .addReg(AArch64::X1)
+                                    .addReg(AArch64::X2)
+                                    .addImm(13));
+
+  BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
+  RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
+  DataflowInfoManager DIM(*BF, &RA, nullptr);
+
+  ASSERT_FALSE(BC->MIB->isReversibleBranch(*I, &DIM));
+}
+
 TEST_P(MCPlusBuilderTester, AArch64_CmpJE) {
   if (GetParam() != Triple::aarch64)
     GTEST_SKIP();

>From 75d8603c42ec4339a70bb8d0ec5205ab61f3efd4 Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Thu, 2 Jul 2026 12:33:22 +0100
Subject: [PATCH 2/9] Run Liveness Analysis upfront and cache it before any CFG
 changes.

Teach the branch-fixing paths to use cached branch liveness information
when deciding whether this expansion is legal. The liveness snapshot is
built before branch relaxation/fixup and is safe for the current users:
they only insert trampolines/stubs between existing CFG edges or invert
branches without changing program semantics.

Since expansion can grow the source basic block, update the local
address accounting used by branch relaxation.
---
 bolt/include/bolt/Core/BinaryFunction.h       |   4 +-
 bolt/include/bolt/Core/MCPlusBuilder.h        |  23 +++-
 bolt/include/bolt/Passes/LongJmp.h            |   4 +-
 bolt/lib/Core/BinaryFunction.cpp              |   6 +-
 bolt/lib/Passes/BinaryPasses.cpp              |  32 ++---
 bolt/lib/Passes/LongJmp.cpp                   | 109 +++++++++++-------
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   |  34 ++++--
 bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp  |   7 +-
 bolt/lib/Target/X86/X86MCPlusBuilder.cpp      |   7 +-
 bolt/unittests/Core/MCPlusBuilder.cpp         |  16 ++-
 10 files changed, 148 insertions(+), 94 deletions(-)

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 51651e41497da..40e7da06666fc 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -65,7 +65,7 @@ class DWARFUnit;
 
 namespace bolt {
 
-class DataflowInfoManager;
+struct BranchLivenessInfo;
 
 using InputOffsetToAddressMapTy = std::unordered_multimap<uint64_t, uint64_t>;
 
@@ -2465,7 +2465,7 @@ class BinaryFunction {
   /// while the second successor - false/fall-through branch.
   ///
   /// When we reverse the branch condition, the CFG is updated accordingly.
-  void fixBranches(DataflowInfoManager *DIM = nullptr);
+  void fixBranches(const BranchLivenessInfo *BLI = nullptr);
 
   /// Mark function as finalized. No further optimizations are permitted.
   void setFinalized() { CurrentState = State::CFG_Finalized; }
diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index c98b56d14523c..c2859561910d5 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -18,6 +18,7 @@
 #include "bolt/Core/Relocation.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
 #include "llvm/MC/MCAsmBackend.h"
@@ -74,6 +75,15 @@ enum BTIKind {
   JC /// Accepting both.
 };
 
+struct BranchLivenessInfo {
+  DenseMap<const MCInst *, bool> FlagsLiveIn;
+
+  bool mustPreserveFlags(const MCInst &Inst) const {
+    auto It = FlagsLiveIn.find(&Inst);
+    return It == FlagsLiveIn.end() || It->second;
+  }
+};
+
 class MCPlusBuilder {
 public:
   using AllocatorIdTy = uint16_t;
@@ -475,9 +485,16 @@ class MCPlusBuilder {
     return false;
   }
 
+  /// Return liveness info required for branch transformations.
+  virtual BranchLivenessInfo
+  createBranchLivenessInfo(BinaryFunction &BF, DataflowInfoManager &DIM) const {
+    return BranchLivenessInfo();
+  }
+
   /// Check whether this conditional branch can be reversed
-  virtual bool isReversibleBranch(const MCInst &Inst,
-                                  DataflowInfoManager *DIM = nullptr) const {
+  virtual bool
+  isReversibleBranch(const MCInst &Inst,
+                     const BranchLivenessInfo *BLI = nullptr) const {
     assert(!isUnsupportedInstruction(Inst) && isConditionalBranch(Inst) &&
            "Instruction is not known conditional branch");
 
@@ -2147,7 +2164,7 @@ class MCPlusBuilder {
   virtual void
   reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
                          const MCSymbol *TBB, MCContext *Ctx,
-                         DataflowInfoManager *DIM = nullptr) const {
+                         const BranchLivenessInfo *BLI = nullptr) const {
     llvm_unreachable("not implemented");
   }
 
diff --git a/bolt/include/bolt/Passes/LongJmp.h b/bolt/include/bolt/Passes/LongJmp.h
index ccdebf9e1ed0a..90e10943c0f70 100644
--- a/bolt/include/bolt/Passes/LongJmp.h
+++ b/bolt/include/bolt/Passes/LongJmp.h
@@ -14,7 +14,7 @@
 namespace llvm {
 namespace bolt {
 
-class DataflowInfoManager;
+struct BranchLivenessInfo;
 
 /// LongJmp is veneer-insertion pass originally written for AArch64 that
 /// compensates for its short-range branches, typically done during linking. We
@@ -76,7 +76,7 @@ class LongJmpPass : public BinaryFunctionPass {
   /// Relax all internal function branches including those between fragments.
   /// Assume that fragments are placed in different sections but are within
   /// 128MB of each other.
-  void relaxLocalBranches(BinaryFunction &BF, DataflowInfoManager *DIM);
+  void relaxLocalBranches(BinaryFunction &BF, const BranchLivenessInfo *BLI);
 
   ///                 -- Layout estimation methods --
   /// Try to do layout before running the emitter, by looking at BinaryFunctions
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index ddfd163be84ec..4fa0290af8b3b 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3650,7 +3650,7 @@ bool BinaryFunction::validateCFG() const {
   return true;
 }
 
-void BinaryFunction::fixBranches(DataflowInfoManager *DIM) {
+void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
   assert(isSimple() && "Expected function with valid CFG.");
 
   auto &MIB = BC.MIB;
@@ -3709,7 +3709,7 @@ void BinaryFunction::fixBranches(DataflowInfoManager *DIM) {
 
       // Reverse branch condition and swap successors.
       auto swapSuccessors = [&]() {
-        if (!MIB->isReversibleBranch(*CondBranch, DIM)) {
+        if (!MIB->isReversibleBranch(*CondBranch, BLI)) {
           if (opts::Verbosity) {
             BC.outs() << "BOLT-INFO: unable to swap successors in " << *this
                       << '\n';
@@ -3720,7 +3720,7 @@ void BinaryFunction::fixBranches(DataflowInfoManager *DIM) {
         BB->swapConditionalSuccessors();
         auto L = BC.scopeLock();
         MIB->reverseBranchCondition(BB, *CondBranch, TSuccessor->getLabel(),
-                                    Ctx, DIM);
+                                    Ctx, BLI);
         return true;
       };
 
diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index 55b88233dcbde..95ebbd8258159 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -547,23 +547,25 @@ bool ReorderBasicBlocks::modifyFunctionLayout(BinaryFunction &BF,
 }
 
 Error FixupBranches::runOnFunctions(BinaryContext &BC) {
-  std::unique_ptr<BinaryFunctionCallGraph> CG;
-  std::unique_ptr<RegAnalysis> RA;
-  std::unique_ptr<DataflowInfoManager> DIM;
+  auto forEachFunction = [&](auto &&Apply) {
+    for (auto &It : BC.getBinaryFunctions()) {
+      BinaryFunction &Function = It.second;
+      if (!BC.shouldEmit(Function) || !Function.isSimple())
+        continue;
+      Apply(Function);
+    }
+  };
 
   if (opts::LivenessAnalysis) {
-    CG = std::make_unique<BinaryFunctionCallGraph>(buildCallGraph(BC));
-    RA = std::make_unique<RegAnalysis>(BC, &BC.getBinaryFunctions(), CG.get());
-  }
-  for (auto &It : BC.getBinaryFunctions()) {
-    BinaryFunction &Function = It.second;
-    if (!BC.shouldEmit(Function) || !Function.isSimple())
-      continue;
-
-    if (opts::LivenessAnalysis)
-      DIM = std::make_unique<DataflowInfoManager>(Function, RA.get(), nullptr);
-    Function.fixBranches(DIM.get());
-  }
+    BinaryFunctionCallGraph CG = buildCallGraph(BC);
+    RegAnalysis RA(BC, &BC.getBinaryFunctions(), &CG);
+    forEachFunction([&](BinaryFunction &BF) {
+      DataflowInfoManager DIM(BF, &RA, nullptr);
+      BranchLivenessInfo Info = BC.MIB->createBranchLivenessInfo(BF, DIM);
+      BF.fixBranches(&Info);
+    });
+  } else
+    forEachFunction([&](BinaryFunction &BF) { BF.fixBranches(nullptr); });
   return Error::success();
 }
 
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 03d5aa4bf3e2c..871791e6b7988 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -665,7 +665,7 @@ Error LongJmpPass::relax(BinaryFunction &Func, bool &Modified) {
 }
 
 void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
-                                     DataflowInfoManager *DIM) {
+                                     const BranchLivenessInfo *BLI) {
   BinaryContext &BC = BF.getBinaryContext();
   auto &MIB = BC.MIB;
 
@@ -711,14 +711,18 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
     DenseMap<const BinaryBasicBlock *, BinaryBasicBlock *> FragmentTrampolines;
 
     // Create a trampoline code after \p BB or at the end of the fragment if BB
-    // is nullptr. If \p UpdateOffsets is true, update FragmentSize and offsets
-    // for basic blocks affected by the insertion of the trampoline.
+    // is nullptr. \p Offset is the fragment size delta caused by the insertion,
+    // including any growth of \p BB before the trampoline.
     auto addTrampolineAfter = [&](BinaryBasicBlock *BB,
                                   BinaryBasicBlock *TargetBB, uint64_t Count,
-                                  bool UpdateOffsets = true) {
+                                  uint64_t Offset) {
       FunctionTrampolines.emplace_back(BB ? BB : FF.back(),
                                        BF.createBasicBlock());
       BinaryBasicBlock *TrampolineBB = FunctionTrampolines.back().second.get();
+      const uint64_t BBGrowth = Offset ? Offset - TrampolineSize : 0;
+
+      if (BB)
+        BB->setOutputEndAddress(BB->getOutputEndAddress() + BBGrowth);
 
       MCInst Inst;
       {
@@ -734,13 +738,23 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       TrampolineBB->setOutputEndAddress(TrampolineAddress + TrampolineSize);
       TrampolineBB->setFragmentNum(FF.getFragmentNum());
 
+      // Shift the fragment-local output address range for blocks at or after
+      // the old end address.
+      auto adjustBasicBlockAddress = [](BinaryBasicBlock *BB, uint64_t Address,
+                                        uint64_t Offset) {
+        if (BB->getOutputStartAddress() < Address)
+          return;
+        BB->setOutputStartAddress(BB->getOutputStartAddress() + Offset);
+        BB->setOutputEndAddress(BB->getOutputEndAddress() + Offset);
+      };
+
       if (!FragmentTrampolines.lookup(TargetBB))
         FragmentTrampolines[TargetBB] = TrampolineBB;
 
-      if (!UpdateOffsets)
+      if (!Offset)
         return TrampolineBB;
 
-      FragmentSize += TrampolineSize;
+      FragmentSize += Offset;
 
       // If the trampoline was added at the end of the fragment, offsets of
       // other fragments should stay intact.
@@ -748,13 +762,8 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
         return TrampolineBB;
 
       // Update offsets for blocks after BB.
-      for (BinaryBasicBlock *IBB : FF) {
-        if (IBB->getOutputStartAddress() >= TrampolineAddress) {
-          IBB->setOutputStartAddress(IBB->getOutputStartAddress() +
-                                     TrampolineSize);
-          IBB->setOutputEndAddress(IBB->getOutputEndAddress() + TrampolineSize);
-        }
-      }
+      for (BinaryBasicBlock *IBB : FF)
+        adjustBasicBlockAddress(IBB, TrampolineAddress - BBGrowth, Offset);
 
       // Update offsets for trampolines in this fragment that are placed after
       // the new trampoline. Note that trampoline blocks are not part of the
@@ -766,11 +775,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
           continue;
         if (IBB == TrampolineBB)
           continue;
-        if (IBB->getOutputStartAddress() >= TrampolineAddress) {
-          IBB->setOutputStartAddress(IBB->getOutputStartAddress() +
-                                     TrampolineSize);
-          IBB->setOutputEndAddress(IBB->getOutputEndAddress() + TrampolineSize);
-        }
+        adjustBasicBlockAddress(IBB, TrampolineAddress - BBGrowth, Offset);
       }
 
       return TrampolineBB;
@@ -791,7 +796,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       BinaryBasicBlock *TargetBB = BB->getSuccessor(TargetSymbol, BI);
 
       BinaryBasicBlock *TrampolineBB =
-          addTrampolineAfter(BB, TargetBB, BI.Count, /*UpdateOffsets*/ false);
+          addTrampolineAfter(BB, TargetBB, BI.Count, /*Offset=*/0);
       BB->replaceSuccessor(TargetBB, TrampolineBB, BI.Count);
     }
 
@@ -824,7 +829,8 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       // case we will need further relaxation.
       const int64_t OffsetToEnd = FragmentSize - InstAddress;
       if (Count == 0 && isBranchOffsetInRange(Inst, OffsetToEnd)) {
-        TrampolineBB = addTrampolineAfter(nullptr, TargetBB, Count);
+        TrampolineBB =
+            addTrampolineAfter(nullptr, TargetBB, Count, TrampolineSize);
         BB->replaceSuccessor(TargetBB, TrampolineBB, Count);
         auto L = BC.scopeLock();
         MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(), BC.Ctx.get());
@@ -835,7 +841,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       // If the other successor is a fall-through, invert the condition code.
       BinaryBasicBlock *NextBB =
           BF->getLayout().getBasicBlockAfter(BB, /*IgnoreSplits*/ false);
-      bool IsReversibleBranch = MIB->isReversibleBranch(Inst, DIM);
+      bool IsReversibleBranch = MIB->isReversibleBranch(Inst, BLI);
       bool ShouldReverseBranch = BB->getConditionalSuccessor(false) == NextBB;
 
       // Create a trampoline basic block for the fall-through target of the
@@ -843,19 +849,26 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       if (ShouldReverseBranch && !IsReversibleBranch) {
         const uint64_t NextCount = BB->getBranchInfo(*NextBB).Count;
         BinaryBasicBlock *FallThrough =
-            addTrampolineAfter(BB, NextBB, NextCount);
+            addTrampolineAfter(BB, NextBB, NextCount, TrampolineSize);
         BB->replaceSuccessor(NextBB, FallThrough, NextCount);
       }
 
-      // Create a trampoline basic block for the taken target of the branch.
-      TrampolineBB = addTrampolineAfter(BB, TargetBB, Count);
-
       if (ShouldReverseBranch && IsReversibleBranch) {
+        const uint64_t OldBBSize = BB->estimateSize();
         BB->swapConditionalSuccessors();
-        auto L = BC.scopeLock();
-        MIB->reverseBranchCondition(BB, Inst, NextBB->getLabel(), BC.Ctx.get(),
-                                    DIM);
+        {
+          auto L = BC.scopeLock();
+          MIB->reverseBranchCondition(BB, Inst, NextBB->getLabel(),
+                                      BC.Ctx.get(), BLI);
+        }
+        const uint64_t NewBBSize = BB->estimateSize();
+
+        // Create a trampoline basic block for the original taken target.
+        TrampolineBB = addTrampolineAfter(
+            BB, TargetBB, Count, TrampolineSize + (NewBBSize - OldBBSize));
       } else {
+        // Create a trampoline basic block for the taken target of the branch.
+        TrampolineBB = addTrampolineAfter(BB, TargetBB, Count, TrampolineSize);
         auto L = BC.scopeLock();
         MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(), BC.Ctx.get());
       }
@@ -870,7 +883,10 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       for (auto BBI = FF.begin(); BBI != FF.end(); ++BBI) {
         BinaryBasicBlock *BB = *BBI;
         uint64_t NextInstOffset = BB->getOutputStartAddress();
-        for (MCInst &Inst : *BB) {
+        // Branch reversal may replace the current instruction with a sequence.
+        // Use an index so the next instruction is reloaded after the mutation.
+        for (size_t I = 0; I < BB->size(); ++I) {
+          MCInst &Inst = *(BB->begin() + I);
           const size_t InstAddress = NextInstOffset;
           if (!MIB->isPseudo(Inst))
             NextInstOffset += 4;
@@ -939,30 +955,37 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
           opts::SplitStrategy != opts::SplitFunctionsStrategy::CDSplit) &&
          "LongJmp cannot work with functions split in more than two fragments");
 
-  std::unique_ptr<BinaryFunctionCallGraph> CG;
-  std::unique_ptr<RegAnalysis> RA;
-  std::unique_ptr<DataflowInfoManager> DIM;
-
+  DenseMap<BinaryFunction *, BranchLivenessInfo> BranchLiveness;
   if (opts::LivenessAnalysis) {
-    CG = std::make_unique<BinaryFunctionCallGraph>(buildCallGraph(BC));
-    RA = std::make_unique<RegAnalysis>(BC, &BC.getBinaryFunctions(), CG.get());
+    BinaryFunctionCallGraph CG = buildCallGraph(BC);
+    RegAnalysis RA(BC, &BC.getBinaryFunctions(), &CG);
+    for (auto &It : BC.getBinaryFunctions()) {
+      BinaryFunction &BF = It.second;
+      if (!BC.shouldEmit(BF) || !BF.isSimple())
+        continue;
+      DataflowInfoManager DIM(BF, &RA, nullptr);
+      BranchLiveness[&BF] = BC.MIB->createBranchLivenessInfo(BF, DIM);
+    }
   }
+  auto getBranchLiveness =
+      [&](BinaryFunction &BF) -> const BranchLivenessInfo * {
+    auto It = BranchLiveness.find(&BF);
+    return It == BranchLiveness.end() ? nullptr : &It->second;
+  };
 
   if (opts::CompactCodeModel) {
     BC.outs()
         << "BOLT-INFO: relaxing branches for compact code model (<128MB)\n";
 
-    ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
-      if (opts::LivenessAnalysis)
-        DIM = std::make_unique<DataflowInfoManager>(BF, RA.get(), nullptr);
-      relaxLocalBranches(BF, DIM.get());
-    };
-
     ParallelUtilities::PredicateTy SkipPredicate =
         [&](const BinaryFunction &BF) {
           return !BC.shouldEmit(BF) || !BF.isSimple();
         };
 
+    ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
+      relaxLocalBranches(BF, getBranchLiveness(BF));
+    };
+
     ParallelUtilities::runOnEachFunction(
         BC, ParallelUtilities::SchedulingPolicy::SP_INST_LINEAR, WorkFun,
         SkipPredicate, "RelaxLocalBranches");
@@ -980,14 +1003,12 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
     tentativeLayout(BC, Sorted);
     updateStubGroups();
     for (BinaryFunction *Func : Sorted) {
-      if (opts::LivenessAnalysis)
-        DIM = std::make_unique<DataflowInfoManager>(*Func, RA.get(), nullptr);
       if (auto E = relax(*Func, Modified))
         return Error(std::move(E));
       // Don't ruin non-simple functions, they can't afford to have the layout
       // changed.
       if (Modified && Func->isSimple())
-        Func->fixBranches(DIM.get());
+        Func->fixBranches(getBranchLiveness(*Func));
     }
   } while (Modified);
   BC.outs() << "BOLT-INFO: Inserted " << NumHotStubs
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index ecb7bac9d7796..8025f75b789e1 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -2191,28 +2191,38 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     }
   }
 
-  bool isReversibleBranch(const MCInst &Inst,
-                          DataflowInfoManager *DIM = nullptr) const override {
+  BranchLivenessInfo
+  createBranchLivenessInfo(BinaryFunction &BF,
+                           DataflowInfoManager &DIM) const override {
+    BranchLivenessInfo Info;
+    LivenessAnalysis &LA = DIM.getLivenessAnalysis();
+    for (BinaryBasicBlock &BB : BF)
+      for (MCInst &Inst : BB)
+        if (isCompAndBranch(Inst))
+          Info.FlagsLiveIn[&Inst] = LA.getLiveIn(Inst).test(getFlagsReg());
+    return Info;
+  }
+
+  bool
+  isReversibleBranch(const MCInst &Inst,
+                     const BranchLivenessInfo *BLI = nullptr) const override {
     if (isCompAndBranch(Inst)) {
-      bool MayClobberFlags =
-          DIM ? DIM->getLivenessAnalysis().getLiveIn(Inst).test(getFlagsReg())
-              : true;
+      bool MustPreserveFlags = BLI ? BLI->mustPreserveFlags(Inst) : true;
       unsigned InvertedOpcode = getInvertedBranchOpcode(Inst.getOpcode());
       if (needsImmDec(InvertedOpcode) && Inst.getOperand(1).getImm() == 0 &&
-          MayClobberFlags)
+          MustPreserveFlags)
         return false;
       if (needsImmInc(InvertedOpcode) && Inst.getOperand(1).getImm() == 63 &&
-          MayClobberFlags)
+          MustPreserveFlags)
         return false;
     }
     return MCPlusBuilder::isReversibleBranch(Inst);
   }
 
-  void
-  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
-                         const MCSymbol *TBB, MCContext *Ctx,
-                         DataflowInfoManager *DIM = nullptr) const override {
-    assert(isReversibleBranch(Inst, DIM) && "Irreversible branch");
+  void reverseBranchCondition(
+      BinaryBasicBlock *Parent, MCInst &Inst, const MCSymbol *TBB,
+      MCContext *Ctx, const BranchLivenessInfo *BLI = nullptr) const override {
+    assert(isReversibleBranch(Inst, BLI) && "Irreversible branch");
 
     if (isTB(Inst) || isCB(Inst) || isCompAndBranch(Inst)) {
       bool ImmediateOutOfBounds = false;
diff --git a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
index b954a392c4f26..34e8392f59743 100644
--- a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
+++ b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
@@ -162,10 +162,9 @@ class RISCVMCPlusBuilder : public MCPlusBuilder {
     }
   }
 
-  void
-  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
-                         const MCSymbol *TBB, MCContext *Ctx,
-                         DataflowInfoManager *DIM = nullptr) const override {
+  void reverseBranchCondition(
+      BinaryBasicBlock *Parent, MCInst &Inst, const MCSymbol *TBB,
+      MCContext *Ctx, const BranchLivenessInfo *BLI = nullptr) const override {
     auto Opcode = getInvertedBranchOpcode(Inst.getOpcode());
     Inst.setOpcode(Opcode);
     replaceBranchTarget(Inst, TBB, Ctx);
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index 88dcf994e5dc7..98b9eb8fdd4cb 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -2811,10 +2811,9 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     Inst.addOperand(MCOperand::createImm(CC));
   }
 
-  void
-  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
-                         const MCSymbol *TBB, MCContext *Ctx,
-                         DataflowInfoManager *DIM = nullptr) const override {
+  void reverseBranchCondition(
+      BinaryBasicBlock *Parent, MCInst &Inst, const MCSymbol *TBB,
+      MCContext *Ctx, const BranchLivenessInfo *BLI = nullptr) const override {
     unsigned InvCC = getInvertedCondCode(getCondCode(Inst));
     assert(InvCC != X86::COND_INVALID && "invalid branch instruction");
     Inst.getOperand(Info->get(Inst.getOpcode()).NumOperands - 1).setImm(InvCC);
diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
index 5ea242d65e533..5fdaf47094770 100644
--- a/bolt/unittests/Core/MCPlusBuilder.cpp
+++ b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -345,10 +345,12 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Underflows) {
   BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
   RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
   DataflowInfoManager DIM(*BF, &RA, nullptr);
+  BranchLivenessInfo BranchLiveness =
+      BC->MIB->createBranchLivenessInfo(*BF, DIM);
 
-  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &DIM));
+  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &BranchLiveness));
   BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
-                                  BC->Ctx.get(), &DIM);
+                                  BC->Ctx.get(), &BranchLiveness);
   I = EntryBB->begin();
   ASSERT_EQ(I->getOpcode(), AArch64::SUBSXri);
   ASSERT_EQ(I->getOperand(0).getReg(), AArch64::XZR);
@@ -384,10 +386,12 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Overflows) {
   BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
   RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
   DataflowInfoManager DIM(*BF, &RA, nullptr);
+  BranchLivenessInfo BranchLiveness =
+      BC->MIB->createBranchLivenessInfo(*BF, DIM);
 
-  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &DIM));
+  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &BranchLiveness));
   BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
-                                  BC->Ctx.get(), &DIM);
+                                  BC->Ctx.get(), &BranchLiveness);
   I = EntryBB->begin();
   ASSERT_EQ(I->getOpcode(), AArch64::SUBSWri);
   ASSERT_EQ(I->getOperand(0).getReg(), AArch64::WZR);
@@ -434,8 +438,10 @@ TEST_P(MCPlusBuilderTester, AArch64_IsReversibleBranch_LiveCondFlags) {
   BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
   RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
   DataflowInfoManager DIM(*BF, &RA, nullptr);
+  BranchLivenessInfo BranchLiveness =
+      BC->MIB->createBranchLivenessInfo(*BF, DIM);
 
-  ASSERT_FALSE(BC->MIB->isReversibleBranch(*I, &DIM));
+  ASSERT_FALSE(BC->MIB->isReversibleBranch(*I, &BranchLiveness));
 }
 
 TEST_P(MCPlusBuilderTester, AArch64_CmpJE) {

>From e1791b36729b17ba4c039621446d65c0d8bb613e Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Mon, 20 Jul 2026 19:40:36 +0100
Subject: [PATCH 3/9] * update offset calculation in addTrampolineAfter * use
 set instead of map for branch liveness * defer liveness analysis until we
 encounter compare-and-branch inst * add minor comment

---
 bolt/include/bolt/Core/MCPlusBuilder.h        |  8 ++---
 bolt/lib/Passes/LongJmp.cpp                   | 33 ++++++++++---------
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   | 13 ++++++--
 3 files changed, 31 insertions(+), 23 deletions(-)

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index c2859561910d5..d0e76032c7222 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -76,11 +76,10 @@ enum BTIKind {
 };
 
 struct BranchLivenessInfo {
-  DenseMap<const MCInst *, bool> FlagsLiveIn;
+  DenseSet<const MCInst *> FlagsDead;
 
   bool mustPreserveFlags(const MCInst &Inst) const {
-    auto It = FlagsLiveIn.find(&Inst);
-    return It == FlagsLiveIn.end() || It->second;
+    return !FlagsDead.count(&Inst);
   }
 };
 
@@ -2160,7 +2159,8 @@ class MCPlusBuilder {
   }
 
   /// Reverses the branch condition in Inst and update its taken target to TBB.
-  /// Assumes that the branch is reversible.
+  /// Assumes that the branch is reversible. It may replace Inst with a longer
+  /// instruction sequence on some targets.
   virtual void
   reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
                          const MCSymbol *TBB, MCContext *Ctx,
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 871791e6b7988..0e8903697612d 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -711,18 +711,20 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
     DenseMap<const BinaryBasicBlock *, BinaryBasicBlock *> FragmentTrampolines;
 
     // Create a trampoline code after \p BB or at the end of the fragment if BB
-    // is nullptr. \p Offset is the fragment size delta caused by the insertion,
-    // including any growth of \p BB before the trampoline.
+    // is nullptr. \p Offset reflects the size delta of BB caused by splitting
+    // unconditional branches, or replacing a branch with a longer instruction
+    // sequence. It is used to update the output addresses of basic blocks
+    // following the trampoline.
     auto addTrampolineAfter = [&](BinaryBasicBlock *BB,
                                   BinaryBasicBlock *TargetBB, uint64_t Count,
-                                  uint64_t Offset) {
+                                  uint64_t Offset = 0) {
       FunctionTrampolines.emplace_back(BB ? BB : FF.back(),
                                        BF.createBasicBlock());
       BinaryBasicBlock *TrampolineBB = FunctionTrampolines.back().second.get();
-      const uint64_t BBGrowth = Offset ? Offset - TrampolineSize : 0;
-
-      if (BB)
-        BB->setOutputEndAddress(BB->getOutputEndAddress() + BBGrowth);
+      const uint64_t OldBBEnd = BB ? BB->getOutputEndAddress() : 0;
+      if (BB && Offset)
+        BB->setOutputEndAddress(OldBBEnd + Offset);
+      Offset += TrampolineSize;
 
       MCInst Inst;
       {
@@ -763,7 +765,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
 
       // Update offsets for blocks after BB.
       for (BinaryBasicBlock *IBB : FF)
-        adjustBasicBlockAddress(IBB, TrampolineAddress - BBGrowth, Offset);
+        adjustBasicBlockAddress(IBB, OldBBEnd, Offset);
 
       // Update offsets for trampolines in this fragment that are placed after
       // the new trampoline. Note that trampoline blocks are not part of the
@@ -775,7 +777,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
           continue;
         if (IBB == TrampolineBB)
           continue;
-        adjustBasicBlockAddress(IBB, TrampolineAddress - BBGrowth, Offset);
+        adjustBasicBlockAddress(IBB, OldBBEnd, Offset);
       }
 
       return TrampolineBB;
@@ -796,7 +798,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       BinaryBasicBlock *TargetBB = BB->getSuccessor(TargetSymbol, BI);
 
       BinaryBasicBlock *TrampolineBB =
-          addTrampolineAfter(BB, TargetBB, BI.Count, /*Offset=*/0);
+          addTrampolineAfter(BB, TargetBB, BI.Count, /*Offset=*/-4);
       BB->replaceSuccessor(TargetBB, TrampolineBB, BI.Count);
     }
 
@@ -829,8 +831,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       // case we will need further relaxation.
       const int64_t OffsetToEnd = FragmentSize - InstAddress;
       if (Count == 0 && isBranchOffsetInRange(Inst, OffsetToEnd)) {
-        TrampolineBB =
-            addTrampolineAfter(nullptr, TargetBB, Count, TrampolineSize);
+        TrampolineBB = addTrampolineAfter(nullptr, TargetBB, Count);
         BB->replaceSuccessor(TargetBB, TrampolineBB, Count);
         auto L = BC.scopeLock();
         MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(), BC.Ctx.get());
@@ -849,7 +850,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       if (ShouldReverseBranch && !IsReversibleBranch) {
         const uint64_t NextCount = BB->getBranchInfo(*NextBB).Count;
         BinaryBasicBlock *FallThrough =
-            addTrampolineAfter(BB, NextBB, NextCount, TrampolineSize);
+            addTrampolineAfter(BB, NextBB, NextCount);
         BB->replaceSuccessor(NextBB, FallThrough, NextCount);
       }
 
@@ -864,11 +865,11 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
         const uint64_t NewBBSize = BB->estimateSize();
 
         // Create a trampoline basic block for the original taken target.
-        TrampolineBB = addTrampolineAfter(
-            BB, TargetBB, Count, TrampolineSize + (NewBBSize - OldBBSize));
+        TrampolineBB =
+            addTrampolineAfter(BB, TargetBB, Count, NewBBSize - OldBBSize);
       } else {
         // Create a trampoline basic block for the taken target of the branch.
-        TrampolineBB = addTrampolineAfter(BB, TargetBB, Count, TrampolineSize);
+        TrampolineBB = addTrampolineAfter(BB, TargetBB, Count);
         auto L = BC.scopeLock();
         MIB->replaceBranchTarget(Inst, TrampolineBB->getLabel(), BC.Ctx.get());
       }
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index 8025f75b789e1..3503cd7df4750 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -2194,12 +2194,19 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
   BranchLivenessInfo
   createBranchLivenessInfo(BinaryFunction &BF,
                            DataflowInfoManager &DIM) const override {
-    BranchLivenessInfo Info;
-    LivenessAnalysis &LA = DIM.getLivenessAnalysis();
+    SmallVector<MCInst *> CompAndBranchInsts;
     for (BinaryBasicBlock &BB : BF)
       for (MCInst &Inst : BB)
         if (isCompAndBranch(Inst))
-          Info.FlagsLiveIn[&Inst] = LA.getLiveIn(Inst).test(getFlagsReg());
+          CompAndBranchInsts.push_back(&Inst);
+    if (CompAndBranchInsts.empty())
+      return {};
+
+    BranchLivenessInfo Info;
+    LivenessAnalysis &LA = DIM.getLivenessAnalysis();
+    for (MCInst *Inst : CompAndBranchInsts)
+      if (!LA.getLiveIn(*Inst).test(getFlagsReg()))
+        Info.FlagsDead.insert(Inst);
     return Info;
   }
 

>From 152117b61c0e5c6aea369e6be022a90d024f7603 Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Mon, 3 Aug 2026 17:24:15 +0100
Subject: [PATCH 4/9] changes:

- decoupled dependency between passes and target library
- renamed the cmdline option
- removed extraneous setOutputEndAddress after refactoring addTrampolineAfter
- removed call graph construction and only invoked register analysis when the function has cmpbr instructions
---
 bolt/include/bolt/Core/BinaryFunction.h       |  5 +-
 bolt/include/bolt/Core/MCPlusBuilder.h        | 28 ++-------
 .../include/bolt/Passes/BranchLivenessUtils.h | 31 ++++++++++
 bolt/include/bolt/Passes/LongJmp.h            |  9 ++-
 bolt/include/bolt/Utils/CommandLineOpts.h     |  2 +-
 bolt/lib/Core/BinaryFunction.cpp              |  9 ++-
 bolt/lib/Passes/BinaryPasses.cpp              | 44 +++++++-------
 bolt/lib/Passes/BranchLivenessUtils.cpp       | 57 +++++++++++++++++++
 bolt/lib/Passes/CMakeLists.txt                |  1 +
 bolt/lib/Passes/LongJmp.cpp                   | 45 ++++++++-------
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   | 35 +++---------
 bolt/lib/Target/AArch64/CMakeLists.txt        |  6 +-
 bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp  |  6 +-
 bolt/lib/Target/X86/X86MCPlusBuilder.cpp      |  6 +-
 bolt/lib/Utils/CommandLineOpts.cpp            |  6 +-
 .../AArch64/compare-and-branch-inversion.S    |  2 +-
 bolt/unittests/Core/CMakeLists.txt            |  1 -
 bolt/unittests/Core/MCPlusBuilder.cpp         | 31 ++--------
 18 files changed, 182 insertions(+), 142 deletions(-)
 create mode 100644 bolt/include/bolt/Passes/BranchLivenessUtils.h
 create mode 100644 bolt/lib/Passes/BranchLivenessUtils.cpp

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 8dc98630d8c39..6128af7f5a941 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -36,6 +36,7 @@
 #include "bolt/Core/MCPlus.h"
 #include "bolt/Utils/NameResolver.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallString.h"
@@ -66,8 +67,6 @@ class DWARFUnit;
 
 namespace bolt {
 
-struct BranchLivenessInfo;
-
 using InputOffsetToAddressMapTy = std::unordered_multimap<uint64_t, uint64_t>;
 
 /// Types of macro-fusion alignment corrections.
@@ -2506,7 +2505,7 @@ class BinaryFunction {
   /// while the second successor - false/fall-through branch.
   ///
   /// When we reverse the branch condition, the CFG is updated accordingly.
-  void fixBranches(const BranchLivenessInfo *BLI = nullptr);
+  void fixBranches(const DenseSet<const MCInst *> *DeadFlagBranches = nullptr);
 
   /// Mark function as finalized. No further optimizations are permitted.
   void setFinalized() { CurrentState = State::CFG_Finalized; }
diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 45e8f82220f9e..0879fa5f74b55 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -18,7 +18,6 @@
 #include "bolt/Core/Relocation.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/BitVector.h"
-#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/CodeGen/TargetOpcodes.h"
 #include "llvm/MC/MCAsmBackend.h"
@@ -53,7 +52,6 @@ namespace bolt {
 class BinaryBasicBlock;
 class BinaryContext;
 class BinaryFunction;
-class DataflowInfoManager;
 
 /// Different types of indirect branches encountered during disassembly.
 enum class IndirectBranchType : char {
@@ -75,14 +73,6 @@ enum BTIKind {
   JC /// Accepting both.
 };
 
-struct BranchLivenessInfo {
-  DenseSet<const MCInst *> FlagsDead;
-
-  bool mustPreserveFlags(const MCInst &Inst) const {
-    return !FlagsDead.count(&Inst);
-  }
-};
-
 class MCPlusBuilder {
 public:
   using AllocatorIdTy = uint16_t;
@@ -484,16 +474,9 @@ class MCPlusBuilder {
     return false;
   }
 
-  /// Return liveness info required for branch transformations.
-  virtual BranchLivenessInfo
-  createBranchLivenessInfo(BinaryFunction &BF, DataflowInfoManager &DIM) const {
-    return BranchLivenessInfo();
-  }
-
   /// Check whether this conditional branch can be reversed
-  virtual bool
-  isReversibleBranch(const MCInst &Inst,
-                     const BranchLivenessInfo *BLI = nullptr) const {
+  virtual bool isReversibleBranch(const MCInst &Inst,
+                                  bool MustPreserveFlags = true) const {
     assert(!isUnsupportedInstruction(Inst) && isConditionalBranch(Inst) &&
            "Instruction is not known conditional branch");
 
@@ -2171,10 +2154,9 @@ class MCPlusBuilder {
   /// Reverses the branch condition in Inst and update its taken target to TBB.
   /// Assumes that the branch is reversible. It may replace Inst with a longer
   /// instruction sequence on some targets.
-  virtual void
-  reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
-                         const MCSymbol *TBB, MCContext *Ctx,
-                         const BranchLivenessInfo *BLI = nullptr) const {
+  virtual void reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                                      const MCSymbol *TBB, MCContext *Ctx,
+                                      bool MustPreserveFlags = true) const {
     llvm_unreachable("not implemented");
   }
 
diff --git a/bolt/include/bolt/Passes/BranchLivenessUtils.h b/bolt/include/bolt/Passes/BranchLivenessUtils.h
new file mode 100644
index 0000000000000..aca9df6ef0b5c
--- /dev/null
+++ b/bolt/include/bolt/Passes/BranchLivenessUtils.h
@@ -0,0 +1,31 @@
+//===- bolt/Passes/BranchLivenessUtils.h ------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BOLT_PASSES_BRANCHLIVENESSUTILS_H
+#define BOLT_PASSES_BRANCHLIVENESSUTILS_H
+
+#include "llvm/ADT/DenseSet.h"
+
+namespace llvm {
+class MCInst;
+
+namespace bolt {
+class BinaryFunction;
+class RegAnalysis;
+
+/// Return true if \p BF has short-range branches.
+bool hasShortRangeBranch(BinaryFunction &BF);
+
+/// Return the branch instructions where the target flags register is dead.
+DenseSet<const MCInst *> computeDeadFlagBranches(BinaryFunction &BF,
+                                                 RegAnalysis &RA);
+
+} // namespace bolt
+} // namespace llvm
+
+#endif
diff --git a/bolt/include/bolt/Passes/LongJmp.h b/bolt/include/bolt/Passes/LongJmp.h
index 90e10943c0f70..9aedfcf44b731 100644
--- a/bolt/include/bolt/Passes/LongJmp.h
+++ b/bolt/include/bolt/Passes/LongJmp.h
@@ -10,11 +10,12 @@
 #define BOLT_PASSES_LONGJMP_H
 
 #include "bolt/Passes/BinaryPasses.h"
+#include "llvm/ADT/DenseSet.h"
 
 namespace llvm {
-namespace bolt {
+class MCInst;
 
-struct BranchLivenessInfo;
+namespace bolt {
 
 /// LongJmp is veneer-insertion pass originally written for AArch64 that
 /// compensates for its short-range branches, typically done during linking. We
@@ -76,7 +77,9 @@ class LongJmpPass : public BinaryFunctionPass {
   /// Relax all internal function branches including those between fragments.
   /// Assume that fragments are placed in different sections but are within
   /// 128MB of each other.
-  void relaxLocalBranches(BinaryFunction &BF, const BranchLivenessInfo *BLI);
+  void relaxLocalBranches(
+      BinaryFunction &BF,
+      const DenseSet<const MCInst *> *DeadFlagBranches = nullptr);
 
   ///                 -- Layout estimation methods --
   /// Try to do layout before running the emitter, by looking at BinaryFunctions
diff --git a/bolt/include/bolt/Utils/CommandLineOpts.h b/bolt/include/bolt/Utils/CommandLineOpts.h
index 88b56217f9512..51ca7702749ce 100644
--- a/bolt/include/bolt/Utils/CommandLineOpts.h
+++ b/bolt/include/bolt/Utils/CommandLineOpts.h
@@ -134,7 +134,7 @@ extern llvm::cl::opt<unsigned> Verbosity;
 
 // Option to control whether liveness analysis should be used by
 // FixupBranches and LongJmpPass. Needed for branch inversion on AArch64.
-extern llvm::cl::opt<bool> LivenessAnalysis;
+extern llvm::cl::opt<bool> FixBranchesWithLiveness;
 
 /// Return true if we should process all functions in the binary.
 bool processAllFunctions();
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index f60f0716f4b64..be5f7f9d4cd0e 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3676,7 +3676,8 @@ bool BinaryFunction::validateCFG() const {
   return true;
 }
 
-void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
+void BinaryFunction::fixBranches(
+    const DenseSet<const MCInst *> *DeadFlagBranches) {
   assert(isSimple() && "Expected function with valid CFG.");
 
   auto &MIB = BC.MIB;
@@ -3735,7 +3736,9 @@ void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
 
       // Reverse branch condition and swap successors.
       auto swapSuccessors = [&]() {
-        if (!MIB->isReversibleBranch(*CondBranch, BLI)) {
+        bool MustPreserveFlags =
+            !DeadFlagBranches || !DeadFlagBranches->count(CondBranch);
+        if (!MIB->isReversibleBranch(*CondBranch, MustPreserveFlags)) {
           if (opts::Verbosity) {
             BC.outs() << "BOLT-INFO: unable to swap successors in " << *this
                       << '\n';
@@ -3746,7 +3749,7 @@ void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
         BB->swapConditionalSuccessors();
         auto L = BC.scopeLock();
         MIB->reverseBranchCondition(BB, *CondBranch, TSuccessor->getLabel(),
-                                    Ctx, BLI);
+                                    Ctx, MustPreserveFlags);
         return true;
       };
 
diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index 0a00458777f8f..2197c7342be6c 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -11,18 +11,21 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/BinaryPasses.h"
-#include "bolt/Core/BinaryFunctionCallGraph.h"
 #include "bolt/Core/FunctionLayout.h"
 #include "bolt/Core/ParallelUtilities.h"
-#include "bolt/Passes/DataflowInfoManager.h"
+#include "bolt/Passes/BranchLivenessUtils.h"
+#include "bolt/Passes/RegAnalysis.h"
 #include "bolt/Passes/ReorderAlgorithm.h"
 #include "bolt/Passes/ReorderFunctions.h"
 #include "bolt/Utils/CommandLineOpts.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/CommandLine.h"
 #include <atomic>
 #include <cmath>
 #include <mutex>
 #include <numeric>
+#include <optional>
 #include <vector>
 
 #define DEBUG_TYPE "bolt-opts"
@@ -555,25 +558,26 @@ bool ReorderBasicBlocks::modifyFunctionLayout(BinaryFunction &BF,
 }
 
 Error FixupBranches::runOnFunctions(BinaryContext &BC) {
-  auto forEachFunction = [&](auto &&Apply) {
-    for (auto &It : BC.getBinaryFunctions()) {
-      BinaryFunction &Function = It.second;
-      if (!BC.shouldEmit(Function) || !Function.isSimple())
-        continue;
-      Apply(Function);
-    }
-  };
+  const bool ShouldRunRegisterAnalysis =
+      opts::FixBranchesWithLiveness &&
+      llvm::any_of(BC.getBinaryFunctions(), [&](auto &It) {
+        BinaryFunction &BF = It.second;
+        return BC.shouldEmit(BF) && BF.isSimple() && hasShortRangeBranch(BF);
+      });
 
-  if (opts::LivenessAnalysis) {
-    BinaryFunctionCallGraph CG = buildCallGraph(BC);
-    RegAnalysis RA(BC, &BC.getBinaryFunctions(), &CG);
-    forEachFunction([&](BinaryFunction &BF) {
-      DataflowInfoManager DIM(BF, &RA, nullptr);
-      BranchLivenessInfo Info = BC.MIB->createBranchLivenessInfo(BF, DIM);
-      BF.fixBranches(&Info);
-    });
-  } else
-    forEachFunction([&](BinaryFunction &BF) { BF.fixBranches(nullptr); });
+  std::optional<RegAnalysis> RA;
+  if (ShouldRunRegisterAnalysis)
+    RA.emplace(BC, nullptr, nullptr);
+
+  for (auto &It : BC.getBinaryFunctions()) {
+    BinaryFunction &BF = It.second;
+    if (!BC.shouldEmit(BF) || !BF.isSimple())
+      continue;
+
+    DenseSet<const MCInst *> DeadFlagBranches =
+        RA ? computeDeadFlagBranches(BF, *RA) : DenseSet<const MCInst *>();
+    BF.fixBranches(&DeadFlagBranches);
+  }
   return Error::success();
 }
 
diff --git a/bolt/lib/Passes/BranchLivenessUtils.cpp b/bolt/lib/Passes/BranchLivenessUtils.cpp
new file mode 100644
index 0000000000000..9ab5de3570509
--- /dev/null
+++ b/bolt/lib/Passes/BranchLivenessUtils.cpp
@@ -0,0 +1,57 @@
+//===- bolt/Passes/BranchLivenessUtils.cpp -------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "bolt/Passes/BranchLivenessUtils.h"
+#include "bolt/Core/BinaryContext.h"
+#include "bolt/Core/BinaryFunction.h"
+#include "bolt/Core/MCPlusBuilder.h"
+#include "bolt/Passes/DataflowInfoManager.h"
+#include "bolt/Passes/RegAnalysis.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/MC/MCRegister.h"
+
+namespace llvm {
+namespace bolt {
+
+bool hasShortRangeBranch(BinaryFunction &BF) {
+  BinaryContext &BC = BF.getBinaryContext();
+  if (!BC.isAArch64())
+    return false;
+
+  return llvm::any_of(BF, [&](BinaryBasicBlock &BB) {
+    return llvm::any_of(
+        BB, [&](MCInst &Inst) { return BC.MIB->isShortRangeBranch(Inst); });
+  });
+}
+
+DenseSet<const MCInst *> computeDeadFlagBranches(BinaryFunction &BF,
+                                                 RegAnalysis &RA) {
+  BinaryContext &BC = BF.getBinaryContext();
+  SmallVector<MCInst *, 4> Insts;
+  if (BC.isAArch64())
+    for (BinaryBasicBlock &BB : BF)
+      for (MCInst &Inst : BB)
+        if (BC.MIB->isShortRangeBranch(Inst))
+          Insts.push_back(&Inst);
+
+  DenseSet<const MCInst *> DeadFlagBranches;
+  if (Insts.empty())
+    return DeadFlagBranches;
+
+  DataflowInfoManager DIM(BF, &RA, nullptr);
+  LivenessAnalysis &LA = DIM.getLivenessAnalysis();
+  const MCPhysReg FlagsReg = BC.MIB->getFlagsReg();
+  for (MCInst *Inst : Insts)
+    if (!LA.getLiveIn(*Inst).test(FlagsReg))
+      DeadFlagBranches.insert(Inst);
+  return DeadFlagBranches;
+}
+
+} // namespace bolt
+} // namespace llvm
diff --git a/bolt/lib/Passes/CMakeLists.txt b/bolt/lib/Passes/CMakeLists.txt
index ec012f05cc498..686dee6987a73 100644
--- a/bolt/lib/Passes/CMakeLists.txt
+++ b/bolt/lib/Passes/CMakeLists.txt
@@ -4,6 +4,7 @@ add_llvm_library(LLVMBOLTPasses
   AllocCombiner.cpp
   AsmDump.cpp
   BinaryPasses.cpp
+  BranchLivenessUtils.cpp
   CMOVConversion.cpp
   CacheMetrics.cpp
   DataflowAnalysis.cpp
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 0e8903697612d..497988728a527 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -11,10 +11,12 @@
 //===----------------------------------------------------------------------===//
 
 #include "bolt/Passes/LongJmp.h"
-#include "bolt/Core/BinaryFunctionCallGraph.h"
 #include "bolt/Core/ParallelUtilities.h"
-#include "bolt/Passes/DataflowInfoManager.h"
+#include "bolt/Passes/BranchLivenessUtils.h"
+#include "bolt/Passes/RegAnalysis.h"
 #include "bolt/Utils/CommandLineOpts.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/MathExtras.h"
 
 #define DEBUG_TYPE "longjmp"
@@ -664,8 +666,8 @@ Error LongJmpPass::relax(BinaryFunction &Func, bool &Modified) {
   return Error::success();
 }
 
-void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
-                                     const BranchLivenessInfo *BLI) {
+void LongJmpPass::relaxLocalBranches(
+    BinaryFunction &BF, const DenseSet<const MCInst *> *DeadFlagBranches) {
   BinaryContext &BC = BF.getBinaryContext();
   auto &MIB = BC.MIB;
 
@@ -792,7 +794,6 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
 
       const MCSymbol *TargetSymbol = MIB->getTargetSymbol(*Inst);
       BB->eraseInstruction(BB->findInstruction(Inst));
-      BB->setOutputEndAddress(BB->getOutputEndAddress() - TrampolineSize);
 
       BinaryBasicBlock::BinaryBranchInfo BI;
       BinaryBasicBlock *TargetBB = BB->getSuccessor(TargetSymbol, BI);
@@ -842,7 +843,10 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       // If the other successor is a fall-through, invert the condition code.
       BinaryBasicBlock *NextBB =
           BF->getLayout().getBasicBlockAfter(BB, /*IgnoreSplits*/ false);
-      bool IsReversibleBranch = MIB->isReversibleBranch(Inst, BLI);
+      bool MustPreserveFlags =
+          !DeadFlagBranches || !DeadFlagBranches->count(&Inst);
+      bool IsReversibleBranch =
+          MIB->isReversibleBranch(Inst, MustPreserveFlags);
       bool ShouldReverseBranch = BB->getConditionalSuccessor(false) == NextBB;
 
       // Create a trampoline basic block for the fall-through target of the
@@ -860,7 +864,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
         {
           auto L = BC.scopeLock();
           MIB->reverseBranchCondition(BB, Inst, NextBB->getLabel(),
-                                      BC.Ctx.get(), BLI);
+                                      BC.Ctx.get(), MustPreserveFlags);
         }
         const uint64_t NewBBSize = BB->estimateSize();
 
@@ -956,22 +960,25 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
           opts::SplitStrategy != opts::SplitFunctionsStrategy::CDSplit) &&
          "LongJmp cannot work with functions split in more than two fragments");
 
-  DenseMap<BinaryFunction *, BranchLivenessInfo> BranchLiveness;
-  if (opts::LivenessAnalysis) {
-    BinaryFunctionCallGraph CG = buildCallGraph(BC);
-    RegAnalysis RA(BC, &BC.getBinaryFunctions(), &CG);
+  DenseMap<BinaryFunction *, DenseSet<const MCInst *>> DeadFlagBranches;
+  if (opts::FixBranchesWithLiveness) {
+    SmallVector<BinaryFunction *> Candidates;
     for (auto &It : BC.getBinaryFunctions()) {
       BinaryFunction &BF = It.second;
       if (!BC.shouldEmit(BF) || !BF.isSimple())
         continue;
-      DataflowInfoManager DIM(BF, &RA, nullptr);
-      BranchLiveness[&BF] = BC.MIB->createBranchLivenessInfo(BF, DIM);
+      if (hasShortRangeBranch(BF))
+        Candidates.push_back(&BF);
+    }
+    if (!Candidates.empty()) {
+      RegAnalysis RA(BC, nullptr, nullptr);
+      for (BinaryFunction *BF : Candidates)
+        DeadFlagBranches[BF] = computeDeadFlagBranches(*BF, RA);
     }
   }
-  auto getBranchLiveness =
-      [&](BinaryFunction &BF) -> const BranchLivenessInfo * {
-    auto It = BranchLiveness.find(&BF);
-    return It == BranchLiveness.end() ? nullptr : &It->second;
+  auto getDeadFlagBranches = [&](BinaryFunction &BF) {
+    auto It = DeadFlagBranches.find(&BF);
+    return It == DeadFlagBranches.end() ? nullptr : &It->second;
   };
 
   if (opts::CompactCodeModel) {
@@ -984,7 +991,7 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
         };
 
     ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
-      relaxLocalBranches(BF, getBranchLiveness(BF));
+      relaxLocalBranches(BF, getDeadFlagBranches(BF));
     };
 
     ParallelUtilities::runOnEachFunction(
@@ -1009,7 +1016,7 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
       // Don't ruin non-simple functions, they can't afford to have the layout
       // changed.
       if (Modified && Func->isSimple())
-        Func->fixBranches(getBranchLiveness(*Func));
+        Func->fixBranches(getDeadFlagBranches(*Func));
     }
   } while (Modified);
   BC.outs() << "BOLT-INFO: Inserted " << NumHotStubs
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index 92b95b341af3f..f9bd174b38814 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -22,7 +22,6 @@
 #include "bolt/Core/BinaryFunction.h"
 #include "bolt/Core/MCInstUtils.h"
 #include "bolt/Core/MCPlusBuilder.h"
-#include "bolt/Passes/DataflowInfoManager.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCDisassembler/MCDisassembler.h"
@@ -2191,30 +2190,9 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     }
   }
 
-  BranchLivenessInfo
-  createBranchLivenessInfo(BinaryFunction &BF,
-                           DataflowInfoManager &DIM) const override {
-    SmallVector<MCInst *> CompAndBranchInsts;
-    for (BinaryBasicBlock &BB : BF)
-      for (MCInst &Inst : BB)
-        if (isCompAndBranch(Inst))
-          CompAndBranchInsts.push_back(&Inst);
-    if (CompAndBranchInsts.empty())
-      return {};
-
-    BranchLivenessInfo Info;
-    LivenessAnalysis &LA = DIM.getLivenessAnalysis();
-    for (MCInst *Inst : CompAndBranchInsts)
-      if (!LA.getLiveIn(*Inst).test(getFlagsReg()))
-        Info.FlagsDead.insert(Inst);
-    return Info;
-  }
-
-  bool
-  isReversibleBranch(const MCInst &Inst,
-                     const BranchLivenessInfo *BLI = nullptr) const override {
+  bool isReversibleBranch(const MCInst &Inst,
+                          bool MustPreserveFlags = true) const override {
     if (isCompAndBranch(Inst)) {
-      bool MustPreserveFlags = BLI ? BLI->mustPreserveFlags(Inst) : true;
       unsigned InvertedOpcode = getInvertedBranchOpcode(Inst.getOpcode());
       if (needsImmDec(InvertedOpcode) && Inst.getOperand(1).getImm() == 0 &&
           MustPreserveFlags)
@@ -2226,10 +2204,11 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     return MCPlusBuilder::isReversibleBranch(Inst);
   }
 
-  void reverseBranchCondition(
-      BinaryBasicBlock *Parent, MCInst &Inst, const MCSymbol *TBB,
-      MCContext *Ctx, const BranchLivenessInfo *BLI = nullptr) const override {
-    assert(isReversibleBranch(Inst, BLI) && "Irreversible branch");
+  void reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                              const MCSymbol *TBB, MCContext *Ctx,
+                              bool MustPreserveFlags = true) const override {
+    assert(isReversibleBranch(Inst, MustPreserveFlags) &&
+           "Irreversible branch");
 
     if (isTB(Inst) || isCB(Inst) || isCompAndBranch(Inst)) {
       bool ImmediateOutOfBounds = false;
diff --git a/bolt/lib/Target/AArch64/CMakeLists.txt b/bolt/lib/Target/AArch64/CMakeLists.txt
index e28ed0bd66ba1..1e171748aece6 100644
--- a/bolt/lib/Target/AArch64/CMakeLists.txt
+++ b/bolt/lib/Target/AArch64/CMakeLists.txt
@@ -29,11 +29,7 @@ add_llvm_library(LLVMBOLTTargetAArch64
   AArch64CommonTableGen
   )
 
-target_link_libraries(LLVMBOLTTargetAArch64 PRIVATE
-  LLVMBOLTCore
-  LLVMBOLTPasses
-  LLVMBOLTUtils
-  )
+target_link_libraries(LLVMBOLTTargetAArch64 PRIVATE LLVMBOLTCore LLVMBOLTUtils)
 
 include_directories(
   ${LLVM_MAIN_SRC_DIR}/lib/Target/AArch64
diff --git a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
index 34e8392f59743..082f9242adf22 100644
--- a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
+++ b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
@@ -162,9 +162,9 @@ class RISCVMCPlusBuilder : public MCPlusBuilder {
     }
   }
 
-  void reverseBranchCondition(
-      BinaryBasicBlock *Parent, MCInst &Inst, const MCSymbol *TBB,
-      MCContext *Ctx, const BranchLivenessInfo *BLI = nullptr) const override {
+  void reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                              const MCSymbol *TBB, MCContext *Ctx,
+                              bool MustPreserveFlags = true) const override {
     auto Opcode = getInvertedBranchOpcode(Inst.getOpcode());
     Inst.setOpcode(Opcode);
     replaceBranchTarget(Inst, TBB, Ctx);
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index 8dbe98e004dd2..ba84617dbd65f 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -2811,9 +2811,9 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     Inst.addOperand(MCOperand::createImm(CC));
   }
 
-  void reverseBranchCondition(
-      BinaryBasicBlock *Parent, MCInst &Inst, const MCSymbol *TBB,
-      MCContext *Ctx, const BranchLivenessInfo *BLI = nullptr) const override {
+  void reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
+                              const MCSymbol *TBB, MCContext *Ctx,
+                              bool MustPreserveFlags = true) const override {
     unsigned InvCC = getInvertedCondCode(getCondCode(Inst));
     assert(InvCC != X86::COND_INVALID && "invalid branch instruction");
     Inst.getOperand(Info->get(Inst.getOpcode()).NumOperands - 1).setImm(InvCC);
diff --git a/bolt/lib/Utils/CommandLineOpts.cpp b/bolt/lib/Utils/CommandLineOpts.cpp
index 82fde5788397a..4eaec32c7daf3 100644
--- a/bolt/lib/Utils/CommandLineOpts.cpp
+++ b/bolt/lib/Utils/CommandLineOpts.cpp
@@ -365,9 +365,9 @@ cl::opt<unsigned>
               cl::init(0), cl::ZeroOrMore, cl::cat(BoltCategory),
               cl::sub(cl::SubCommand::getAll()));
 
-cl::opt<bool> LivenessAnalysis(
-    "liveness-analysis",
-    cl::desc("use liveness analysis in FixupBranches and LongJmpPass"
+cl::opt<bool> FixBranchesWithLiveness(
+    "fix-branches-with-liveness",
+    cl::desc("use liveness analysis during branch fixup "
              "(needed for branch inversion on AArch64)"),
     cl::init(false), cl::cat(BoltCategory));
 
diff --git a/bolt/test/AArch64/compare-and-branch-inversion.S b/bolt/test/AArch64/compare-and-branch-inversion.S
index 0ea084e904101..871fce6c67f5f 100644
--- a/bolt/test/AArch64/compare-and-branch-inversion.S
+++ b/bolt/test/AArch64/compare-and-branch-inversion.S
@@ -15,7 +15,7 @@
 # RUN: llvm-objdump -d %t.bolt | FileCheck %s --check-prefix=COMMON --check-prefix=NO-LIVENESS
 #
 # RUN: llvm-bolt -v=1 %t -o %t.bolt --data %t.fdata --reorder-blocks=ext-tsp --compact-code-model \
-# RUN:   --liveness-analysis | FileCheck %s --check-prefix=BOLT-INFO-LIVENESS
+# RUN:   --fix-branches-with-liveness | FileCheck %s --check-prefix=BOLT-INFO-LIVENESS
 # RUN: llvm-objdump -d %t.bolt | FileCheck %s --check-prefix=COMMON --check-prefix=LIVENESS
 
 # COMMON: Disassembly of section .text:
diff --git a/bolt/unittests/Core/CMakeLists.txt b/bolt/unittests/Core/CMakeLists.txt
index b755afc8a66da..297dec7449202 100644
--- a/bolt/unittests/Core/CMakeLists.txt
+++ b/bolt/unittests/Core/CMakeLists.txt
@@ -24,7 +24,6 @@ target_link_libraries(CoreTests
   PRIVATE
   LLVMBOLTCore
   LLVMBOLTRewrite
-  LLVMBOLTPasses
   LLVMBOLTProfile
   LLVMBOLTUtils
   )
diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
index 5fdaf47094770..1c206152c1fe8 100644
--- a/bolt/unittests/Core/MCPlusBuilder.cpp
+++ b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -18,9 +18,6 @@
 
 #include "bolt/Core/BinaryBasicBlock.h"
 #include "bolt/Core/BinaryFunction.h"
-#include "bolt/Core/BinaryFunctionCallGraph.h"
-#include "bolt/Passes/BinaryPasses.h"
-#include "bolt/Passes/DataflowInfoManager.h"
 #include "bolt/Rewrite/RewriteInstance.h"
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/DebugInfo/DWARF/DWARFContext.h"
@@ -342,15 +339,9 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Underflows) {
                                   .addImm(0)
                                   .addExpr(MCSymbolRefExpr::create(
                                       TargetBB->getLabel(), *BC->Ctx.get())));
-  BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
-  RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
-  DataflowInfoManager DIM(*BF, &RA, nullptr);
-  BranchLivenessInfo BranchLiveness =
-      BC->MIB->createBranchLivenessInfo(*BF, DIM);
-
-  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &BranchLiveness));
+  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, /*PreserveFlags=*/false));
   BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
-                                  BC->Ctx.get(), &BranchLiveness);
+                                  BC->Ctx.get(), /*PreserveFlags=*/false);
   I = EntryBB->begin();
   ASSERT_EQ(I->getOpcode(), AArch64::SUBSXri);
   ASSERT_EQ(I->getOperand(0).getReg(), AArch64::XZR);
@@ -383,15 +374,9 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Overflows) {
                                   .addImm(63)
                                   .addExpr(MCSymbolRefExpr::create(
                                       TargetBB->getLabel(), *BC->Ctx.get())));
-  BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
-  RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
-  DataflowInfoManager DIM(*BF, &RA, nullptr);
-  BranchLivenessInfo BranchLiveness =
-      BC->MIB->createBranchLivenessInfo(*BF, DIM);
-
-  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, &BranchLiveness));
+  ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, /*PreserveFlags=*/false));
   BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
-                                  BC->Ctx.get(), &BranchLiveness);
+                                  BC->Ctx.get(), /*PreserveFlags=*/false);
   I = EntryBB->begin();
   ASSERT_EQ(I->getOpcode(), AArch64::SUBSWri);
   ASSERT_EQ(I->getOperand(0).getReg(), AArch64::WZR);
@@ -435,13 +420,7 @@ TEST_P(MCPlusBuilderTester, AArch64_IsReversibleBranch_LiveCondFlags) {
                                     .addReg(AArch64::X2)
                                     .addImm(13));
 
-  BinaryFunctionCallGraph CG(buildCallGraph(*BC.get()));
-  RegAnalysis RA(*BC.get(), &BC->getBinaryFunctions(), &CG);
-  DataflowInfoManager DIM(*BF, &RA, nullptr);
-  BranchLivenessInfo BranchLiveness =
-      BC->MIB->createBranchLivenessInfo(*BF, DIM);
-
-  ASSERT_FALSE(BC->MIB->isReversibleBranch(*I, &BranchLiveness));
+  ASSERT_FALSE(BC->MIB->isReversibleBranch(*I, /*PreserveFlags=*/true));
 }
 
 TEST_P(MCPlusBuilderTester, AArch64_CmpJE) {

>From 01278d952e31ecdad43a5c06a10205a30b654273 Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Tue, 4 Aug 2026 13:07:18 +0100
Subject: [PATCH 5/9] wrap dense set in branch liveness info

---
 bolt/include/bolt/Core/BinaryFunction.h       |  4 ++-
 bolt/include/bolt/Core/BranchLiveness.h       | 30 +++++++++++++++++++
 .../include/bolt/Passes/BranchLivenessUtils.h | 13 ++++----
 bolt/include/bolt/Passes/LongJmp.h            | 10 +++----
 bolt/lib/Core/BinaryFunction.cpp              |  6 ++--
 bolt/lib/Passes/BinaryPasses.cpp              | 10 +++----
 bolt/lib/Passes/BranchLivenessUtils.cpp       | 15 +++++-----
 bolt/lib/Passes/LongJmp.cpp                   | 24 +++++++--------
 8 files changed, 68 insertions(+), 44 deletions(-)
 create mode 100644 bolt/include/bolt/Core/BranchLiveness.h

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 6128af7f5a941..4cd4c842a7153 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -67,6 +67,8 @@ class DWARFUnit;
 
 namespace bolt {
 
+struct BranchLivenessInfo;
+
 using InputOffsetToAddressMapTy = std::unordered_multimap<uint64_t, uint64_t>;
 
 /// Types of macro-fusion alignment corrections.
@@ -2505,7 +2507,7 @@ class BinaryFunction {
   /// while the second successor - false/fall-through branch.
   ///
   /// When we reverse the branch condition, the CFG is updated accordingly.
-  void fixBranches(const DenseSet<const MCInst *> *DeadFlagBranches = nullptr);
+  void fixBranches(const BranchLivenessInfo *BLI = nullptr);
 
   /// Mark function as finalized. No further optimizations are permitted.
   void setFinalized() { CurrentState = State::CFG_Finalized; }
diff --git a/bolt/include/bolt/Core/BranchLiveness.h b/bolt/include/bolt/Core/BranchLiveness.h
new file mode 100644
index 0000000000000..80b9f3b272aac
--- /dev/null
+++ b/bolt/include/bolt/Core/BranchLiveness.h
@@ -0,0 +1,30 @@
+//===- bolt/Core/BranchLiveness.h -------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef BOLT_CORE_BRANCHLIVENESS_H
+#define BOLT_CORE_BRANCHLIVENESS_H
+
+#include "llvm/ADT/DenseSet.h"
+
+namespace llvm {
+class MCInst;
+
+namespace bolt {
+
+struct BranchLivenessInfo {
+  DenseSet<const MCInst *> BranchesWithDeadFlags;
+
+  bool mustPreserveFlags(const MCInst &Inst) const {
+    return !BranchesWithDeadFlags.count(&Inst);
+  }
+};
+
+} // namespace bolt
+} // namespace llvm
+
+#endif
diff --git a/bolt/include/bolt/Passes/BranchLivenessUtils.h b/bolt/include/bolt/Passes/BranchLivenessUtils.h
index aca9df6ef0b5c..e94cde4295afe 100644
--- a/bolt/include/bolt/Passes/BranchLivenessUtils.h
+++ b/bolt/include/bolt/Passes/BranchLivenessUtils.h
@@ -9,21 +9,18 @@
 #ifndef BOLT_PASSES_BRANCHLIVENESSUTILS_H
 #define BOLT_PASSES_BRANCHLIVENESSUTILS_H
 
-#include "llvm/ADT/DenseSet.h"
+#include "bolt/Core/BranchLiveness.h"
 
 namespace llvm {
-class MCInst;
-
 namespace bolt {
 class BinaryFunction;
 class RegAnalysis;
 
-/// Return true if \p BF has short-range branches.
-bool hasShortRangeBranch(BinaryFunction &BF);
+/// Return true if \p BF needs liveness info for branch transformations.
+bool needsBranchLiveness(BinaryFunction &BF);
 
-/// Return the branch instructions where the target flags register is dead.
-DenseSet<const MCInst *> computeDeadFlagBranches(BinaryFunction &BF,
-                                                 RegAnalysis &RA);
+/// Return liveness info required for branch transformations.
+BranchLivenessInfo computeBranchLiveness(BinaryFunction &BF, RegAnalysis &RA);
 
 } // namespace bolt
 } // namespace llvm
diff --git a/bolt/include/bolt/Passes/LongJmp.h b/bolt/include/bolt/Passes/LongJmp.h
index 9aedfcf44b731..e9f2b29ff9588 100644
--- a/bolt/include/bolt/Passes/LongJmp.h
+++ b/bolt/include/bolt/Passes/LongJmp.h
@@ -10,13 +10,12 @@
 #define BOLT_PASSES_LONGJMP_H
 
 #include "bolt/Passes/BinaryPasses.h"
-#include "llvm/ADT/DenseSet.h"
 
 namespace llvm {
-class MCInst;
-
 namespace bolt {
 
+struct BranchLivenessInfo;
+
 /// LongJmp is veneer-insertion pass originally written for AArch64 that
 /// compensates for its short-range branches, typically done during linking. We
 /// pull this pass inside BOLT because here we can do a better job at stub
@@ -77,9 +76,8 @@ class LongJmpPass : public BinaryFunctionPass {
   /// Relax all internal function branches including those between fragments.
   /// Assume that fragments are placed in different sections but are within
   /// 128MB of each other.
-  void relaxLocalBranches(
-      BinaryFunction &BF,
-      const DenseSet<const MCInst *> *DeadFlagBranches = nullptr);
+  void relaxLocalBranches(BinaryFunction &BF,
+                          const BranchLivenessInfo *BLI = nullptr);
 
   ///                 -- Layout estimation methods --
   /// Try to do layout before running the emitter, by looking at BinaryFunctions
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index be5f7f9d4cd0e..f9c48c56f9775 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -12,6 +12,7 @@
 
 #include "bolt/Core/BinaryFunction.h"
 #include "bolt/Core/BinaryBasicBlock.h"
+#include "bolt/Core/BranchLiveness.h"
 #include "bolt/Core/DynoStats.h"
 #include "bolt/Core/HashUtilities.h"
 #include "bolt/Core/MCPlusBuilder.h"
@@ -3676,8 +3677,7 @@ bool BinaryFunction::validateCFG() const {
   return true;
 }
 
-void BinaryFunction::fixBranches(
-    const DenseSet<const MCInst *> *DeadFlagBranches) {
+void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
   assert(isSimple() && "Expected function with valid CFG.");
 
   auto &MIB = BC.MIB;
@@ -3737,7 +3737,7 @@ void BinaryFunction::fixBranches(
       // Reverse branch condition and swap successors.
       auto swapSuccessors = [&]() {
         bool MustPreserveFlags =
-            !DeadFlagBranches || !DeadFlagBranches->count(CondBranch);
+            BLI ? BLI->mustPreserveFlags(*CondBranch) : true;
         if (!MIB->isReversibleBranch(*CondBranch, MustPreserveFlags)) {
           if (opts::Verbosity) {
             BC.outs() << "BOLT-INFO: unable to swap successors in " << *this
diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index 2197c7342be6c..2637993ae4d4d 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -18,7 +18,6 @@
 #include "bolt/Passes/ReorderAlgorithm.h"
 #include "bolt/Passes/ReorderFunctions.h"
 #include "bolt/Utils/CommandLineOpts.h"
-#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/CommandLine.h"
 #include <atomic>
@@ -562,7 +561,7 @@ Error FixupBranches::runOnFunctions(BinaryContext &BC) {
       opts::FixBranchesWithLiveness &&
       llvm::any_of(BC.getBinaryFunctions(), [&](auto &It) {
         BinaryFunction &BF = It.second;
-        return BC.shouldEmit(BF) && BF.isSimple() && hasShortRangeBranch(BF);
+        return BC.shouldEmit(BF) && BF.isSimple() && needsBranchLiveness(BF);
       });
 
   std::optional<RegAnalysis> RA;
@@ -574,9 +573,10 @@ Error FixupBranches::runOnFunctions(BinaryContext &BC) {
     if (!BC.shouldEmit(BF) || !BF.isSimple())
       continue;
 
-    DenseSet<const MCInst *> DeadFlagBranches =
-        RA ? computeDeadFlagBranches(BF, *RA) : DenseSet<const MCInst *>();
-    BF.fixBranches(&DeadFlagBranches);
+    BranchLivenessInfo BLI;
+    if (RA)
+      BLI = computeBranchLiveness(BF, *RA);
+    BF.fixBranches(&BLI);
   }
   return Error::success();
 }
diff --git a/bolt/lib/Passes/BranchLivenessUtils.cpp b/bolt/lib/Passes/BranchLivenessUtils.cpp
index 9ab5de3570509..6bc1c159377bc 100644
--- a/bolt/lib/Passes/BranchLivenessUtils.cpp
+++ b/bolt/lib/Passes/BranchLivenessUtils.cpp
@@ -19,7 +19,7 @@
 namespace llvm {
 namespace bolt {
 
-bool hasShortRangeBranch(BinaryFunction &BF) {
+bool needsBranchLiveness(BinaryFunction &BF) {
   BinaryContext &BC = BF.getBinaryContext();
   if (!BC.isAArch64())
     return false;
@@ -30,27 +30,26 @@ bool hasShortRangeBranch(BinaryFunction &BF) {
   });
 }
 
-DenseSet<const MCInst *> computeDeadFlagBranches(BinaryFunction &BF,
-                                                 RegAnalysis &RA) {
+BranchLivenessInfo computeBranchLiveness(BinaryFunction &BF, RegAnalysis &RA) {
   BinaryContext &BC = BF.getBinaryContext();
-  SmallVector<MCInst *, 4> Insts;
+  SmallVector<MCInst *> Insts;
   if (BC.isAArch64())
     for (BinaryBasicBlock &BB : BF)
       for (MCInst &Inst : BB)
         if (BC.MIB->isShortRangeBranch(Inst))
           Insts.push_back(&Inst);
 
-  DenseSet<const MCInst *> DeadFlagBranches;
+  BranchLivenessInfo BLI;
   if (Insts.empty())
-    return DeadFlagBranches;
+    return BLI;
 
   DataflowInfoManager DIM(BF, &RA, nullptr);
   LivenessAnalysis &LA = DIM.getLivenessAnalysis();
   const MCPhysReg FlagsReg = BC.MIB->getFlagsReg();
   for (MCInst *Inst : Insts)
     if (!LA.getLiveIn(*Inst).test(FlagsReg))
-      DeadFlagBranches.insert(Inst);
-  return DeadFlagBranches;
+      BLI.BranchesWithDeadFlags.insert(Inst);
+  return BLI;
 }
 
 } // namespace bolt
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 497988728a527..1b2cb912c574e 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -15,7 +15,6 @@
 #include "bolt/Passes/BranchLivenessUtils.h"
 #include "bolt/Passes/RegAnalysis.h"
 #include "bolt/Utils/CommandLineOpts.h"
-#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/MathExtras.h"
 
@@ -666,8 +665,8 @@ Error LongJmpPass::relax(BinaryFunction &Func, bool &Modified) {
   return Error::success();
 }
 
-void LongJmpPass::relaxLocalBranches(
-    BinaryFunction &BF, const DenseSet<const MCInst *> *DeadFlagBranches) {
+void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
+                                     const BranchLivenessInfo *BLI) {
   BinaryContext &BC = BF.getBinaryContext();
   auto &MIB = BC.MIB;
 
@@ -843,8 +842,7 @@ void LongJmpPass::relaxLocalBranches(
       // If the other successor is a fall-through, invert the condition code.
       BinaryBasicBlock *NextBB =
           BF->getLayout().getBasicBlockAfter(BB, /*IgnoreSplits*/ false);
-      bool MustPreserveFlags =
-          !DeadFlagBranches || !DeadFlagBranches->count(&Inst);
+      bool MustPreserveFlags = BLI ? BLI->mustPreserveFlags(Inst) : true;
       bool IsReversibleBranch =
           MIB->isReversibleBranch(Inst, MustPreserveFlags);
       bool ShouldReverseBranch = BB->getConditionalSuccessor(false) == NextBB;
@@ -960,25 +958,25 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
           opts::SplitStrategy != opts::SplitFunctionsStrategy::CDSplit) &&
          "LongJmp cannot work with functions split in more than two fragments");
 
-  DenseMap<BinaryFunction *, DenseSet<const MCInst *>> DeadFlagBranches;
+  DenseMap<BinaryFunction *, BranchLivenessInfo> BranchLiveness;
   if (opts::FixBranchesWithLiveness) {
     SmallVector<BinaryFunction *> Candidates;
     for (auto &It : BC.getBinaryFunctions()) {
       BinaryFunction &BF = It.second;
       if (!BC.shouldEmit(BF) || !BF.isSimple())
         continue;
-      if (hasShortRangeBranch(BF))
+      if (needsBranchLiveness(BF))
         Candidates.push_back(&BF);
     }
     if (!Candidates.empty()) {
       RegAnalysis RA(BC, nullptr, nullptr);
       for (BinaryFunction *BF : Candidates)
-        DeadFlagBranches[BF] = computeDeadFlagBranches(*BF, RA);
+        BranchLiveness[BF] = computeBranchLiveness(*BF, RA);
     }
   }
-  auto getDeadFlagBranches = [&](BinaryFunction &BF) {
-    auto It = DeadFlagBranches.find(&BF);
-    return It == DeadFlagBranches.end() ? nullptr : &It->second;
+  auto getBranchLiveness = [&](BinaryFunction &BF) {
+    auto It = BranchLiveness.find(&BF);
+    return It == BranchLiveness.end() ? nullptr : &It->second;
   };
 
   if (opts::CompactCodeModel) {
@@ -991,7 +989,7 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
         };
 
     ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
-      relaxLocalBranches(BF, getDeadFlagBranches(BF));
+      relaxLocalBranches(BF, getBranchLiveness(BF));
     };
 
     ParallelUtilities::runOnEachFunction(
@@ -1016,7 +1014,7 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
       // Don't ruin non-simple functions, they can't afford to have the layout
       // changed.
       if (Modified && Func->isSimple())
-        Func->fixBranches(getDeadFlagBranches(*Func));
+        Func->fixBranches(getBranchLiveness(*Func));
     }
   } while (Modified);
   BC.outs() << "BOLT-INFO: Inserted " << NumHotStubs

>From 42c20fdc5d034269b879df2df868461cf351960c Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Tue, 4 Aug 2026 13:41:06 +0100
Subject: [PATCH 6/9] nfc, stylistic changes

---
 bolt/lib/Core/BinaryFunction.cpp |  7 +++----
 bolt/lib/Passes/LongJmp.cpp      | 19 ++++++++-----------
 2 files changed, 11 insertions(+), 15 deletions(-)

diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index f9c48c56f9775..121a5e52c8c50 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3736,9 +3736,8 @@ void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
 
       // Reverse branch condition and swap successors.
       auto swapSuccessors = [&]() {
-        bool MustPreserveFlags =
-            BLI ? BLI->mustPreserveFlags(*CondBranch) : true;
-        if (!MIB->isReversibleBranch(*CondBranch, MustPreserveFlags)) {
+        bool PreserveFlags = BLI ? BLI->mustPreserveFlags(*CondBranch) : true;
+        if (!MIB->isReversibleBranch(*CondBranch, PreserveFlags)) {
           if (opts::Verbosity) {
             BC.outs() << "BOLT-INFO: unable to swap successors in " << *this
                       << '\n';
@@ -3749,7 +3748,7 @@ void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
         BB->swapConditionalSuccessors();
         auto L = BC.scopeLock();
         MIB->reverseBranchCondition(BB, *CondBranch, TSuccessor->getLabel(),
-                                    Ctx, MustPreserveFlags);
+                                    Ctx, PreserveFlags);
         return true;
       };
 
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 1b2cb912c574e..39743e2e17d90 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -842,9 +842,8 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
       // If the other successor is a fall-through, invert the condition code.
       BinaryBasicBlock *NextBB =
           BF->getLayout().getBasicBlockAfter(BB, /*IgnoreSplits*/ false);
-      bool MustPreserveFlags = BLI ? BLI->mustPreserveFlags(Inst) : true;
-      bool IsReversibleBranch =
-          MIB->isReversibleBranch(Inst, MustPreserveFlags);
+      bool PreserveFlags = BLI ? BLI->mustPreserveFlags(Inst) : true;
+      bool IsReversibleBranch = MIB->isReversibleBranch(Inst, PreserveFlags);
       bool ShouldReverseBranch = BB->getConditionalSuccessor(false) == NextBB;
 
       // Create a trampoline basic block for the fall-through target of the
@@ -862,7 +861,7 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
         {
           auto L = BC.scopeLock();
           MIB->reverseBranchCondition(BB, Inst, NextBB->getLabel(),
-                                      BC.Ctx.get(), MustPreserveFlags);
+                                      BC.Ctx.get(), PreserveFlags);
         }
         const uint64_t NewBBSize = BB->estimateSize();
 
@@ -963,9 +962,7 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
     SmallVector<BinaryFunction *> Candidates;
     for (auto &It : BC.getBinaryFunctions()) {
       BinaryFunction &BF = It.second;
-      if (!BC.shouldEmit(BF) || !BF.isSimple())
-        continue;
-      if (needsBranchLiveness(BF))
+      if (BC.shouldEmit(BF) && BF.isSimple() && needsBranchLiveness(BF))
         Candidates.push_back(&BF);
     }
     if (!Candidates.empty()) {
@@ -983,15 +980,15 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
     BC.outs()
         << "BOLT-INFO: relaxing branches for compact code model (<128MB)\n";
 
+    ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
+      relaxLocalBranches(BF, getBranchLiveness(BF));
+    };
+
     ParallelUtilities::PredicateTy SkipPredicate =
         [&](const BinaryFunction &BF) {
           return !BC.shouldEmit(BF) || !BF.isSimple();
         };
 
-    ParallelUtilities::WorkFuncTy WorkFun = [&](BinaryFunction &BF) {
-      relaxLocalBranches(BF, getBranchLiveness(BF));
-    };
-
     ParallelUtilities::runOnEachFunction(
         BC, ParallelUtilities::SchedulingPolicy::SP_INST_LINEAR, WorkFun,
         SkipPredicate, "RelaxLocalBranches");

>From e7529da65cdbcae97161744e2caf15b6a407c03d Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Tue, 4 Aug 2026 16:08:39 +0100
Subject: [PATCH 7/9] remove unused include

---
 bolt/include/bolt/Core/BinaryFunction.h | 1 -
 1 file changed, 1 deletion(-)

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 4cd4c842a7153..8dc98630d8c39 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -36,7 +36,6 @@
 #include "bolt/Core/MCPlus.h"
 #include "bolt/Utils/NameResolver.h"
 #include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallString.h"

>From a4311cf8cf6c250b4e6471398b7a55b85ea08095 Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Wed, 5 Aug 2026 12:51:37 +0100
Subject: [PATCH 8/9] Return reversed branch sequences from MCPlusBuilder

Make reverseBranchCondition return an instruction sequence instead of
mutating the parent basic block directly. Callers now splice the returned
replacement at the original instruction, keeping MCPlusBuilder out of
basic block mutation and avoiding surprises from instruction storage
invalidation.
---
 bolt/include/bolt/Core/MCPlusBuilder.h        | 13 +--
 bolt/lib/Core/BinaryFunction.cpp              |  5 +-
 bolt/lib/Passes/BinaryPasses.cpp              | 15 ++-
 bolt/lib/Passes/BranchLivenessUtils.cpp       |  9 +-
 bolt/lib/Passes/LongJmp.cpp                   |  5 +-
 .../Target/AArch64/AArch64MCPlusBuilder.cpp   | 10 +-
 bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp  |  7 +-
 bolt/lib/Target/X86/X86MCPlusBuilder.cpp      |  7 +-
 bolt/unittests/Core/MCPlusBuilder.cpp         | 93 ++++++++++---------
 9 files changed, 91 insertions(+), 73 deletions(-)

diff --git a/bolt/include/bolt/Core/MCPlusBuilder.h b/bolt/include/bolt/Core/MCPlusBuilder.h
index 0879fa5f74b55..be0d58af14fc4 100644
--- a/bolt/include/bolt/Core/MCPlusBuilder.h
+++ b/bolt/include/bolt/Core/MCPlusBuilder.h
@@ -2151,12 +2151,13 @@ class MCPlusBuilder {
     llvm_unreachable("not implemented");
   }
 
-  /// Reverses the branch condition in Inst and update its taken target to TBB.
-  /// Assumes that the branch is reversible. It may replace Inst with a longer
-  /// instruction sequence on some targets.
-  virtual void reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
-                                      const MCSymbol *TBB, MCContext *Ctx,
-                                      bool MustPreserveFlags = true) const {
+  /// Return the instruction sequence for the reversed branch condition of
+  /// \p Inst and update its taken target to \p TBB. Assumes that the branch is
+  /// reversible. It may replace Inst with a longer instruction sequence on some
+  /// targets.
+  virtual InstructionListType
+  reverseBranchCondition(MCInst Inst, const MCSymbol *TBB, MCContext *Ctx,
+                         bool MustPreserveFlags = true) const {
     llvm_unreachable("not implemented");
   }
 
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 121a5e52c8c50..fa040d2a41d17 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3747,8 +3747,9 @@ void BinaryFunction::fixBranches(const BranchLivenessInfo *BLI) {
         std::swap(TSuccessor, FSuccessor);
         BB->swapConditionalSuccessors();
         auto L = BC.scopeLock();
-        MIB->reverseBranchCondition(BB, *CondBranch, TSuccessor->getLabel(),
-                                    Ctx, PreserveFlags);
+        InstructionListType Code = MIB->reverseBranchCondition(
+            *CondBranch, TSuccessor->getLabel(), Ctx, PreserveFlags);
+        BB->replaceInstruction(BB->findInstruction(CondBranch), Code);
         return true;
       };
 
diff --git a/bolt/lib/Passes/BinaryPasses.cpp b/bolt/lib/Passes/BinaryPasses.cpp
index 2637993ae4d4d..83c8843d6c6b3 100644
--- a/bolt/lib/Passes/BinaryPasses.cpp
+++ b/bolt/lib/Passes/BinaryPasses.cpp
@@ -573,9 +573,12 @@ Error FixupBranches::runOnFunctions(BinaryContext &BC) {
     if (!BC.shouldEmit(BF) || !BF.isSimple())
       continue;
 
-    BranchLivenessInfo BLI;
-    if (RA)
-      BLI = computeBranchLiveness(BF, *RA);
+    if (!RA) {
+      BF.fixBranches();
+      continue;
+    }
+
+    BranchLivenessInfo BLI = computeBranchLiveness(BF, *RA);
     BF.fixBranches(&BLI);
   }
   return Error::success();
@@ -987,7 +990,11 @@ uint64_t SimplifyConditionalTailCalls::fixTailCalls(BinaryFunction &BF) {
       uint64_t Count = 0;
       if (CondSucc != BB) {
         // Patch the new target address into the conditional branch.
-        MIB->reverseBranchCondition(PredBB, *CondBranch, CalleeSymbol, Ctx);
+        InstructionListType Code =
+            MIB->reverseBranchCondition(*CondBranch, CalleeSymbol, Ctx);
+        auto II = PredBB->replaceInstruction(
+            PredBB->findInstruction(CondBranch), Code);
+        CondBranch = &*(II);
         // Since we reversed the condition on the branch we need to change
         // the target for the unconditional branch or add a unconditional
         // branch to the old target.  This has to be done manually since
diff --git a/bolt/lib/Passes/BranchLivenessUtils.cpp b/bolt/lib/Passes/BranchLivenessUtils.cpp
index 6bc1c159377bc..63a0da405fc88 100644
--- a/bolt/lib/Passes/BranchLivenessUtils.cpp
+++ b/bolt/lib/Passes/BranchLivenessUtils.cpp
@@ -25,8 +25,10 @@ bool needsBranchLiveness(BinaryFunction &BF) {
     return false;
 
   return llvm::any_of(BF, [&](BinaryBasicBlock &BB) {
-    return llvm::any_of(
-        BB, [&](MCInst &Inst) { return BC.MIB->isShortRangeBranch(Inst); });
+    return llvm::any_of(BB, [&](MCInst &Inst) {
+      return BC.MIB->isShortRangeBranch(Inst) &&
+             !BC.MIB->isReversibleBranch(Inst);
+    });
   });
 }
 
@@ -36,7 +38,8 @@ BranchLivenessInfo computeBranchLiveness(BinaryFunction &BF, RegAnalysis &RA) {
   if (BC.isAArch64())
     for (BinaryBasicBlock &BB : BF)
       for (MCInst &Inst : BB)
-        if (BC.MIB->isShortRangeBranch(Inst))
+        if (BC.MIB->isShortRangeBranch(Inst) &&
+            !BC.MIB->isReversibleBranch(Inst))
           Insts.push_back(&Inst);
 
   BranchLivenessInfo BLI;
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 39743e2e17d90..06b36fd59f0e0 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -860,8 +860,9 @@ void LongJmpPass::relaxLocalBranches(BinaryFunction &BF,
         BB->swapConditionalSuccessors();
         {
           auto L = BC.scopeLock();
-          MIB->reverseBranchCondition(BB, Inst, NextBB->getLabel(),
-                                      BC.Ctx.get(), PreserveFlags);
+          InstructionListType Code = MIB->reverseBranchCondition(
+              Inst, NextBB->getLabel(), BC.Ctx.get(), PreserveFlags);
+          BB->replaceInstruction(BB->findInstruction(&Inst), Code);
         }
         const uint64_t NewBBSize = BB->estimateSize();
 
diff --git a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
index f9bd174b38814..d3eb0d4d905aa 100644
--- a/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
+++ b/bolt/lib/Target/AArch64/AArch64MCPlusBuilder.cpp
@@ -2204,9 +2204,9 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
     return MCPlusBuilder::isReversibleBranch(Inst);
   }
 
-  void reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
-                              const MCSymbol *TBB, MCContext *Ctx,
-                              bool MustPreserveFlags = true) const override {
+  InstructionListType
+  reverseBranchCondition(MCInst Inst, const MCSymbol *TBB, MCContext *Ctx,
+                         bool MustPreserveFlags = true) const override {
     assert(isReversibleBranch(Inst, MustPreserveFlags) &&
            "Irreversible branch");
 
@@ -2255,8 +2255,7 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
         Code.emplace_back(MCInstBuilder(AArch64::Bcc)
                               .addImm(getInvertedCC(Inst.getOpcode()))
                               .addExpr(MCSymbolRefExpr::create(TBB, *Ctx)));
-        Parent->replaceInstruction(Parent->findInstruction(&Inst), Code);
-        return;
+        return Code;
       }
       Inst.setOpcode(InvertedOpcode);
     } else if (Inst.getOpcode() == AArch64::Bcc) {
@@ -2270,6 +2269,7 @@ class AArch64MCPlusBuilder : public MCPlusBuilder {
       llvm_unreachable("Unrecognized branch instruction");
     }
     replaceBranchTarget(Inst, TBB, Ctx);
+    return {Inst};
   }
 
   int getPCRelEncodingSize(const MCInst &Inst) const override {
diff --git a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
index 082f9242adf22..7422320b9a28c 100644
--- a/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
+++ b/bolt/lib/Target/RISCV/RISCVMCPlusBuilder.cpp
@@ -162,12 +162,13 @@ class RISCVMCPlusBuilder : public MCPlusBuilder {
     }
   }
 
-  void reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
-                              const MCSymbol *TBB, MCContext *Ctx,
-                              bool MustPreserveFlags = true) const override {
+  InstructionListType
+  reverseBranchCondition(MCInst Inst, const MCSymbol *TBB, MCContext *Ctx,
+                         bool MustPreserveFlags = true) const override {
     auto Opcode = getInvertedBranchOpcode(Inst.getOpcode());
     Inst.setOpcode(Opcode);
     replaceBranchTarget(Inst, TBB, Ctx);
+    return {Inst};
   }
 
   void replaceBranchTarget(MCInst &Inst, const MCSymbol *TBB,
diff --git a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
index ba84617dbd65f..684bedacde3e9 100644
--- a/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
+++ b/bolt/lib/Target/X86/X86MCPlusBuilder.cpp
@@ -2811,14 +2811,15 @@ class X86MCPlusBuilder : public MCPlusBuilder {
     Inst.addOperand(MCOperand::createImm(CC));
   }
 
-  void reverseBranchCondition(BinaryBasicBlock *Parent, MCInst &Inst,
-                              const MCSymbol *TBB, MCContext *Ctx,
-                              bool MustPreserveFlags = true) const override {
+  InstructionListType
+  reverseBranchCondition(MCInst Inst, const MCSymbol *TBB, MCContext *Ctx,
+                         bool MustPreserveFlags = true) const override {
     unsigned InvCC = getInvertedCondCode(getCondCode(Inst));
     assert(InvCC != X86::COND_INVALID && "invalid branch instruction");
     Inst.getOperand(Info->get(Inst.getOpcode()).NumOperands - 1).setImm(InvCC);
     Inst.getOperand(0) =
         MCOperand::createExpr(MCSymbolRefExpr::create(TBB, *Ctx));
+    return {Inst};
   }
 
   bool replaceBranchCondition(MCInst &Inst, const MCSymbol *TBB, MCContext *Ctx,
diff --git a/bolt/unittests/Core/MCPlusBuilder.cpp b/bolt/unittests/Core/MCPlusBuilder.cpp
index 1c206152c1fe8..a692f45f551eb 100644
--- a/bolt/unittests/Core/MCPlusBuilder.cpp
+++ b/bolt/unittests/Core/MCPlusBuilder.cpp
@@ -233,10 +233,11 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                            .addExpr(MCSymbolRefExpr::create(
                                TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(NeedsImmInc));
-  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, NeedsImmInc,
-                                  TargetBB->getLabel(), BC->Ctx.get());
-  ASSERT_EQ(NeedsImmInc.getOpcode(), AArch64::CBLTXri);
-  ASSERT_EQ(NeedsImmInc.getOperand(1).getImm(), 1);
+  auto NeedsImmIncCode = BC->MIB->reverseBranchCondition(
+      NeedsImmInc, TargetBB->getLabel(), BC->Ctx.get());
+  ASSERT_EQ(NeedsImmIncCode.size(), 1u);
+  ASSERT_EQ(NeedsImmIncCode[0].getOpcode(), AArch64::CBLTXri);
+  ASSERT_EQ(NeedsImmIncCode[0].getOperand(1).getImm(), 1);
 
   // Compare register with immediate and branch.
   // Inversion requires decrementing the immediate value.
@@ -247,10 +248,11 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                            .addExpr(MCSymbolRefExpr::create(
                                TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(NeedsImmDec));
-  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, NeedsImmDec,
-                                  TargetBB->getLabel(), BC->Ctx.get());
-  ASSERT_EQ(NeedsImmDec.getOpcode(), AArch64::CBHIXri);
-  ASSERT_EQ(NeedsImmDec.getOperand(1).getImm(), 0);
+  auto NeedsImmDecCode = BC->MIB->reverseBranchCondition(
+      NeedsImmDec, TargetBB->getLabel(), BC->Ctx.get());
+  ASSERT_EQ(NeedsImmDecCode.size(), 1u);
+  ASSERT_EQ(NeedsImmDecCode[0].getOpcode(), AArch64::CBHIXri);
+  ASSERT_EQ(NeedsImmDecCode[0].getOperand(1).getImm(), 0);
 
   // Compare registers and branch.
   // Inversion requires swapping registers.
@@ -261,11 +263,12 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                                    .addExpr(MCSymbolRefExpr::create(
                                        TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(CompRegNeedsRegSwap));
-  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, CompRegNeedsRegSwap,
-                                  TargetBB->getLabel(), BC->Ctx.get());
-  ASSERT_EQ(CompRegNeedsRegSwap.getOpcode(), AArch64::CBGTXrr);
-  ASSERT_EQ(CompRegNeedsRegSwap.getOperand(0).getReg(), AArch64::X1);
-  ASSERT_EQ(CompRegNeedsRegSwap.getOperand(1).getReg(), AArch64::X0);
+  auto CompRegCode = BC->MIB->reverseBranchCondition(
+      CompRegNeedsRegSwap, TargetBB->getLabel(), BC->Ctx.get());
+  ASSERT_EQ(CompRegCode.size(), 1u);
+  ASSERT_EQ(CompRegCode[0].getOpcode(), AArch64::CBGTXrr);
+  ASSERT_EQ(CompRegCode[0].getOperand(0).getReg(), AArch64::X1);
+  ASSERT_EQ(CompRegCode[0].getOperand(1).getReg(), AArch64::X0);
 
   // Compare bytes and branch.
   // Inversion requires swapping registers.
@@ -276,11 +279,12 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                                     .addExpr(MCSymbolRefExpr::create(
                                         TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(CompByteNeedsRegSwap));
-  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, CompByteNeedsRegSwap,
-                                  TargetBB->getLabel(), BC->Ctx.get());
-  ASSERT_EQ(CompByteNeedsRegSwap.getOpcode(), AArch64::CBBHSWrr);
-  ASSERT_EQ(CompByteNeedsRegSwap.getOperand(0).getReg(), AArch64::W1);
-  ASSERT_EQ(CompByteNeedsRegSwap.getOperand(1).getReg(), AArch64::W0);
+  auto CompByteCode = BC->MIB->reverseBranchCondition(
+      CompByteNeedsRegSwap, TargetBB->getLabel(), BC->Ctx.get());
+  ASSERT_EQ(CompByteCode.size(), 1u);
+  ASSERT_EQ(CompByteCode[0].getOpcode(), AArch64::CBBHSWrr);
+  ASSERT_EQ(CompByteCode[0].getOperand(0).getReg(), AArch64::W1);
+  ASSERT_EQ(CompByteCode[0].getOperand(1).getReg(), AArch64::W0);
 
   // Compare halfwords and branch.
   // Inversion requires swapping registers.
@@ -291,11 +295,12 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch) {
                                     .addExpr(MCSymbolRefExpr::create(
                                         TargetBB->getLabel(), *BC->Ctx.get()));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(CompHalfNeedsRegSwap));
-  BC->MIB->reverseBranchCondition(/*ParentBB*/ nullptr, CompHalfNeedsRegSwap,
-                                  TargetBB->getLabel(), BC->Ctx.get());
-  ASSERT_EQ(CompHalfNeedsRegSwap.getOpcode(), AArch64::CBHHIWrr);
-  ASSERT_EQ(CompHalfNeedsRegSwap.getOperand(0).getReg(), AArch64::W1);
-  ASSERT_EQ(CompHalfNeedsRegSwap.getOperand(1).getReg(), AArch64::W0);
+  auto CompHalfCode = BC->MIB->reverseBranchCondition(
+      CompHalfNeedsRegSwap, TargetBB->getLabel(), BC->Ctx.get());
+  ASSERT_EQ(CompHalfCode.size(), 1u);
+  ASSERT_EQ(CompHalfCode[0].getOpcode(), AArch64::CBHHIWrr);
+  ASSERT_EQ(CompHalfCode[0].getOperand(0).getReg(), AArch64::W1);
+  ASSERT_EQ(CompHalfCode[0].getOperand(1).getReg(), AArch64::W0);
 
   // Compare register with immediate and branch.
   // Inversion not possible, immediate value underflows.
@@ -340,17 +345,16 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Underflows) {
                                   .addExpr(MCSymbolRefExpr::create(
                                       TargetBB->getLabel(), *BC->Ctx.get())));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, /*PreserveFlags=*/false));
-  BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
-                                  BC->Ctx.get(), /*PreserveFlags=*/false);
-  I = EntryBB->begin();
-  ASSERT_EQ(I->getOpcode(), AArch64::SUBSXri);
-  ASSERT_EQ(I->getOperand(0).getReg(), AArch64::XZR);
-  ASSERT_EQ(I->getOperand(1).getReg(), AArch64::X0);
-  ASSERT_EQ(I->getOperand(2).getImm(), 0);
-  ASSERT_EQ(I->getOperand(3).getImm(), 0);
-  I++;
-  ASSERT_EQ(I->getOpcode(), AArch64::Bcc);
-  ASSERT_EQ(I->getOperand(0).getImm(), AArch64CC::GE);
+  auto Code = BC->MIB->reverseBranchCondition(
+      *I, TargetBB->getLabel(), BC->Ctx.get(), /*PreserveFlags=*/false);
+  ASSERT_EQ(Code.size(), 2u);
+  ASSERT_EQ(Code[0].getOpcode(), AArch64::SUBSXri);
+  ASSERT_EQ(Code[0].getOperand(0).getReg(), AArch64::XZR);
+  ASSERT_EQ(Code[0].getOperand(1).getReg(), AArch64::X0);
+  ASSERT_EQ(Code[0].getOperand(2).getImm(), 0);
+  ASSERT_EQ(Code[0].getOperand(3).getImm(), 0);
+  ASSERT_EQ(Code[1].getOpcode(), AArch64::Bcc);
+  ASSERT_EQ(Code[1].getOperand(0).getImm(), AArch64CC::GE);
 }
 
 TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Overflows) {
@@ -375,17 +379,16 @@ TEST_P(MCPlusBuilderTester, AArch64_ReverseCompAndBranch_Overflows) {
                                   .addExpr(MCSymbolRefExpr::create(
                                       TargetBB->getLabel(), *BC->Ctx.get())));
   ASSERT_TRUE(BC->MIB->isReversibleBranch(*I, /*PreserveFlags=*/false));
-  BC->MIB->reverseBranchCondition(EntryBB, *I, TargetBB->getLabel(),
-                                  BC->Ctx.get(), /*PreserveFlags=*/false);
-  I = EntryBB->begin();
-  ASSERT_EQ(I->getOpcode(), AArch64::SUBSWri);
-  ASSERT_EQ(I->getOperand(0).getReg(), AArch64::WZR);
-  ASSERT_EQ(I->getOperand(1).getReg(), AArch64::W0);
-  ASSERT_EQ(I->getOperand(2).getImm(), 63);
-  ASSERT_EQ(I->getOperand(3).getImm(), 0);
-  I++;
-  ASSERT_EQ(I->getOpcode(), AArch64::Bcc);
-  ASSERT_EQ(I->getOperand(0).getImm(), AArch64CC::LS);
+  auto Code = BC->MIB->reverseBranchCondition(
+      *I, TargetBB->getLabel(), BC->Ctx.get(), /*PreserveFlags=*/false);
+  ASSERT_EQ(Code.size(), 2u);
+  ASSERT_EQ(Code[0].getOpcode(), AArch64::SUBSWri);
+  ASSERT_EQ(Code[0].getOperand(0).getReg(), AArch64::WZR);
+  ASSERT_EQ(Code[0].getOperand(1).getReg(), AArch64::W0);
+  ASSERT_EQ(Code[0].getOperand(2).getImm(), 63);
+  ASSERT_EQ(Code[0].getOperand(3).getImm(), 0);
+  ASSERT_EQ(Code[1].getOpcode(), AArch64::Bcc);
+  ASSERT_EQ(Code[1].getOperand(0).getImm(), AArch64CC::LS);
 }
 
 TEST_P(MCPlusBuilderTester, AArch64_IsReversibleBranch_LiveCondFlags) {

>From 459c2f455d2aa2428c108946b6ba5b04816c50cb Mon Sep 17 00:00:00 2001
From: Alexandros Lamprineas <alexandros.lamprineas at arm.com>
Date: Wed, 5 Aug 2026 14:45:24 +0100
Subject: [PATCH 9/9] Store branch liveness info as annotations

---
 bolt/include/bolt/Core/BinaryFunction.h |  2 +-
 bolt/include/bolt/Core/BranchLiveness.h | 26 +++++++---
 bolt/include/bolt/Passes/LongJmp.h      |  2 +-
 bolt/lib/Core/BranchLiveness.cpp        | 69 +++++++++++++++++++++++++
 bolt/lib/Core/CMakeLists.txt            |  1 +
 bolt/lib/Passes/BranchLivenessUtils.cpp |  4 +-
 bolt/lib/Passes/LongJmp.cpp             |  2 +-
 7 files changed, 94 insertions(+), 12 deletions(-)
 create mode 100644 bolt/lib/Core/BranchLiveness.cpp

diff --git a/bolt/include/bolt/Core/BinaryFunction.h b/bolt/include/bolt/Core/BinaryFunction.h
index 8dc98630d8c39..2846b562a1c1a 100644
--- a/bolt/include/bolt/Core/BinaryFunction.h
+++ b/bolt/include/bolt/Core/BinaryFunction.h
@@ -66,7 +66,7 @@ class DWARFUnit;
 
 namespace bolt {
 
-struct BranchLivenessInfo;
+class BranchLivenessInfo;
 
 using InputOffsetToAddressMapTy = std::unordered_multimap<uint64_t, uint64_t>;
 
diff --git a/bolt/include/bolt/Core/BranchLiveness.h b/bolt/include/bolt/Core/BranchLiveness.h
index 80b9f3b272aac..584533e69ed96 100644
--- a/bolt/include/bolt/Core/BranchLiveness.h
+++ b/bolt/include/bolt/Core/BranchLiveness.h
@@ -9,19 +9,31 @@
 #ifndef BOLT_CORE_BRANCHLIVENESS_H
 #define BOLT_CORE_BRANCHLIVENESS_H
 
-#include "llvm/ADT/DenseSet.h"
-
 namespace llvm {
 class MCInst;
 
 namespace bolt {
+class BinaryFunction;
+
+class BranchLivenessInfo {
+  BinaryFunction *BF;
+  unsigned AnnotationIndex;
+
+  void swap(BranchLivenessInfo &Other) noexcept;
+
+public:
+  explicit BranchLivenessInfo(BinaryFunction &BF);
+  ~BranchLivenessInfo();
+
+  // Copies would create multiple owners for removing the same annotations.
+  BranchLivenessInfo(const BranchLivenessInfo &) = delete;
+  BranchLivenessInfo &operator=(const BranchLivenessInfo &) = delete;
 
-struct BranchLivenessInfo {
-  DenseSet<const MCInst *> BranchesWithDeadFlags;
+  BranchLivenessInfo(BranchLivenessInfo &&Other) noexcept;
+  BranchLivenessInfo &operator=(BranchLivenessInfo &&Other) noexcept;
 
-  bool mustPreserveFlags(const MCInst &Inst) const {
-    return !BranchesWithDeadFlags.count(&Inst);
-  }
+  bool mustPreserveFlags(const MCInst &Inst) const;
+  void setFlagsDead(MCInst &Inst);
 };
 
 } // namespace bolt
diff --git a/bolt/include/bolt/Passes/LongJmp.h b/bolt/include/bolt/Passes/LongJmp.h
index e9f2b29ff9588..0a7c4d5f33334 100644
--- a/bolt/include/bolt/Passes/LongJmp.h
+++ b/bolt/include/bolt/Passes/LongJmp.h
@@ -14,7 +14,7 @@
 namespace llvm {
 namespace bolt {
 
-struct BranchLivenessInfo;
+class BranchLivenessInfo;
 
 /// LongJmp is veneer-insertion pass originally written for AArch64 that
 /// compensates for its short-range branches, typically done during linking. We
diff --git a/bolt/lib/Core/BranchLiveness.cpp b/bolt/lib/Core/BranchLiveness.cpp
new file mode 100644
index 0000000000000..db0d24aae98c5
--- /dev/null
+++ b/bolt/lib/Core/BranchLiveness.cpp
@@ -0,0 +1,69 @@
+//===- bolt/Core/BranchLiveness.cpp --------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "bolt/Core/BranchLiveness.h"
+#include "bolt/Core/BinaryBasicBlock.h"
+#include "bolt/Core/BinaryContext.h"
+#include "bolt/Core/BinaryFunction.h"
+#include "bolt/Core/MCPlusBuilder.h"
+#include "llvm/MC/MCInst.h"
+#include <cassert>
+#include <utility>
+
+namespace llvm {
+namespace bolt {
+
+BranchLivenessInfo::BranchLivenessInfo(BinaryFunction &BF)
+    : BF(&BF),
+      AnnotationIndex(
+          BF.getBinaryContext().MIB->getOrCreateAnnotationIndex("DeadFlags")) {}
+
+BranchLivenessInfo::~BranchLivenessInfo() {
+  if (!BF)
+    return;
+
+  MCPlusBuilder &MIB = *BF->getBinaryContext().MIB;
+  for (BinaryBasicBlock &BB : *BF)
+    for (MCInst &Inst : BB)
+      MIB.removeAnnotation(Inst, AnnotationIndex);
+}
+
+BranchLivenessInfo::BranchLivenessInfo(BranchLivenessInfo &&Other) noexcept
+    : BF(nullptr), AnnotationIndex(0) {
+  swap(Other);
+}
+
+BranchLivenessInfo &
+BranchLivenessInfo::operator=(BranchLivenessInfo &&Other) noexcept {
+  BranchLivenessInfo Tmp(std::move(Other));
+  swap(Tmp);
+  return *this;
+}
+
+void BranchLivenessInfo::swap(BranchLivenessInfo &Other) noexcept {
+  std::swap(BF, Other.BF);
+  std::swap(AnnotationIndex, Other.AnnotationIndex);
+}
+
+bool BranchLivenessInfo::mustPreserveFlags(const MCInst &Inst) const {
+  if (!BF)
+    return true;
+
+  return !BF->getBinaryContext().MIB->hasAnnotation(Inst, AnnotationIndex);
+}
+
+void BranchLivenessInfo::setFlagsDead(MCInst &Inst) {
+  assert(BF && "branch liveness info is not initialized");
+
+  MCPlusBuilder &MIB = *BF->getBinaryContext().MIB;
+  if (!MIB.hasAnnotation(Inst, AnnotationIndex))
+    MIB.addAnnotation(Inst, AnnotationIndex, true);
+}
+
+} // namespace bolt
+} // namespace llvm
diff --git a/bolt/lib/Core/CMakeLists.txt b/bolt/lib/Core/CMakeLists.txt
index 58cfcab370f16..459151b86721b 100644
--- a/bolt/lib/Core/CMakeLists.txt
+++ b/bolt/lib/Core/CMakeLists.txt
@@ -18,6 +18,7 @@ add_llvm_library(LLVMBOLTCore
   BinaryData.cpp
   BinaryEmitter.cpp
   BinaryFunction.cpp
+  BranchLiveness.cpp
   BinaryFunctionCallGraph.cpp
   BinaryFunctionProfile.cpp
   BinarySection.cpp
diff --git a/bolt/lib/Passes/BranchLivenessUtils.cpp b/bolt/lib/Passes/BranchLivenessUtils.cpp
index 63a0da405fc88..a187ca3c98757 100644
--- a/bolt/lib/Passes/BranchLivenessUtils.cpp
+++ b/bolt/lib/Passes/BranchLivenessUtils.cpp
@@ -42,7 +42,7 @@ BranchLivenessInfo computeBranchLiveness(BinaryFunction &BF, RegAnalysis &RA) {
             !BC.MIB->isReversibleBranch(Inst))
           Insts.push_back(&Inst);
 
-  BranchLivenessInfo BLI;
+  BranchLivenessInfo BLI(BF);
   if (Insts.empty())
     return BLI;
 
@@ -51,7 +51,7 @@ BranchLivenessInfo computeBranchLiveness(BinaryFunction &BF, RegAnalysis &RA) {
   const MCPhysReg FlagsReg = BC.MIB->getFlagsReg();
   for (MCInst *Inst : Insts)
     if (!LA.getLiveIn(*Inst).test(FlagsReg))
-      BLI.BranchesWithDeadFlags.insert(Inst);
+      BLI.setFlagsDead(*Inst);
   return BLI;
 }
 
diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index 06b36fd59f0e0..f4b975f2ff4f9 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -969,7 +969,7 @@ Error LongJmpPass::runOnFunctions(BinaryContext &BC) {
     if (!Candidates.empty()) {
       RegAnalysis RA(BC, nullptr, nullptr);
       for (BinaryFunction *BF : Candidates)
-        BranchLiveness[BF] = computeBranchLiveness(*BF, RA);
+        BranchLiveness.try_emplace(BF, computeBranchLiveness(*BF, RA));
     }
   }
   auto getBranchLiveness = [&](BinaryFunction &BF) {



More information about the llvm-commits mailing list