[BOLT] Split getBasicBlocksAfter cache into a distinct function

This enables future re-use in other code that calls getBasicBlockAfter
in loops, though for now those uses aren't introduced.

Author: Mark-Simulacrum

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,17 +3591,7 @@ void BinaryFunction::fixBranches() {
   auto &MIB = BC.MIB;
   MCContext *Ctx = BC.Ctx.get();
 
-  // Caches `FunctionLayout::nextBasicBlock(IgnoreSplits = false)`.
-  // nextBasicBlock uses linear search to find the next block, so the loop
-  // below becomes O(n^2). This avoids that.
-  DenseMap<BinaryBasicBlock *, BinaryBasicBlock *> nextBasicBlock(
-      Layout.block_size());
-  for (size_t i = 0; i + 1 < Layout.block_size(); i++) {
-    auto current = Layout.block_begin() + i;
-    auto next = Layout.block_begin() + i + 1;
-    if (next != Layout.getFragment((*current)->getFragmentNum()).end())
-      nextBasicBlock.insert(std::pair(*current, *next));
-  }
+  auto NextBasicBlock = Layout.getBasicBlocksAfter(/* IgnoreSplits */ false);
 
   for (BinaryBasicBlock *BB : BasicBlocks) {
     const MCSymbol *TBB = nullptr;
@@ -3617,7 +3607,7 @@ void BinaryFunction::fixBranches() {
 
     // Basic block that follows the current one in the final layout.
     const BinaryBasicBlock *const NextBB =
-        nextBasicBlock.lookup_or(BB, nullptr);
+        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(); });