Fix the count of pruning incoming values

Author: dianqk

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -5836,6 +5836,8 @@ bool SimplifyCFGOpt::turnSwitchRangeIntoICmp(SwitchInst *SI,
   BasicBlock *ContiguousDest = nullptr;
   BasicBlock *OtherDest = nullptr;
   bool IsWrapping = false;
+  SmallVectorImpl<ConstantInt *> *ContiguousCases = &CasesA;
+  SmallVectorImpl<ConstantInt *> *OtherCases = &CasesB;
   if (!CasesA.empty() &&
       casesAreContiguous(SI->getCondition(), CasesA, ContiguousCasesMin,
                          ContiguousCasesMax, IsWrapping)) {
@@ -5845,28 +5847,29 @@ bool SimplifyCFGOpt::turnSwitchRangeIntoICmp(SwitchInst *SI,
                                 ContiguousCasesMax, IsWrapping)) {
     ContiguousDest = DestB;
     OtherDest = DestA;
+    std::swap(ContiguousCases, OtherCases);
   } else
     return false;
 
-  if (IsWrapping)
+  if (IsWrapping) {
     std::swap(ContiguousDest, OtherDest);
+    std::swap(ContiguousCases, OtherCases);
+  }
 
   // Start building the compare and branch.
 
-  auto ContiguousCasesSize =
-      (ContiguousCasesMax->getValue() - ContiguousCasesMin->getValue())
-          .getZExtValue() +
-      1;
   Constant *Offset = ConstantExpr::getNeg(ContiguousCasesMin);
-  Constant *NumCases = ConstantInt::get(Offset->getType(), ContiguousCasesSize);
+  Constant *NumCases = ConstantInt::get(Offset->getType(),
+                                        ContiguousCasesMax->getValue() -
+                                            ContiguousCasesMin->getValue() + 1);
 
   Value *Sub = SI->getCondition();
   if (!Offset->isNullValue())
     Sub = Builder.CreateAdd(Sub, Offset, Sub->getName() + ".off");
 
   Value *Cmp;
   // If NumCases overflowed, then all possible values jump to the successor.
-  if (NumCases->isNullValue() && ContiguousCasesSize != 0)
+  if (NumCases->isNullValue() && !ContiguousCases->empty())
     Cmp = ConstantInt::getTrue(SI->getContext());
   else
     Cmp = Builder.CreateICmpULT(Sub, NumCases, "switch");
@@ -5896,14 +5899,14 @@ bool SimplifyCFGOpt::turnSwitchRangeIntoICmp(SwitchInst *SI,
 
   // Prune obsolete incoming values off the successors' PHI nodes.
   for (auto BBI = ContiguousDest->begin(); isa<PHINode>(BBI); ++BBI) {
-    unsigned PreviousEdges = ContiguousCasesSize;
+    unsigned PreviousEdges = ContiguousCases->size();
     if (ContiguousDest == SI->getDefaultDest())
       ++PreviousEdges;
     for (unsigned I = 0, E = PreviousEdges - 1; I != E; ++I)
       cast<PHINode>(BBI)->removeIncomingValue(SI->getParent());
   }
   for (auto BBI = OtherDest->begin(); isa<PHINode>(BBI); ++BBI) {
-    unsigned PreviousEdges = SI->getNumCases() - ContiguousCasesSize;
+    unsigned PreviousEdges = OtherCases->size();
     if (OtherDest == SI->getDefaultDest())
       ++PreviousEdges;
     for (unsigned I = 0, E = PreviousEdges - 1; I != E; ++I)

llvm/test/Transforms/SimplifyCFG/switch-range-to-icmp.ll

@@ -249,6 +249,28 @@ else:
   ret i32 %i1
 }
 
+define i8 @wrapping_range_phi(i8 %arg) {
+; CHECK-LABEL: @wrapping_range_phi(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ARG_OFF:%.*]] = add i8 [[ARG:%.*]], -1
+; CHECK-NEXT:    [[SWITCH:%.*]] = icmp ult i8 [[ARG_OFF]], -2
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[SWITCH]], i8 0, i8 1
+; CHECK-NEXT:    ret i8 [[SPEC_SELECT]]
+;
+entry:
+  switch i8 %arg, label %else [
+  i8 0, label %if
+  i8 -1, label %if
+  ]
+
+if:
+  %i = phi i8 [ 0, %else ], [ 1, %entry ], [ 1, %entry ]
+  ret i8 %i
+
+else:
+  br label %if
+}
+
 define i32 @no_continuous_wrapping_range(i8 %arg) {
 ; CHECK-LABEL: @no_continuous_wrapping_range(
 ; CHECK-NEXT:    switch i8 [[ARG:%.*]], label [[ELSE:%.*]] [