[llvm] Branch folding: Do not use both TII::analyzeBranch and MBB::isSuccessor (PR #104240)

[via llvm-commits]

https://github.com/v01dXYZ updated https://github.com/llvm/llvm-project/pull/104240

>From e6b35ba2148a883a16e85beebff005f3851f441f Mon Sep 17 00:00:00 2001
From: v01dxyz <v01dxyz at v01d.xyz>
Date: Thu, 27 Jun 2024 21:33:57 +0200
Subject: [PATCH] Branch folding: Do not use both TII::analyzeBranch and
 MBB::isSuccessor

Avoid infinite loop with EHPad blocks or indirect target blocks.

After a ``!TII->analyzeBranch(PrevBB, PrevTBB, PrevFBB, PrevCond,
true)``, the code previously used ``PrevBB.isSuccessor(&*MBB)`` to
check if MBB is PrevTBB or PrevFBB.

But ``TII->analyzeBranch(...)`` only supports unconditional/conditional
branches while INVOKE or CALLBR could also add blocks as successors (EHPad
blocks for INVOKE, Indirect Target Blocks for CALLBR).
---
 llvm/lib/CodeGen/BranchFolding.cpp             | 15 ++-------------
 ...-loop-with-callbr-indirect-target-blocks.ll | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 13 deletions(-)
 create mode 100644 llvm/test/CodeGen/X86/branchfolding-no-inf-loop-with-callbr-indirect-target-blocks.ll

diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index 1dc278586f1178..db6ad189cd80dc 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -1422,8 +1422,7 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
     // This has to check PrevBB->succ_size() because EH edges are ignored by
     // analyzeBranch.
     if (PriorCond.empty() && !PriorTBB && MBB->pred_size() == 1 &&
-        PrevBB.succ_size() == 1 && PrevBB.isSuccessor(MBB) &&
-        !MBB->hasAddressTaken() && !MBB->isEHPad()) {
+        PrevBB.succ_size() == 1 && !MBB->hasAddressTaken()) {
       LLVM_DEBUG(dbgs() << "\nMerging into block: " << PrevBB
                         << "From MBB: " << *MBB);
       // Remove redundant DBG_VALUEs first.
@@ -1756,21 +1755,11 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
       // removed, move this block to the end of the function. There is no real
       // advantage in "falling through" to an EH block, so we don't want to
       // perform this transformation for that case.
-      //
-      // Also, Windows EH introduced the possibility of an arbitrary number of
-      // successors to a given block.  The analyzeBranch call does not consider
-      // exception handling and so we can get in a state where a block
-      // containing a call is followed by multiple EH blocks that would be
-      // rotated infinitely at the end of the function if the transformation
-      // below were performed for EH "FallThrough" blocks.  Therefore, even if
-      // that appears not to be happening anymore, we should assume that it is
-      // possible and not remove the "!FallThrough()->isEHPad" condition below.
       MachineBasicBlock *PrevTBB = nullptr, *PrevFBB = nullptr;
       SmallVector<MachineOperand, 4> PrevCond;
       if (FallThrough != MF.end() &&
-          !FallThrough->isEHPad() &&
           !TII->analyzeBranch(PrevBB, PrevTBB, PrevFBB, PrevCond, true) &&
-          PrevBB.isSuccessor(&*FallThrough)) {
+          ((PrevTBB == &*FallThrough) || (PrevFBB == &*FallThrough))) {
         MBB->moveAfter(&MF.back());
         MadeChange = true;
         return MadeChange;
diff --git a/llvm/test/CodeGen/X86/branchfolding-no-inf-loop-with-callbr-indirect-target-blocks.ll b/llvm/test/CodeGen/X86/branchfolding-no-inf-loop-with-callbr-indirect-target-blocks.ll
new file mode 100644
index 00000000000000..5a9557c3e343f8
--- /dev/null
+++ b/llvm/test/CodeGen/X86/branchfolding-no-inf-loop-with-callbr-indirect-target-blocks.ll
@@ -0,0 +1,18 @@
+; RUN: llc -mtriple x86_64 %s -stop-after=branch-folder
+
+; Check branch-folder do not keep swapping %indirect.target.block.1 and
+; %indirect.target.block.2
+define void @no_inf_loop() {
+entry:
+  br label %bb1
+
+bb1:
+  callbr void asm "", "!i,!i"() to label %bb1
+    [label %indirect.target.block.1, label %indirect.target.block.2]
+
+indirect.target.block.1:
+  br label %bb1
+
+indirect.target.block.2:
+  br label %bb1
+}