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

[via llvm-commits]

================
@@ -553,42 +572,118 @@ struct AllSwitchPaths {
       return;
     }
 
-    for (const PathType &Path : LoopPaths) {
-      ThreadingPath TPath;
-
-      const BasicBlock *PrevBB = Path.back();
-      for (const BasicBlock *BB : Path) {
-        if (StateDef.contains(BB)) {
-          const PHINode *Phi = dyn_cast<PHINode>(StateDef[BB]);
-          assert(Phi && "Expected a state-defining instr to be a phi node.");
-
-          const Value *V = Phi->getIncomingValueForBlock(PrevBB);
-          if (const ConstantInt *C = dyn_cast<const ConstantInt>(V)) {
-            TPath.setExitValue(C);
-            TPath.setDeterminator(BB);
-            TPath.setPath(Path);
-          }
-        }
+    auto *SwitchPhi = cast<PHINode>(Switch->getOperand(0));
+    auto *SwitchPhiDefBB = SwitchPhi->getParent();
+    VisitedBlocks VB;
+    // Get paths from the determinator BBs to SwitchPhiDefBB
+    std::vector<ThreadingPath> PathsToPhiDef =
+        getPathsFromStateDefMap(StateDef, SwitchPhi, VB);
+    if (SwitchPhiDefBB == SwitchBlock) {
+      TPaths = std::move(PathsToPhiDef);
+      return;
+    }
 
-        // Switch block is the determinator, this is the final exit value.
-        if (TPath.isExitValueSet() && BB == Path.front())
-          break;
+    // Find and append paths from SwitchPhiDefBB to SwitchBlock.
+    PathsType PathsToSwitchBB =
+        paths(SwitchPhiDefBB, SwitchBlock, VB, /* PathDepth = */ 1);
+    if (PathsToSwitchBB.empty())
+      return;
 
-        PrevBB = BB;
+    std::vector<ThreadingPath> TempList;
+    for (const ThreadingPath &Path : PathsToPhiDef) {
+      for (const PathType &PathToSw : PathsToSwitchBB) {
+        ThreadingPath PathCopy(Path);
+        PathCopy.appendExcludingFirst(PathToSw);
+        TempList.push_back(PathCopy);
       }
-
-      if (TPath.isExitValueSet() && isSupported(TPath))
-        TPaths.push_back(TPath);
     }
+    TPaths = std::move(TempList);
   }
 
 private:
   // Value: an instruction that defines a switch state;
   // Key: the parent basic block of that instruction.
   typedef DenseMap<const BasicBlock *, const PHINode *> StateDefMap;
+  std::vector<ThreadingPath> getPathsFromStateDefMap(StateDefMap &StateDef,
+                                                     PHINode *Phi,
+                                                     VisitedBlocks &VB) {
+    std::vector<ThreadingPath> Res;
+    auto *PhiBB = Phi->getParent();
+    VB.insert(PhiBB);
+
+    VisitedBlocks UniqueBlocks;
+    for (auto *IncomingBB : Phi->blocks()) {
+      if (!UniqueBlocks.insert(IncomingBB).second)
+        continue;
+      if (!SwitchOuterLoop->contains(IncomingBB))
+        continue;
+
+      Value *IncomingValue = Phi->getIncomingValueForBlock(IncomingBB);
+      // We found the determinator. The is the start of our path.
+      if (auto *C = dyn_cast<ConstantInt>(IncomingValue)) {
+        // SwitchBlock is the determinator, unsupported unless its also the def.
+        if (PhiBB == SwitchBlock &&
+            SwitchBlock != cast<PHINode>(Switch->getOperand(0))->getParent())
+          continue;
+        ThreadingPath NewPath;
+        NewPath.setDeterminator(PhiBB);
+        NewPath.setExitValue(C);
+        // Don't add SwitchBlock at the start, this is handled later.
+        if (IncomingBB != SwitchBlock)
+          NewPath.push_back(IncomingBB);
+        NewPath.push_back(PhiBB);
+        Res.push_back(NewPath);
+        continue;
+      }
+      // Don't get into a cycle.
+      if (VB.contains(IncomingBB) || IncomingBB == SwitchBlock)
+        continue;
+      // Recurse up the PHI chain.
+      auto *IncomingPhi = dyn_cast<PHINode>(IncomingValue);
+      if (!IncomingPhi)
+        continue;
+      auto *IncomingPhiDefBB = IncomingPhi->getParent();
+      if (!StateDef.contains(IncomingPhiDefBB))
+        continue;
+
+      // Direct prececessor, just add to the path.
+      if (IncomingPhiDefBB == IncomingBB) {
+        std::vector<ThreadingPath> PredPaths =
+            getPathsFromStateDefMap(StateDef, IncomingPhi, VB);
+        for (ThreadingPath &Path : PredPaths) {
+          Path.push_back(PhiBB);
+          Res.push_back(std::move(Path));
+        }
+        continue;
+      }
+      // Not a direct prececessor, find intermediate paths to append to the
+      // existing path.
----------------
XChy wrote:

Why do we handle `incomingPhiDefBB` separately for direct predecessor and non-direct predecessor? Is there any performance impact when handling it generally?

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