[llvm] [BOLT] Fix local out-of-range stub issue that leads to infinite loop in LongJmp pass (PR #73918)

[via llvm-commits]

https://github.com/linsinan1995 updated https://github.com/llvm/llvm-project/pull/73918

>From f03a10ddebc09578570a3767a4be36b07eb64640 Mon Sep 17 00:00:00 2001
From: Lin Sinan <linsinan.lsn at alibaba-inc.com>
Date: Tue, 28 Nov 2023 00:51:19 +0800
Subject: [PATCH] [BOLT] Fix local out-of-range stub issue that leads to
 infinite loop in LongJmp pass

fix an unexpected behavior in replaceTargetWithStub at LongJmp pass if
the local stub is out-of-range.
---
 bolt/lib/Passes/LongJmp.cpp | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index a81689bc37469..8528b82bd0d21 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -202,10 +202,21 @@ LongJmpPass::replaceTargetWithStub(BinaryBasicBlock &BB, MCInst &Inst,
     }
   } else if (LocalStubsIter != Stubs.end() &&
              LocalStubsIter->second.count(TgtBB)) {
-    // If we are replacing a local stub (because it is now out of range),
-    // use its target instead of creating a stub to jump to another stub
+    // The TgtBB and TgtSym now are the local out-of-range stub and its label.
+    // So, we are attempting to restore BB to its previous state without using
+    // this stub.
     TgtSym = BC.MIB->getTargetSymbol(*TgtBB->begin());
-    TgtBB = BB.getSuccessor(TgtSym, BI);
+    BinaryBasicBlock *TgtBBSucc = TgtBB->getSuccessor(TgtSym, BI);
+
+    // TgtBB might have no successor. e.g. a stub for a function call.
+    if (TgtBBSucc) {
+      BB.replaceSuccessor(TgtBB, TgtBBSucc, BI.Count, BI.MispredictedCount);
+      assert(TgtBB->getExecutionCount() >= BI.Count &&
+             "At least equal or greater than the branch count.");
+      TgtBB->setExecutionCount(TgtBB->getExecutionCount() - BI.Count);
+    }
+
+    TgtBB = TgtBBSucc;
   }
 
   BinaryBasicBlock *StubBB = lookupLocalStub(BB, Inst, TgtSym, DotAddress);