From 607b22594c68201c988482c15123c8b42cdd36f1 Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks@fb.com>
Date: Tue, 23 Sep 2025 12:25:08 -0700
Subject: [PATCH] [BOLT] Avoid n^2 complexity in fixBranches(). NFCI

Iterator implementation of PR #156243:

This improves BOLT runtime when optimizing rustc_driver.so from 15
minutes to 7 minutes (or 49 minutes to 37 minutes of userspace time).

Co-authored-by: Mark-Simulacrum <mark.simulacrum@gmail.com>
---
 bolt/include/bolt/Core/FunctionLayout.h | 21 ++++++++++++++++-----
 bolt/lib/Core/BinaryFunction.cpp        |  6 ++++--
 bolt/lib/Core/FunctionLayout.cpp        | 20 +++++++++++++-------
 3 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/bolt/include/bolt/Core/FunctionLayout.h b/bolt/include/bolt/Core/FunctionLayout.h
index ee4dd689b8dd6..240d5138a093a 100644
--- a/bolt/include/bolt/Core/FunctionLayout.h
+++ b/bolt/include/bolt/Core/FunctionLayout.h
@@ -232,8 +232,24 @@ class FunctionLayout {
     return Blocks[Index];
   }
 
+  /// Return the basic block after the given basic block iterator in the layout
+  /// or nullptr if the last basic block iterator is given.
+  const BinaryBasicBlock *getBasicBlockAfter(block_const_iterator BlockIt,
+                                             bool IgnoreSplits = true) const;
+
+  /// Returns the basic block after the given basic block in the layout or
+  /// nullptr if the last basic block is given.
+  ///
+  /// Note: prefer the version that takes the iterator as this function uses
+  /// linear basic block lookup.
+  const BinaryBasicBlock *getBasicBlockAfter(const BinaryBasicBlock *BB,
+                                             bool IgnoreSplits = true) const;
+
   /// Returns the basic block after the given basic block in the layout or
   /// nullptr if the last basic block is given.
+  ///
+  /// Note: prefer the version that takes the iterator as this function uses
+  /// linear basic block lookup.
   BinaryBasicBlock *getBasicBlockAfter(const BinaryBasicBlock *const BB,
                                        const bool IgnoreSplits = true) {
     return const_cast<BinaryBasicBlock *>(
@@ -241,11 +257,6 @@ class FunctionLayout {
             BB, IgnoreSplits));
   }
 
-  /// Returns the basic block after the given basic block in the layout or
-  /// nullptr if the last basic block is given.
-  const BinaryBasicBlock *getBasicBlockAfter(const BinaryBasicBlock *BB,
-                                             bool IgnoreSplits = true) const;
-
   /// True if the layout contains at least two non-empty fragments.
   bool isSplit() const;
 
diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 578a87dc6c09d..7102054a366fa 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3598,7 +3598,9 @@ void BinaryFunction::fixBranches() {
   auto &MIB = BC.MIB;
   MCContext *Ctx = BC.Ctx.get();
 
-  for (BinaryBasicBlock *BB : BasicBlocks) {
+  for (auto BBI = Layout.block_begin(), BBE = Layout.block_end(); BBI != BBE;
+       ++BBI) {
+    BinaryBasicBlock *BB = *BBI;
     const MCSymbol *TBB = nullptr;
     const MCSymbol *FBB = nullptr;
     MCInst *CondBranch = nullptr;
@@ -3612,7 +3614,7 @@ void BinaryFunction::fixBranches() {
 
     // Basic block that follows the current one in the final layout.
     const BinaryBasicBlock *const NextBB =
-        Layout.getBasicBlockAfter(BB, /*IgnoreSplits=*/false);
+        Layout.getBasicBlockAfter(BBI, /*IgnoreSplits*/ false);
 
     if (BB->succ_size() == 1) {
       // __builtin_unreachable() could create a conditional branch that
diff --git a/bolt/lib/Core/FunctionLayout.cpp b/bolt/lib/Core/FunctionLayout.cpp
index 4498fc44da954..98ed6e1320b3e 100644
--- a/bolt/lib/Core/FunctionLayout.cpp
+++ b/bolt/lib/Core/FunctionLayout.cpp
@@ -224,23 +224,29 @@ void FunctionLayout::clear() {
 }
 
 const BinaryBasicBlock *
-FunctionLayout::getBasicBlockAfter(const BinaryBasicBlock *BB,
+FunctionLayout::getBasicBlockAfter(block_const_iterator BBIter,
                                    bool IgnoreSplits) const {
-  const block_const_iterator BBPos = find(blocks(), BB);
-  if (BBPos == block_end())
-    return nullptr;
-
-  const block_const_iterator BlockAfter = std::next(BBPos);
+  const block_const_iterator BlockAfter = std::next(BBIter);
   if (BlockAfter == block_end())
     return nullptr;
 
   if (!IgnoreSplits)
-    if (BlockAfter == getFragment(BB->getFragmentNum()).end())
+    if (BlockAfter == getFragment((*BBIter)->getFragmentNum()).end())
       return nullptr;
 
   return *BlockAfter;
 }
 
+const BinaryBasicBlock *
+FunctionLayout::getBasicBlockAfter(const BinaryBasicBlock *BB,
+                                   bool IgnoreSplits) const {
+  const block_const_iterator BBPos = find(blocks(), BB);
+  if (BBPos == block_end())
+    return nullptr;
+
+  return getBasicBlockAfter(BBPos, IgnoreSplits);
+}
+
 bool FunctionLayout::isSplit() const {
   const unsigned NonEmptyFragCount = llvm::count_if(
       fragments(), [](const FunctionFragment &FF) { return !FF.empty(); });