[BOLT] Avoid n^2 complexity in fixBranches(). NFCI

Iterator implementation of PR llvm#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 <[email protected]>

Author: maksfb

bolt/include/bolt/Core/FunctionLayout.h

@@ -232,20 +232,31 @@ 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 *>(
         static_cast<const FunctionLayout &>(*this).getBasicBlockAfter(
             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;
 

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

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(); });