Bug type: crash
Component: DFAJumpThreading
Bug location: llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp: unfold
Issue title: [DFAJumpThreading] crash since b167ada
Labels: crash
llvm:transforms
Reported by: JonPsson1

Seeing many assertions that go away if I revert b167ada "[DFAJumpThreading] Rewrite the way paths are enumerated (#96127)".

```
clang-20 -cc1 -triple s390x-unknown-linux-gnu -target-cpu z15 -O3 -w -mllvm -enable-dfa-jump-thread -o /dev/null -emit-llvm -x ir  tc_dfa-jt_invalidbbarg.ll
...
10 0x000002aa0a73e544 (anonymous namespace)::DFAJumpThreading::run
...
```

```
clang-20 -cc1 -triple s390x-unknown-linux-gnu   -target-cpu z16 -O3 -w -mllvm -enable-dfa-jump-thread -o /dev/null -emit-llvm -x ir  tc_select_mustbedeadnow.ll
...
#10 0x000002aa1b0c0968 llvm::DFAJumpThreadingPass::run
...

```
[tcs_dfa-jt.tar.gz](https://github.com/user-attachments/files/16750050/tcs_dfa-jt.tar.gz)


Comments: author: mikaelholmen
body: Perhaps the same assertion I reported here:
 https://github.com/llvm/llvm-project/pull/96127#issuecomment-2303862037
?
author: JonPsson1
body: > Perhaps the same assertion I reported here: [#96127 (comment)](https://github.com/llvm/llvm-project/pull/96127#issuecomment-2303862037) ?

Yes, same assertion, and in addition my second test case above triggered another assertion.

Fix patch:
commit d4a38c8ff5c993e14c42895b51a47272fb03a857
Author: Usman Nadeem <mnadeem@quicinc.com>
Date:   Tue Sep 24 08:54:36 2024 -0700

    [DFAJumpThreading] Handle select unfolding when user phi is not a dir… (#109511)
    
    …ect successor
    
    Previously the code assumed that the select instruction is defined in a
    block that is a direct predecessor of the block where the PHINode uses
    it. So, we were hitting an assertion when we tried to access the def
    block as an incoming block for the user phi node.
    
    This patch handles that case by using the correct end block and creating
    a new phi node that aggregates both the values of the select in that end
    block, and then using that new unfolded phi to overwrite the original
    user phi node.
    
    Fixes #106083
    
    Change-Id: Ie471994cca232318f74a6e6438efa21e561c2dc0

diff --git a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
index ef9c264482a6..0e2b5c925a6a 100644
--- a/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/DFAJumpThreading.cpp
@@ -194,7 +194,6 @@ void unfold(DomTreeUpdater *DTU, LoopInfo *LI, SelectInstToUnfold SIToUnfold,
   SelectInst *SI = SIToUnfold.getInst();
   PHINode *SIUse = SIToUnfold.getUse();
   BasicBlock *StartBlock = SI->getParent();
-  BasicBlock *EndBlock = SIUse->getParent();
   BranchInst *StartBlockTerm =
       dyn_cast<BranchInst>(StartBlock->getTerminator());
 
@@ -202,6 +201,7 @@ void unfold(DomTreeUpdater *DTU, LoopInfo *LI, SelectInstToUnfold SIToUnfold,
   assert(SI->hasOneUse());
 
   if (StartBlockTerm->isUnconditional()) {
+    BasicBlock *EndBlock = StartBlock->getUniqueSuccessor();
     // Arbitrarily choose the 'false' side for a new input value to the PHI.
     BasicBlock *NewBlock = BasicBlock::Create(
         SI->getContext(), Twine(SI->getName(), ".si.unfold.false"),
@@ -223,32 +223,44 @@ void unfold(DomTreeUpdater *DTU, LoopInfo *LI, SelectInstToUnfold SIToUnfold,
                                       NewBlock->getFirstInsertionPt());
     NewPhi->addIncoming(SIOp2, StartBlock);
 
-    if (auto *OpSi = dyn_cast<SelectInst>(SIOp1))
-      NewSIsToUnfold->push_back(SelectInstToUnfold(OpSi, SIUse));
-    if (auto *OpSi = dyn_cast<SelectInst>(SIOp2))
-      NewSIsToUnfold->push_back(SelectInstToUnfold(OpSi, NewPhi));
-
-    // Update the phi node of SI.
-    for (unsigned Idx = 0; Idx < SIUse->getNumIncomingValues(); ++Idx) {
-      if (SIUse->getIncomingBlock(Idx) == StartBlock)
-        SIUse->setIncomingValue(Idx, SIOp1);
+    // Update any other PHI nodes in EndBlock.
+    for (PHINode &Phi : EndBlock->phis()) {
+      if (SIUse == &Phi)
+        continue;
+      Phi.addIncoming(Phi.getIncomingValueForBlock(StartBlock), NewBlock);
     }
-    SIUse->addIncoming(NewPhi, NewBlock);
 
-    // Update any other PHI nodes in EndBlock.
-    for (auto II = EndBlock->begin(); PHINode *Phi = dyn_cast<PHINode>(II);
-         ++II) {
-      if (Phi != SIUse)
-        Phi->addIncoming(Phi->getIncomingValueForBlock(StartBlock), NewBlock);
+    // Update the phi node of SI, which is its only use.
+    if (EndBlock == SIUse->getParent()) {
+      SIUse->addIncoming(NewPhi, NewBlock);
+      SIUse->replaceUsesOfWith(SI, SIOp1);
+    } else {
+      PHINode *EndPhi = PHINode::Create(SIUse->getType(), pred_size(EndBlock),
+                                        Twine(SI->getName(), ".si.unfold.phi"),
+                                        EndBlock->getFirstInsertionPt());
+      for (BasicBlock *Pred : predecessors(EndBlock)) {
+        if (Pred != StartBlock && Pred != NewBlock)
+          EndPhi->addIncoming(EndPhi, Pred);
+      }
+
+      EndPhi->addIncoming(SIOp1, StartBlock);
+      EndPhi->addIncoming(NewPhi, NewBlock);
+      SIUse->replaceUsesOfWith(SI, EndPhi);
+      SIUse = EndPhi;
     }
 
-    StartBlockTerm->eraseFromParent();
+    if (auto *OpSi = dyn_cast<SelectInst>(SIOp1))
+      NewSIsToUnfold->push_back(SelectInstToUnfold(OpSi, SIUse));
+    if (auto *OpSi = dyn_cast<SelectInst>(SIOp2))
+      NewSIsToUnfold->push_back(SelectInstToUnfold(OpSi, NewPhi));
 
     // Insert the real conditional branch based on the original condition.
+    StartBlockTerm->eraseFromParent();
     BranchInst::Create(EndBlock, NewBlock, SI->getCondition(), StartBlock);
     DTU->applyUpdates({{DominatorTree::Insert, StartBlock, EndBlock},
                        {DominatorTree::Insert, StartBlock, NewBlock}});
   } else {
+    BasicBlock *EndBlock = SIUse->getParent();
     BasicBlock *NewBlockT = BasicBlock::Create(
         SI->getContext(), Twine(SI->getName(), ".si.unfold.true"),
         EndBlock->getParent(), EndBlock);


Regression tests:
file: llvm/test/Transforms/DFAJumpThreading/dfa-jump-threading-transform.ll
commands: opt -S -passes=dfa-jump-threading %s
tests: test_name: test2
test_body: define i32 @test2(i32 %init) {
entry:
  %cmp = icmp eq i32 %init, 0
  %sel = select i1 %cmp, i32 0, i32 2
  br label %loop.1

loop.1:                                           ; preds = %loop.1.backedge, %entry
  %state.1 = phi i32 [ %sel, %entry ], [ %state.1.be2, %loop.1.backedge ]
  br label %loop.2

loop.2:                                           ; preds = %loop.2.backedge, %loop.1
  %state.2 = phi i32 [ %state.1, %loop.1 ], [ %state.2.be, %loop.2.backedge ]
  br label %loop.3

loop.3:                                           ; preds = %case2, %loop.2
  %state = phi i32 [ %state.2, %loop.2 ], [ 3, %case2 ]
  switch i32 %state, label %infloop.i [
    i32 2, label %case2
    i32 3, label %case3
    i32 4, label %case4
    i32 0, label %case0
    i32 1, label %case1
  ]

case2:                                            ; preds = %loop.3
  br i1 %cmp, label %loop.3, label %loop.1.backedge

case3:                                            ; preds = %loop.3
  br i1 %cmp, label %loop.2.backedge, label %case4

case4:                                            ; preds = %case3, %loop.3
  br i1 %cmp, label %loop.2.backedge, label %loop.1.backedge

loop.1.backedge:                                  ; preds = %case4, %case2
  %state.1.be = phi i32 [ 2, %case4 ], [ 4, %case2 ]
  %state.1.be2 = select i1 %cmp, i32 1, i32 %state.1.be
  br label %loop.1

loop.2.backedge:                                  ; preds = %case4, %case3
  %state.2.be = phi i32 [ 3, %case4 ], [ 0, %case3 ]
  br label %loop.2

case0:                                            ; preds = %loop.3
  br label %exit

case1:                                            ; preds = %loop.3
  br label %exit

infloop.i:                                        ; preds = %infloop.i, %loop.3
  br label %infloop.i

exit:                                             ; preds = %case1, %case0
  ret i32 0
}

test_name: pr106083_select_dead_uses
test_body: define void @pr106083_select_dead_uses(i1 %cmp1, i1 %not, ptr %p) {
bb:
  %spec.select = select i1 %cmp1, i32 0, i32 1
  br label %.loopexit6

.loopexit6:                                       ; preds = %select.unfold, %bb
  br i1 %not, label %select.unfold, label %bb1

bb1:                                              ; preds = %.loopexit6
  %i = load i32, ptr %p, align 4
  %not2 = icmp eq i32 0, 0
  %spec.select7 = select i1 %not2, i32 %spec.select, i32 0
  br label %select.unfold

select.unfold:                                    ; preds = %bb1, %.loopexit6
  %_2 = phi i32 [ 0, %.loopexit6 ], [ %spec.select7, %bb1 ]
  switch i32 %_2, label %bb2 [
    i32 0, label %.preheader.preheader
    i32 1, label %.loopexit6
  ]

.preheader.preheader:                             ; preds = %select.unfold
  ret void

bb2:                                              ; preds = %select.unfold
  unreachable
}

test_name: test1
test_body: define i32 @test1(i32 %num) {
entry:
  br label %for.body

for.body:                                         ; preds = %for.inc, %entry
  %count = phi i32 [ 0, %entry ], [ %inc, %for.inc ]
  %state = phi i32 [ 1, %entry ], [ %state.next, %for.inc ]
  switch i32 %state, label %for.inc [
    i32 1, label %case1
    i32 2, label %case2
  ]

case1:                                            ; preds = %for.body
  br label %for.inc

case2:                                            ; preds = %for.body
  %cmp = icmp eq i32 %count, 50
  %sel = select i1 %cmp, i32 1, i32 2
  br label %for.inc

for.inc:                                          ; preds = %case2, %case1, %for.body
  %state.next = phi i32 [ %sel, %case2 ], [ 1, %for.body ], [ 2, %case1 ]
  %inc = add nsw i32 %count, 1
  %cmp.exit = icmp slt i32 %inc, %num
  br i1 %cmp.exit, label %for.body, label %for.end

for.end:                                          ; preds = %for.inc
  ret i32 0
}

test_name: pr106083_invalidBBarg_fold
test_body: define void @pr106083_invalidBBarg_fold(i1 %cmp1, i1 %cmp2, i1 %not, ptr %d) {
bb:
  %sel = select i1 %cmp1, i32 0, i32 1
  br label %BB1

BB1:                                              ; preds = %BB1.backedge, %BB7, %bb
  %i = phi i16 [ 0, %BB1.backedge ], [ 0, %bb ], [ 1, %BB7 ]
  br i1 %not, label %BB7, label %BB2

BB2:                                              ; preds = %BB1
  store i16 0, ptr %d, align 2
  %spec.select = select i1 %cmp2, i32 %sel, i32 0
  br label %BB7

BB7:                                              ; preds = %BB2, %BB1
  %d.promoted4 = phi i16 [ 0, %BB1 ], [ 1, %BB2 ]
  %_3 = phi i32 [ 0, %BB1 ], [ %spec.select, %BB2 ]
  switch i32 %_3, label %BB1.backedge [
    i32 0, label %BB1
    i32 1, label %BB8
  ]

BB1.backedge:                                     ; preds = %BB7
  br label %BB1

BB8:                                              ; preds = %BB7
  ret void
}

test_name: pr78059_bitwidth
test_body: define void @pr78059_bitwidth() {
.split.preheader:
  br label %.split

.split:                                           ; preds = %.split, %.split.preheader
  %0 = phi i128 [ 0, %.split.preheader ], [ -1, %.split ]
  switch i128 %0, label %end [
    i128 -1, label %end
    i128 0, label %.split
  ]

end:                                              ; preds = %.split, %.split
  ret void
}