[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 93a4fc9e237c1312b2f28a47c6d0b22dd453594a 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            | 28 ++++---------------
 ...loop-with-callbr-indirect-target-blocks.ll | 18 ++++++++++++
 2 files changed, 24 insertions(+), 22 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..8620d59c182f00 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -1354,20 +1354,14 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
   // explicitly.  Landing pads should not do this since the landing-pad table
   // points to this block.  Blocks with their addresses taken shouldn't be
   // optimized away.
-  if (IsEmptyBlock(MBB) && !MBB->isEHPad() && !MBB->hasAddressTaken() &&
-      SameEHScope) {
+  if (IsEmptyBlock(MBB) && !MBB->hasAddressTaken() && SameEHScope) {
     salvageDebugInfoFromEmptyBlock(TII, *MBB);
     // Dead block?  Leave for cleanup later.
     if (MBB->pred_empty()) return MadeChange;
 
     if (FallThrough == MF.end()) {
       // TODO: Simplify preds to not branch here if possible!
-    } else if (FallThrough->isEHPad()) {
-      // Don't rewrite to a landing pad fallthough.  That could lead to the case
-      // where a BB jumps to more than one landing pad.
-      // TODO: Is it ever worth rewriting predecessors which don't already
-      // jump to a landing pad, and so can safely jump to the fallthrough?
-    } else if (MBB->isSuccessor(&*FallThrough)) {
+    } else if ((CurTBB == &*FallThrough) || (CurFBB == &*FallThrough)) {
       // Rewrite all predecessors of the old block to go to the fallthrough
       // instead.
       while (!MBB->pred_empty()) {
@@ -1422,8 +1416,8 @@ 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 && ((PriorTBB == MBB) || (PriorFBB == MBB)) &&
+        !MBB->hasAddressTaken()) {
       LLVM_DEBUG(dbgs() << "\nMerging into block: " << PrevBB
                         << "From MBB: " << *MBB);
       // Remove redundant DBG_VALUEs first.
@@ -1612,7 +1606,7 @@ bool BranchFolder::OptimizeBlock(MachineBasicBlock *MBB) {
       if (MBB->empty()) {
         bool PredHasNoFallThrough = !PrevBB.canFallThrough();
         if (PredHasNoFallThrough || !PriorUnAnalyzable ||
-            !PrevBB.isSuccessor(MBB)) {
+            !((CurTBB == MBB) || (CurFBB == MBB))) {
           // If the prior block falls through into us, turn it into an
           // explicit branch to us to make updates simpler.
           if (!PredHasNoFallThrough && PrevBB.isSuccessor(MBB) &&
@@ -1756,21 +1750,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
+}