[ConstantRange] Support for ashr in ConstantRange computation

[jskumari]

Summary

Extend the ConstantRange implementation to compute the range of possible values resulting from an arithmetic right shift operation.
There will be a follow up patch to leverage this constant range infrastructure in LazyValueInfo.

[mkazantsev]

IR/ConstantRange.cpp

963 ↗

(On Diff #125670)

positive -> non-negative

[mkazantsev]

LGTM

[jskumari]

Incorporated review comments.

include/llvm/IR/ConstantRange.h

public:
   /// from a left shift of a value in this range by a value in \p Other.
   /// TODO: This isn't fully implemented yet.
   ConstantRange shl(const ConstantRange &Other) const;
 
   /// Return a new range representing the possible values resulting from a
   /// logical right shift of a value in this range and a value in \p Other.
   ConstantRange lshr(const ConstantRange &Other) const;
 
+  /// Return a new range representing the possible values resulting from a
+  /// arithmetic right shift of a value in this range and a value in \p Other.
+  ConstantRange ashr(const ConstantRange &Other) const;
+
   /// Return a new range that is the logical not of the current set.
   ConstantRange inverse() const;
 
   /// Print out the bounds to a stream.
   void print(raw_ostream &OS) const;
 
   /// Allow printing from a debugger easily.
   void dump() const;

lib/IR/ConstantRange.cpp

ConstantRange ConstantRange::binaryOp(Instruction::BinaryOps BinOp,
   case Instruction::Mul:
     return multiply(Other);
   case Instruction::UDiv:
     return udiv(Other);
   case Instruction::Shl:
     return shl(Other);
   case Instruction::LShr:
     return lshr(Other);
+  case Instruction::AShr:
+    return ashr(Other);
   case Instruction::And:
     return binaryAnd(Other);
   case Instruction::Or:
     return binaryOr(Other);
   // Note: floating point operations applied to abstract ranges are just
   // ideal integer operations with a lossy representation
   case Instruction::FAdd:
     return add(Other);
ConstantRange::lshr(const ConstantRange &Other) const {
   APInt max = getUnsignedMax().lshr(Other.getUnsignedMin()) + 1;
   APInt min = getUnsignedMin().lshr(Other.getUnsignedMax());
   if (min == max)
     return ConstantRange(getBitWidth(), /*isFullSet=*/true);
 
   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);
   if (isEmptySet())
     return ConstantRange(getBitWidth(), /*isFullSet=*/true);
   return ConstantRange(Upper, Lower);
 }
 

unittests/IR/ConstantRangeTest.cpp

TEST_F(ConstantRangeTest, Lshr) {
   EXPECT_EQ(One.lshr(Some), ConstantRange(APInt(16, 0)));
   EXPECT_EQ(One.lshr(Wrap), ConstantRange(APInt(16, 0), APInt(16, 0xb)));
   EXPECT_EQ(Some.lshr(Some), ConstantRange(APInt(16, 0),
                                            APInt(16, (0xaaa >> 0xa) + 1)));
   EXPECT_EQ(Some.lshr(Wrap), ConstantRange(APInt(16, 0), APInt(16, 0xaaa)));
   EXPECT_EQ(Wrap.lshr(Wrap), Full);
 }
 
+TEST_F(ConstantRangeTest, Ashr) {
+  EXPECT_EQ(Full.ashr(Full), Full);
+  EXPECT_EQ(Full.ashr(Empty), Empty);
+  EXPECT_EQ(Full.ashr(One), ConstantRange(APInt(16, 0xffe0),
+                                          APInt(16, (0x7fff >> 0xa) + 1 )));
+  ConstantRange Small(APInt(16, 0xa), APInt(16, 0xb));
+  EXPECT_EQ(Full.ashr(Small), ConstantRange(APInt(16, 0xffe0),
+                                           APInt(16, (0x7fff >> 0xa) + 1 )));
+  EXPECT_EQ(Full.ashr(Some), ConstantRange(APInt(16, 0xffe0),
+                                           APInt(16, (0x7fff >> 0xa) + 1 )));
+  EXPECT_EQ(Full.ashr(Wrap), Full);
+  EXPECT_EQ(Empty.ashr(Empty), Empty);
+  EXPECT_EQ(Empty.ashr(One), Empty);
+  EXPECT_EQ(Empty.ashr(Some), Empty);
+  EXPECT_EQ(Empty.ashr(Wrap), Empty);
+  EXPECT_EQ(One.ashr(One), ConstantRange(APInt(16, 0)));
+  EXPECT_EQ(One.ashr(Some), ConstantRange(APInt(16, 0)));
+  EXPECT_EQ(One.ashr(Wrap), ConstantRange(APInt(16, 0), APInt(16, 0xb)));
+  EXPECT_EQ(Some.ashr(Some), ConstantRange(APInt(16, 0),
+                                           APInt(16, (0xaaa >> 0xa) + 1)));
+  EXPECT_EQ(Some.ashr(Wrap), ConstantRange(APInt(16, 0), APInt(16, 0xaaa)));
+  EXPECT_EQ(Wrap.ashr(Wrap), Full);
+  ConstantRange Neg(APInt(16, 0xf3f0, true), APInt(16, 0xf7f8, true));
+  EXPECT_EQ(Neg.ashr(Small), ConstantRange(APInt(16, 0xfffc, true),
+                                           APInt(16, 0xfffe, true)));
+}
+
 TEST(ConstantRange, MakeAllowedICmpRegion) {
   // PR8250
   ConstantRange SMax = ConstantRange(APInt::getSignedMaxValue(32));
   EXPECT_TRUE(ConstantRange::makeAllowedICmpRegion(ICmpInst::ICMP_SGT, SMax)
                   .isEmptySet());
 }
 
 TEST(ConstantRange, MakeSatisfyingICmpRegion) {