[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 e3f587f8d4d0ce55633b4d179c4e08346f5d0c0b Mon Sep 17 00:00:00 2001
From: Sinan Lin <sinan.lin at linux.alibaba.com>
Date: Tue, 28 Nov 2023 00:51:19 +0800
Subject: [PATCH] [BOLT] Fix local out-of-range stub issue in LongJmp

If a local stub is out-of-range, at LongJmp we will try to find another
local stub first. However, The original implementation do not work as
expected and it leads to an infinite loop between replaceTargetWithStub
and fixBranches.

After this patch, we first convert the target of BB back to the target
of the local stub, and then look up for other valid local stubs and so
on.
---
 bolt/lib/Passes/LongJmp.cpp | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index a81689bc37469..ded0db2cd30b6 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -202,10 +202,23 @@ 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);
+    assert(TgtSym &&
+           "First instruction is expected to contain a target symbol.");
+    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);