[llvm] [DFAJumpThreading] Rewrite the way paths are enumerated (PR #96127)

[via llvm-commits]

================
@@ -212,101 +200,130 @@ void unfold(DomTreeUpdater *DTU, LoopInfo *LI, SelectInstToUnfold SIToUnfold,
   BranchInst *StartBlockTerm =
       dyn_cast<BranchInst>(StartBlock->getTerminator());
 
-  assert(StartBlockTerm && StartBlockTerm->isUnconditional());
+  assert(StartBlockTerm);
   assert(SI->hasOneUse());
 
-  // These are the new basic blocks for the conditional branch.
-  // At least one will become an actual new basic block.
-  BasicBlock *TrueBlock = nullptr;
-  BasicBlock *FalseBlock = nullptr;
-  BranchInst *TrueBranch = nullptr;
-  BranchInst *FalseBranch = nullptr;
-
-  // Sink select instructions to be able to unfold them later.
-  if (SelectInst *SIOp = dyn_cast<SelectInst>(SI->getTrueValue())) {
-    createBasicBlockAndSinkSelectInst(DTU, SI, SIUse, SIOp, EndBlock,
-                                      "si.unfold.true", &TrueBlock, &TrueBranch,
-                                      NewSIsToUnfold, NewBBs);
-  }
-  if (SelectInst *SIOp = dyn_cast<SelectInst>(SI->getFalseValue())) {
-    createBasicBlockAndSinkSelectInst(DTU, SI, SIUse, SIOp, EndBlock,
-                                      "si.unfold.false", &FalseBlock,
-                                      &FalseBranch, NewSIsToUnfold, NewBBs);
-  }
-
-  // If there was nothing to sink, then arbitrarily choose the 'false' side
-  // for a new input value to the PHI.
-  if (!TrueBlock && !FalseBlock) {
-    FalseBlock = BasicBlock::Create(SI->getContext(), "si.unfold.false",
-                                    EndBlock->getParent(), EndBlock);
-    NewBBs->push_back(FalseBlock);
-    BranchInst::Create(EndBlock, FalseBlock);
-    DTU->applyUpdates({{DominatorTree::Insert, FalseBlock, EndBlock}});
-  }
-
-  // Insert the real conditional branch based on the original condition.
-  // If we did not create a new block for one of the 'true' or 'false' paths
-  // of the condition, it means that side of the branch goes to the end block
-  // directly and the path originates from the start block from the point of
-  // view of the new PHI.
-  BasicBlock *TT = EndBlock;
-  BasicBlock *FT = EndBlock;
-  if (TrueBlock && FalseBlock) {
-    // A diamond.
-    TT = TrueBlock;
-    FT = FalseBlock;
-
-    // Update the phi node of SI.
-    SIUse->addIncoming(SI->getTrueValue(), TrueBlock);
-    SIUse->addIncoming(SI->getFalseValue(), FalseBlock);
-
-    // Update any other PHI nodes in EndBlock.
-    for (PHINode &Phi : EndBlock->phis()) {
-      if (&Phi != SIUse) {
-        Value *OrigValue = Phi.getIncomingValueForBlock(StartBlock);
-        Phi.addIncoming(OrigValue, TrueBlock);
-        Phi.addIncoming(OrigValue, FalseBlock);
-      }
-
-      // Remove incoming place of original StartBlock, which comes in a indirect
-      // way (through TrueBlock and FalseBlock) now.
-      Phi.removeIncomingValue(StartBlock, /* DeletePHIIfEmpty = */ false);
-    }
-  } else {
-    BasicBlock *NewBlock = nullptr;
+  if (StartBlockTerm->isUnconditional()) {
+    // 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"),
+        EndBlock->getParent(), EndBlock);
+    NewBBs->push_back(NewBlock);
+    BranchInst::Create(EndBlock, NewBlock);
+    DTU->applyUpdates({{DominatorTree::Insert, NewBlock, EndBlock}});
+
+    // StartBlock
+    //   |  \
+    //   |  NewBlock
+    //   |  /
+    // EndBlock
     Value *SIOp1 = SI->getTrueValue();
     Value *SIOp2 = SI->getFalseValue();
 
-    // A triangle pointing right.
-    if (!TrueBlock) {
-      NewBlock = FalseBlock;
-      FT = FalseBlock;
-    }
-    // A triangle pointing left.
-    else {
-      NewBlock = TrueBlock;
-      TT = TrueBlock;
-      std::swap(SIOp1, SIOp2);
-    }
+    PHINode *NewPhi = PHINode::Create(SIUse->getType(), 1,
+                                      Twine(SIOp2->getName(), ".si.unfold.phi"),
+                                      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);
     }
-    SIUse->addIncoming(SIOp2, 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);
     }
+
+    StartBlockTerm->eraseFromParent();
+
+    // Insert the real conditional branch based on the original condition.
+    BranchInst::Create(EndBlock, NewBlock, SI->getCondition(), StartBlock);
+    DTU->applyUpdates({{DominatorTree::Insert, StartBlock, EndBlock},
+                       {DominatorTree::Insert, StartBlock, NewBlock}});
+  } else {
+    BasicBlock *NewBlockT = BasicBlock::Create(
+        SI->getContext(), Twine(SI->getName(), ".si.unfold.true"),
+        EndBlock->getParent(), EndBlock);
+    BasicBlock *NewBlockF = BasicBlock::Create(
+        SI->getContext(), Twine(SI->getName(), ".si.unfold.false"),
+        EndBlock->getParent(), EndBlock);
+
+    NewBBs->push_back(NewBlockT);
+    NewBBs->push_back(NewBlockF);
+
+    // Def only has one use in EndBlock.
+    // Before transformation:
+    // StartBlock(Def)
+    //   |      \
+    // EndBlock  OtherBlock
+    //  (Use)
+    //
+    // After transformation:
+    // StartBlock(Def)
+    //   |      \
+    //   |       OtherBlock
+    // NewBlockT
+    //   |     \
+    //   |   NewBlockF
+    //   |      /
+    //   |     /
+    // EndBlock
+    //  (Use)
+    BranchInst::Create(EndBlock, NewBlockF);
+    // Insert the real conditional branch based on the original condition.
+    BranchInst::Create(EndBlock, NewBlockF, SI->getCondition(), NewBlockT);
+    DTU->applyUpdates({{DominatorTree::Insert, NewBlockT, NewBlockF},
+                       {DominatorTree::Insert, NewBlockT, EndBlock},
+                       {DominatorTree::Insert, NewBlockF, EndBlock}});
+
+    Value *TrueVal = SI->getTrueValue();
+    Value *FalseVal = SI->getFalseValue();
+
+    PHINode *NewPhiT = PHINode::Create(
+        SIUse->getType(), 1, Twine(TrueVal->getName(), ".si.unfold.phi"),
+        NewBlockT->getFirstInsertionPt());
+    PHINode *NewPhiF = PHINode::Create(
+        SIUse->getType(), 1, Twine(FalseVal->getName(), ".si.unfold.phi"),
+        NewBlockF->getFirstInsertionPt());
+    NewPhiT->addIncoming(TrueVal, StartBlock);
+    NewPhiF->addIncoming(FalseVal, NewBlockT);
+
+    if (auto *TrueSI = dyn_cast<SelectInst>(TrueVal))
+      NewSIsToUnfold->push_back(SelectInstToUnfold(TrueSI, NewPhiT));
+    if (auto *FalseSi = dyn_cast<SelectInst>(FalseVal))
+      NewSIsToUnfold->push_back(SelectInstToUnfold(FalseSi, NewPhiF));
+
+    SIUse->addIncoming(NewPhiT, NewBlockT);
+    SIUse->addIncoming(NewPhiF, NewBlockF);
+    SIUse->removeIncomingValue(StartBlock);
+
+    // Update any other PHI nodes in EndBlock.
+    for (auto II = EndBlock->begin(); PHINode *Phi = dyn_cast<PHINode>(II);
+         ++II) {
----------------
XChy wrote:

```suggestion
    for (PHINode& Phi : EndBlock->phis()) {
```

https://github.com/llvm/llvm-project/pull/96127