[llvm] [BOLT] a local out-of-range stub might lead to infinite loop in LongJmp (PR #73918)

[via llvm-commits]

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

>From cd5c4f4ab7ca7e9db6a8908e3e8f00a40031f80b 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 might lead 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 | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index a81689bc37469..34a15461e15a0 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -202,10 +202,20 @@ 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);
+    auto *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);