[BOLT] Optimize basic block loops to avoid n^2 loop

bolt/include/bolt/Core/FunctionLayout.h

@@ -243,9 +243,18 @@ class FunctionLayout {
 
   /// Returns the basic block after the given basic block in the layout or
   /// nullptr if the last basic block is given.
+  ///
+  /// Note that this performs a linear search for BB.
   const BinaryBasicBlock *getBasicBlockAfter(const BinaryBasicBlock *BB,
                                              bool IgnoreSplits = true) const;
 
+  /// Returns a mapping from BB -> getBasicBlockAfter(BB).
+  ///
+  /// This should be preferred in loops that call getBasicBlockAfter without
+  /// changes to the function layout. Caching the results avoid n^2 lookup cost.
+  DenseMap<BinaryBasicBlock *, BinaryBasicBlock *>
+  getBasicBlocksAfter(bool IgnoreSplits = true) const;
+
   /// True if the layout contains at least two non-empty fragments.
   bool isSplit() const;
 

bolt/lib/Core/BinaryFunction.cpp

@@ -3591,6 +3591,8 @@ void BinaryFunction::fixBranches() {
   auto &MIB = BC.MIB;
   MCContext *Ctx = BC.Ctx.get();
 
+  auto NextBasicBlock = Layout.getBasicBlocksAfter(/* IgnoreSplits */ false);
+
   for (BinaryBasicBlock *BB : BasicBlocks) {
     const MCSymbol *TBB = nullptr;
     const MCSymbol *FBB = nullptr;
@@ -3605,7 +3607,7 @@ void BinaryFunction::fixBranches() {
 
     // Basic block that follows the current one in the final layout.
     const BinaryBasicBlock *const NextBB =
-        Layout.getBasicBlockAfter(BB, /*IgnoreSplits=*/false);
+        NextBasicBlock.lookup_or(BB, nullptr);
 
     if (BB->succ_size() == 1) {
       // __builtin_unreachable() could create a conditional branch that

bolt/lib/Core/FunctionLayout.cpp

@@ -241,6 +241,22 @@ FunctionLayout::getBasicBlockAfter(const BinaryBasicBlock *BB,
   return *BlockAfter;
 }
 
+DenseMap<BinaryBasicBlock *, BinaryBasicBlock *>
+FunctionLayout::getBasicBlocksAfter(bool IgnoreSplits) const {
+  DenseMap<BinaryBasicBlock *, BinaryBasicBlock *> NextBasicBlock(block_size());
+  for (size_t i = 0; i + 1 < block_size(); i++) {
+    auto Current = block_begin() + i;
+    auto Next = block_begin() + i + 1;
+
+    if (IgnoreSplits) {
+      NextBasicBlock.insert(std::pair(*Current, *Next));
+    } else if (Next != getFragment((*Current)->getFragmentNum()).end()) {
+      NextBasicBlock.insert(std::pair(*Current, *Next));
+    }
+  }
+  return NextBasicBlock;
+}
+
 bool FunctionLayout::isSplit() const {
   const unsigned NonEmptyFragCount = llvm::count_if(
       fragments(), [](const FunctionFragment &FF) { return !FF.empty(); });