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

[via llvm-commits]

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

If a local stub is out-of-range, at LongJmp we will try to find another local stub first. However, there seems to be a problem with the code logic.


https://github.com/llvm/llvm-project/blob/main/bolt/lib/Passes/LongJmp.cpp#L203-L209
```
  } 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
    TgtSym = BC.MIB->getTargetSymbol(*TgtBB->begin()); // (1)
    TgtBB = BB.getSuccessor(TgtSym, BI);                              // (2)
  }
```

`TgtSym` now is the target of the local stub (statement 1), and thus it is not a successor of `BB` (statement 2), and thus TgtBB will be set to nullptr.


>From 06f4b4dee335ef76febbae645061fb91328daa6b 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 | 27 ++++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/bolt/lib/Passes/LongJmp.cpp b/bolt/lib/Passes/LongJmp.cpp
index a81689bc37469a4..a9aed870617ca05 100644
--- a/bolt/lib/Passes/LongJmp.cpp
+++ b/bolt/lib/Passes/LongJmp.cpp
@@ -202,10 +202,31 @@ 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 BB and label of an out-of-range
+    // local stub. We replace with the target of the local stub instead
+    // of creating a stub to jump to another stub.
+    // e.g.
+    // change the out-of-range stub
+    //    .Ltmp1:
+    //        ldr	x8, [x8, #0x10]
+    //        tbz	w0, #0x0, .LStub1111
+    //        ...
+    //    .LStub1111:
+    //        b	.Ltmp22
+    //  into the old target(target of the stub)
+    //    .Ltmp1:
+    //        ldr	x8, [x8, #0x10]
+    //        tbz	w0, #0x0, .Ltmp22
+    //        ...
+    //    .LStub1111:
+    //        b	.Ltmp22
     TgtSym = BC.MIB->getTargetSymbol(*TgtBB->begin());
-    TgtBB = BB.getSuccessor(TgtSym, BI);
+    BB.replaceSuccessor(TgtBB, TgtBB->getSuccessor(TgtSym, BI), 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 = TgtBB->getSuccessor(TgtSym, BI);
   }
 
   BinaryBasicBlock *StubBB = lookupLocalStub(BB, Inst, TgtSym, DotAddress);