1 files changed, 108 insertions, 28 deletions
diff --git a/lib/IR/ConstantRange.cpp b/lib/IR/ConstantRange.cpp
index 4bd17257016d..48d16f334ba3 100644
--- a/lib/IR/ConstantRange.cpp
+++ b/lib/IR/ConstantRange.cpp
@@ -199,39 +199,63 @@ ConstantRange::makeGuaranteedNoWrapRegion(Instruction::BinaryOps BinOp,
          "NoWrapKind invalid!");
 
   unsigned BitWidth = Other.getBitWidth();
-  if (BinOp != Instruction::Add)
+  ConstantRange Result(BitWidth);
+
+  switch (BinOp) {
+  default:
     // Conservative answer: empty set
     return ConstantRange(BitWidth, false);
 
-  if (auto *C = Other.getSingleElement())
-    if (C->isNullValue())
-      // Full set: nothing signed / unsigned wraps when added to 0.
-      return ConstantRange(BitWidth);
-
-  ConstantRange Result(BitWidth);
+  case Instruction::Add:
+    if (auto *C = Other.getSingleElement())
+      if (C->isNullValue())
+        // Full set: nothing signed / unsigned wraps when added to 0.
+        return ConstantRange(BitWidth);
+    if (NoWrapKind & OBO::NoUnsignedWrap)
+      Result =
+          SubsetIntersect(Result, ConstantRange(APInt::getNullValue(BitWidth),
+                                                -Other.getUnsignedMax()));
+    if (NoWrapKind & OBO::NoSignedWrap) {
+      const APInt &SignedMin = Other.getSignedMin();
+      const APInt &SignedMax = Other.getSignedMax();
+      if (SignedMax.isStrictlyPositive())
+        Result = SubsetIntersect(
+            Result,
+            ConstantRange(APInt::getSignedMinValue(BitWidth),
+                          APInt::getSignedMinValue(BitWidth) - SignedMax));
+      if (SignedMin.isNegative())
+        Result = SubsetIntersect(
+            Result,
+            ConstantRange(APInt::getSignedMinValue(BitWidth) - SignedMin,
+                          APInt::getSignedMinValue(BitWidth)));
+    }
+    return Result;
 
-  if (NoWrapKind & OBO::NoUnsignedWrap)
-    Result =
-        SubsetIntersect(Result, ConstantRange(APInt::getNullValue(BitWidth),
-                                              -Other.getUnsignedMax()));
-
-  if (NoWrapKind & OBO::NoSignedWrap) {
-    const APInt &SignedMin = Other.getSignedMin();
-    const APInt &SignedMax = Other.getSignedMax();
-
-    if (SignedMax.isStrictlyPositive())
-      Result = SubsetIntersect(
-          Result,
-          ConstantRange(APInt::getSignedMinValue(BitWidth),
-                        APInt::getSignedMinValue(BitWidth) - SignedMax));
-
-    if (SignedMin.isNegative())
-      Result = SubsetIntersect(
-          Result, ConstantRange(APInt::getSignedMinValue(BitWidth) - SignedMin,
-                                APInt::getSignedMinValue(BitWidth)));
+  case Instruction::Sub:
+    if (auto *C = Other.getSingleElement())
+      if (C->isNullValue())
+        // Full set: nothing signed / unsigned wraps when subtracting 0.
+        return ConstantRange(BitWidth);
+    if (NoWrapKind & OBO::NoUnsignedWrap)
+      Result =
+          SubsetIntersect(Result, ConstantRange(Other.getUnsignedMax(),
+                                                APInt::getMinValue(BitWidth)));
+    if (NoWrapKind & OBO::NoSignedWrap) {
+      const APInt &SignedMin = Other.getSignedMin();
+      const APInt &SignedMax = Other.getSignedMax();
+      if (SignedMax.isStrictlyPositive())
+        Result = SubsetIntersect(
+            Result,
+            ConstantRange(APInt::getSignedMinValue(BitWidth) + SignedMax,
+                          APInt::getSignedMinValue(BitWidth)));
+      if (SignedMin.isNegative())
+        Result = SubsetIntersect(
+            Result,
+            ConstantRange(APInt::getSignedMinValue(BitWidth),
+                          APInt::getSignedMinValue(BitWidth) + SignedMin));
+    }
+    return Result;
   }
-
-  return Result;
 }
 
 bool ConstantRange::isFullSet() const {
@@ -656,6 +680,8 @@ ConstantRange ConstantRange::binaryOp(Instruction::BinaryOps BinOp,
     return shl(Other);
   case Instruction::LShr:
     return lshr(Other);
+  case Instruction::AShr:
+    return ashr(Other);
   case Instruction::And:
     return binaryAnd(Other);
   case Instruction::Or:
@@ -922,6 +948,60 @@ ConstantRange::lshr(const ConstantRange &Other) const {
   return ConstantRange(std::move(min), std::move(max));
 }
 
+ConstantRange
+ConstantRange::ashr(const ConstantRange &Other) const {
+  if (isEmptySet() || Other.isEmptySet())
+    return ConstantRange(getBitWidth(), /*isFullSet=*/false);
+
+  // May straddle zero, so handle both positive and negative cases.
+  // 'PosMax' is the upper bound of the result of the ashr
+  // operation, when Upper of the LHS of ashr is a non-negative.
+  // number. Since ashr of a non-negative number will result in a
+  // smaller number, the Upper value of LHS is shifted right with
+  // the minimum value of 'Other' instead of the maximum value.
+  APInt PosMax = getSignedMax().ashr(Other.getUnsignedMin()) + 1;
+
+  // 'PosMin' is the lower bound of the result of the ashr
+  // operation, when Lower of the LHS is a non-negative number.
+  // Since ashr of a non-negative number will result in a smaller
+  // number, the Lower value of LHS is shifted right with the
+  // maximum value of 'Other'.
+  APInt PosMin = getSignedMin().ashr(Other.getUnsignedMax());
+
+  // 'NegMax' is the upper bound of the result of the ashr
+  // operation, when Upper of the LHS of ashr is a negative number.
+  // Since 'ashr' of a negative number will result in a bigger
+  // number, the Upper value of LHS is shifted right with the
+  // maximum value of 'Other'.
+  APInt NegMax = getSignedMax().ashr(Other.getUnsignedMax()) + 1;
+
+  // 'NegMin' is the lower bound of the result of the ashr
+  // operation, when Lower of the LHS of ashr is a negative number.
+  // Since 'ashr' of a negative number will result in a bigger
+  // number, the Lower value of LHS is shifted right with the
+  // minimum value of 'Other'.
+  APInt NegMin = getSignedMin().ashr(Other.getUnsignedMin());
+
+  APInt max, min;
+  if (getSignedMin().isNonNegative()) {
+    // Upper and Lower of LHS are non-negative.
+    min = PosMin;
+    max = PosMax;
+  } else if (getSignedMax().isNegative()) {
+    // Upper and Lower of LHS are negative.
+    min = NegMin;
+    max = NegMax;
+  } else {
+    // Upper is non-negative and Lower is negative.
+    min = NegMin;
+    max = PosMax;
+  }
+  if (min == max)
+    return ConstantRange(getBitWidth(), /*isFullSet=*/true);
+
+  return ConstantRange(std::move(min), std::move(max));
+}
+
 ConstantRange ConstantRange::inverse() const {
   if (isFullSet())
     return ConstantRange(getBitWidth(), /*isFullSet=*/false);