[InstCombine] Fix const folding of switched with default case

[danilaml]

Summary

In case phi was in the default block it could lead to multi-edge.
Fixes #55721.

[nikic]

llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp

1308

I don't think this is sufficient. Can you please add a test where the default doesn't directly go to the phi block, but have one dummy block in between? (That is, go to a sw.19 block from both 19 and default.)

I think the right fix would be adding a ++SuccCount[SI->getDefaultDest()].

[danilaml]

llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp

1308

I'll add this test case, but I'm pretty sure that this is covered by the dominance check

DT.dominates(BasicBlockEdge(IDom, It->second),
                          BasicBlockEdge(Pred, BB));

I think this check will only succeed when both edges are the "same".

Then again, I think your proposed fix would work as well (at the expense of slightly more work), so I can switch to it.

[danilaml]

Addressed the comments

[nikic]

LGTM

llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp

 //===- InstCombinePHI.cpp -------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 static Value *simplifyUsingControlFlow(InstCombiner &Self, PHINode &PN,
                                        const DominatorTree &DT) {
   // Simplify the following patterns:
   //       if (cond)
   //       /       \
   //      ...      ...
   //       \       /
   //    phi [true] [false]
+  // and
+  //        switch (cond)
+  // case v1: /       \ case v2:
+  //         ...      ...
+  //          \       /
+  //       phi [v1] [v2]
   // Make sure all inputs are constants.
   if (!all_of(PN.operands(), [](Value *V) { return isa<ConstantInt>(V); }))
     return nullptr;
 
   BasicBlock *BB = PN.getParent();
   // Do not bother with unreachable instructions.
   if (!DT.isReachableFromEntry(BB))
     return nullptr;
static Value *simplifyUsingControlFlow(InstCombiner &Self, PHINode &PN,
   if (auto *BI = dyn_cast<BranchInst>(IDom->getTerminator())) {
     if (BI->isUnconditional())
       return nullptr;
 
     Cond = BI->getCondition();
     AddSucc(ConstantInt::getTrue(Context), BI->getSuccessor(0));
     AddSucc(ConstantInt::getFalse(Context), BI->getSuccessor(1));
   } else if (auto *SI = dyn_cast<SwitchInst>(IDom->getTerminator())) {
+    // Can't generally replace Phi with Cond if we can get there from default
+    // case.
+    if (SI->getDefaultDest() == PN.getParent())
+      return nullptr;
     Cond = SI->getCondition();
     for (auto Case : SI->cases())
       AddSucc(Case.getCaseValue(), Case.getCaseSuccessor());
   } else {
     return nullptr;
   }
 
   if (Cond->getType() != PN.getType())

llvm/test/Transforms/InstCombine/simple_phi_condition.ll

 
 sw.19:
   br label %merge
 
 merge:
   %ret = phi i8 [ 1, %sw.1 ], [ 7, %sw.7 ], [ 19, %sw.19 ], [ 42, %entry ]
   ret i8 %ret
 }
+
+define i8 @test_switch_default_edge_direct(i8 %cond) {
+; CHECK-LABEL: @test_switch_default_edge_direct(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    switch i8 [[COND:%.*]], label [[MERGE:%.*]] [
+; CHECK-NEXT:    i8 1, label [[SW_1:%.*]]
+; CHECK-NEXT:    i8 7, label [[SW_7:%.*]]
+; CHECK-NEXT:    i8 19, label [[MERGE]]
+; CHECK-NEXT:    ]
+; CHECK:       sw.1:
+; CHECK-NEXT:    br label [[MERGE]]
+; CHECK:       sw.7:
+; CHECK-NEXT:    br label [[MERGE]]
+; CHECK:       merge:
+; CHECK-NEXT:    [[RET:%.*]] = phi i8 [ 1, [[SW_1]] ], [ 7, [[SW_7]] ], [ 19, [[ENTRY:%.*]] ], [ 19, [[ENTRY]] ]
+; CHECK-NEXT:    ret i8 [[RET]]
+;
+entry:
+  switch i8 %cond, label %merge [
+  i8 1, label %sw.1
+  i8 7, label %sw.7
+  i8 19, label %merge
+  ]
+sw.1:
+  br label %merge
+sw.7:
+  br label %merge
+merge:
+  %ret = phi i8 [ 1, %sw.1 ], [ 7, %sw.7 ], [ 19, %entry ], [ 19, %entry ]
+  ret i8 %ret
+}