[llvm] [BOLT] Optimize basic block loops to avoid n^2 loop (PR #156243)

[Mark Rousskov via llvm-commits]

https://github.com/Mark-Simulacrum updated https://github.com/llvm/llvm-project/pull/156243

>From 480d48e42720915de2f9958b69ad05df7c0bf802 Mon Sep 17 00:00:00 2001
From: Mark Rousskov <mark.simulacrum at gmail.com>
Date: Sun, 31 Aug 2025 09:03:27 -0400
Subject: [PATCH] [BOLT] Optimize basic block loops to avoid n^2 loop

This improves BOLT runtime when optimizing rustc_driver.so from 15
minutes to 7 minutes (49 minutes to 37 minutes of userspace time).
---
 bolt/lib/Core/BinaryFunction.cpp | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/bolt/lib/Core/BinaryFunction.cpp b/bolt/lib/Core/BinaryFunction.cpp
index 6cac2d0cca2cb..a86e204cae974 100644
--- a/bolt/lib/Core/BinaryFunction.cpp
+++ b/bolt/lib/Core/BinaryFunction.cpp
@@ -3591,6 +3591,18 @@ 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));
+  }
+
   for (BinaryBasicBlock *BB : BasicBlocks) {
     const MCSymbol *TBB = nullptr;
     const MCSymbol *FBB = nullptr;
@@ -3605,7 +3617,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