[SimplifyCFG] Fold the contiguous wrapping cases into ICmp.

Author: dianqk

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -5718,15 +5718,49 @@ bool SimplifyCFGOpt::simplifyUnreachable(UnreachableInst *UI) {
   return Changed;
 }
 
-static bool casesAreContiguous(SmallVectorImpl<ConstantInt *> &Cases) {
+static bool casesAreContiguous(Value *Condition,
+                               SmallVectorImpl<ConstantInt *> &Cases,
+                               ConstantInt *&ContiguousCasesMin,
+                               ConstantInt *&ContiguousCasesMax,
+                               bool &IsWrapping) {
   assert(Cases.size() >= 1);
 
   array_pod_sort(Cases.begin(), Cases.end(), constantIntSortPredicate);
-  for (size_t I = 1, E = Cases.size(); I != E; ++I) {
-    if (Cases[I - 1]->getValue() != Cases[I]->getValue() + 1)
+  auto Min = Cases.back()->getValue();
+  auto Max = Cases.front()->getValue();
+  auto Offset = Max - Min;
+  auto ContiguousOffset = Cases.size() - 1;
+  if (Offset == ContiguousOffset) {
+    ContiguousCasesMin = Cases.back();
+    ContiguousCasesMax = Cases.front();
+    return true;
+  }
+  ConstantRange CR = computeConstantRange(Condition, /*ForSigned=*/false);
+  // If this is a wrapping contiguous range, that is, [Min, OtherMin] +
+  // [OtherMax, Max] (also [OtherMax, OtherMin]), [OtherMin+1, OtherMax-1] is a
+  // contiguous range for the other destination. N.B. If CR is not a full range,
+  // Max+1 is not equal to Min. It's not continuous in arithmetic.
+  if (Max == CR.getUnsignedMax() && Min == CR.getUnsignedMin()) {
+    assert(Cases.size() >= 2);
+    auto *It =
+        std::adjacent_find(Cases.begin(), Cases.end(), [](auto L, auto R) {
+          return L->getValue() != R->getValue() + 1;
+        });
+    if (It == Cases.end())
       return false;
+    auto *OtherMax = *It;
+    auto *OtherMin = *(It + 1);
+    if ((Max - OtherMax->getValue()) + (OtherMin->getValue() - Min) ==
+        Cases.size() - 2) {
+      ContiguousCasesMin = cast<ConstantInt>(
+          ConstantInt::get(OtherMin->getType(), OtherMin->getValue() + 1));
+      ContiguousCasesMax = cast<ConstantInt>(
+          ConstantInt::get(OtherMax->getType(), OtherMax->getValue() - 1));
+      IsWrapping = true;
+      return true;
+    }
   }
-  return true;
+  return false;
 }
 
 static void createUnreachableSwitchDefault(SwitchInst *Switch,
@@ -5797,33 +5831,42 @@ bool SimplifyCFGOpt::turnSwitchRangeIntoICmp(SwitchInst *SI,
   assert(!CasesA.empty() || HasDefault);
 
   // Figure out if one of the sets of cases form a contiguous range.
-  SmallVectorImpl<ConstantInt *> *ContiguousCases = nullptr;
+  ConstantInt *ContiguousCasesMin = nullptr;
+  ConstantInt *ContiguousCasesMax = nullptr;
   BasicBlock *ContiguousDest = nullptr;
   BasicBlock *OtherDest = nullptr;
-  if (!CasesA.empty() && casesAreContiguous(CasesA)) {
-    ContiguousCases = &CasesA;
+  bool IsWrapping = false;
+  if (!CasesA.empty() &&
+      casesAreContiguous(SI->getCondition(), CasesA, ContiguousCasesMin,
+                         ContiguousCasesMax, IsWrapping)) {
     ContiguousDest = DestA;
     OtherDest = DestB;
-  } else if (casesAreContiguous(CasesB)) {
-    ContiguousCases = &CasesB;
+  } else if (casesAreContiguous(SI->getCondition(), CasesB, ContiguousCasesMin,
+                                ContiguousCasesMax, IsWrapping)) {
     ContiguousDest = DestB;
     OtherDest = DestA;
   } else
     return false;
 
+  if (IsWrapping)
+    std::swap(ContiguousDest, OtherDest);
+
   // Start building the compare and branch.
 
-  Constant *Offset = ConstantExpr::getNeg(ContiguousCases->back());
-  Constant *NumCases =
-      ConstantInt::get(Offset->getType(), ContiguousCases->size());
+  auto ContiguousCasesSize =
+      (ContiguousCasesMax->getValue() - ContiguousCasesMin->getValue())
+          .getZExtValue() +
+      1;
+  Constant *Offset = ConstantExpr::getNeg(ContiguousCasesMin);
+  Constant *NumCases = ConstantInt::get(Offset->getType(), ContiguousCasesSize);
 
   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() && !ContiguousCases->empty())
+  if (NumCases->isNullValue() && ContiguousCasesSize != 0)
     Cmp = ConstantInt::getTrue(SI->getContext());
   else
     Cmp = Builder.CreateICmpULT(Sub, NumCases, "switch");
@@ -5853,14 +5896,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 = ContiguousCases->size();
+    unsigned PreviousEdges = ContiguousCasesSize;
     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() - ContiguousCases->size();
+    unsigned PreviousEdges = SI->getNumCases() - ContiguousCasesSize;
     if (OtherDest == SI->getDefaultDest())
       ++PreviousEdges;
     for (unsigned I = 0, E = PreviousEdges - 1; I != E; ++I)

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

@@ -190,11 +190,9 @@ exit:
 
 define i32 @wrapping_known_range(i8 range(i8 0, 6) %arg) {
 ; CHECK-LABEL: @wrapping_known_range(
-; CHECK-NEXT:    switch i8 [[ARG:%.*]], label [[ELSE:%.*]] [
-; CHECK-NEXT:      i8 0, label [[IF:%.*]]
-; CHECK-NEXT:      i8 4, label [[IF]]
-; CHECK-NEXT:      i8 5, label [[IF]]
-; CHECK-NEXT:    ]
+; CHECK-NEXT:    [[ARG_OFF:%.*]] = add i8 [[ARG:%.*]], -1
+; CHECK-NEXT:    [[SWITCH:%.*]] = icmp ult i8 [[ARG_OFF]], 3
+; CHECK-NEXT:    br i1 [[SWITCH]], label [[ELSE:%.*]], label [[IF:%.*]]
 ; CHECK:       common.ret:
 ; CHECK-NEXT:    [[COMMON_RET_OP:%.*]] = phi i32 [ [[I0:%.*]], [[IF]] ], [ [[I1:%.*]], [[ELSE]] ]
 ; CHECK-NEXT:    ret i32 [[COMMON_RET_OP]]
@@ -222,12 +220,9 @@ else:
 
 define i32 @wrapping_range(i8 %arg) {
 ; CHECK-LABEL: @wrapping_range(
-; CHECK-NEXT:    switch i8 [[ARG:%.*]], label [[ELSE:%.*]] [
-; CHECK-NEXT:      i8 0, label [[IF:%.*]]
-; CHECK-NEXT:      i8 -3, label [[IF]]
-; CHECK-NEXT:      i8 -2, label [[IF]]
-; CHECK-NEXT:      i8 -1, label [[IF]]
-; CHECK-NEXT:    ]
+; CHECK-NEXT:    [[ARG_OFF:%.*]] = add i8 [[ARG:%.*]], -1
+; CHECK-NEXT:    [[SWITCH:%.*]] = icmp ult i8 [[ARG_OFF]], -4
+; CHECK-NEXT:    br i1 [[SWITCH]], label [[ELSE:%.*]], label [[IF:%.*]]
 ; CHECK:       common.ret:
 ; CHECK-NEXT:    [[COMMON_RET_OP:%.*]] = phi i32 [ [[I0:%.*]], [[IF]] ], [ [[I1:%.*]], [[ELSE]] ]
 ; CHECK-NEXT:    ret i32 [[COMMON_RET_OP]]