[CVP] Keep ReachableCaseCount in sync with range of condition (#142302)

https://github.com/llvm/llvm-project/pull/79993 assumes that a reachable
case must be contained by `CR`. However, it doesn't hold for some edge
cases. This patch adds additional checks to ensure `ReachableCaseCount`
is correct.

Note: Similar optimization in SCCP isn't affected by this bug because it
uses `CR` to compute `ReachableCaseCount`.

Closes https://github.com/llvm/llvm-project/issues/142286.

llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp

@@ -370,15 +370,30 @@ static bool processSwitch(SwitchInst *I, LazyValueInfo *LVI,
   { // Scope for SwitchInstProfUpdateWrapper. It must not live during
     // ConstantFoldTerminator() as the underlying SwitchInst can be changed.
     SwitchInstProfUpdateWrapper SI(*I);
+    ConstantRange CR =
+        LVI->getConstantRangeAtUse(I->getOperandUse(0), /*UndefAllowed=*/false);
     unsigned ReachableCaseCount = 0;
 
     for (auto CI = SI->case_begin(), CE = SI->case_end(); CI != CE;) {
       ConstantInt *Case = CI->getCaseValue();
+      std::optional<bool> Predicate = std::nullopt;
+      if (!CR.contains(Case->getValue()))
+        Predicate = false;
+      else if (CR.isSingleElement() &&
+               *CR.getSingleElement() == Case->getValue())
+        Predicate = true;
+      if (!Predicate) {
+        // Handle missing cases, e.g., the range has a hole.
         auto *Res = dyn_cast_or_null<ConstantInt>(
             LVI->getPredicateAt(CmpInst::ICMP_EQ, Cond, Case, I,
-                              /* UseBlockValue */ true));
+                                /* UseBlockValue=*/true));
+        if (Res && Res->isZero())
+          Predicate = false;
+        else if (Res && Res->isOne())
+          Predicate = true;
+      }
 
-      if (Res && Res->isZero()) {
+      if (Predicate && !*Predicate) {
         // This case never fires - remove it.
         BasicBlock *Succ = CI->getCaseSuccessor();
         Succ->removePredecessor(BB);
@@ -395,7 +410,7 @@ static bool processSwitch(SwitchInst *I, LazyValueInfo *LVI,
           DTU.applyUpdatesPermissive({{DominatorTree::Delete, BB, Succ}});
         continue;
       }
-      if (Res && Res->isOne()) {
+      if (Predicate && *Predicate) {
         // This case always fires.  Arrange for the switch to be turned into an
         // unconditional branch by replacing the switch condition with the case
         // value.
@@ -410,13 +425,10 @@ static bool processSwitch(SwitchInst *I, LazyValueInfo *LVI,
       ++ReachableCaseCount;
     }
 
-    BasicBlock *DefaultDest = SI->getDefaultDest();
-    if (ReachableCaseCount > 1 &&
-        !isa<UnreachableInst>(DefaultDest->getFirstNonPHIOrDbg())) {
-      ConstantRange CR = LVI->getConstantRangeAtUse(I->getOperandUse(0),
-                                                    /*UndefAllowed*/ false);
     // The default dest is unreachable if all cases are covered.
-      if (!CR.isSizeLargerThan(ReachableCaseCount)) {
+    if (!SI->defaultDestUndefined() &&
+        !CR.isSizeLargerThan(ReachableCaseCount)) {
+      BasicBlock *DefaultDest = SI->getDefaultDest();
       BasicBlock *NewUnreachableBB =
           BasicBlock::Create(BB->getContext(), "default.unreachable",
                              BB->getParent(), DefaultDest);
@@ -433,7 +445,6 @@ static bool processSwitch(SwitchInst *I, LazyValueInfo *LVI,
       Changed = true;
     }
   }
-  }
 
   if (Changed)
     // If the switch has been simplified to the point where it can be replaced

llvm/test/Transforms/CorrelatedValuePropagation/switch.ll

@@ -294,6 +294,42 @@ cleanup:
   ret i32 %retval.0
 }
 
+; Make sure that we don't branch into unreachable.
+
+define void @pr142286() {
+; CHECK-LABEL: define void @pr142286() {
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    br label [[LOOP2:%.*]]
+; CHECK:       loop2:
+; CHECK-NEXT:    br label [[LOOP3:%.*]]
+; CHECK:       loop3:
+; CHECK-NEXT:    br label [[EXIT:%.*]]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+start:
+  br label %loop
+
+loop:
+  %phi = phi i8 [ -1, %start ], [ 0, %loop3 ]
+  br label %loop2
+
+loop2:
+  br label %loop3
+
+loop3:
+  switch i8 %phi, label %exit [
+  i8 0, label %loop3
+  i8 1, label %loop2
+  i8 2, label %loop
+  ]
+
+exit:
+  ret void
+}
+
 declare i32 @call0()
 declare i32 @call1()
 declare i32 @call2()