[ConstantRange] Add `shlWithNoWrap()` method

[lebedev.ri]

Summary

This modelling isn't precise, much like the baseline shl() modelling,
so only conservative-correctness exaustive tests are performed.
Also, pending resolution of disscussion in D70043 whether or not
we should bother to ensure production of empty-set ranges for all-overflow inputs i've gone ahead with presumption that we will want that.

So all in all this is pretty similar to ConstantRange::mulWithNoWrap() modelling in D70043.

[lebedev.ri]

Simplify nsw all-wrap detection - we only need to check the .back() value
returned by getClosestToZero() - either there was a single value to begin with
and we are good either way, or there were two values and we are looking at the negative one.

[nikic]

Unlike the mul case, it should be possible to efficiently compute a precise shl range. This is the implementation I came up with for the nuw case:

if (NoWrapKind == OBO::NoUnsignedWrap) {
  APInt Min = getUnsignedMin(), Max = getUnsignedMax();
  unsigned MinLeadingZeros = Max.countLeadingZeros();
  unsigned MaxLeadingZeros = Min.countLeadingZeros();
  unsigned ShAmtMax = ShAmt.getUnsignedMax().getLimitedValue();
  unsigned ShAmtMin = ShAmt.getUnsignedMin().getLimitedValue();
  if (ShAmtMin > MaxLeadingZeros)
    // All possible shifts will overflow.
    return getEmpty();

  APInt ResMin = Min.shl(ShAmtMin);
  // We'll pick the larger of two possible upper bounds, initialize to zero.
  APInt ResMax = APInt::getNullValue(BitWidth);
  // Shift the maximum as far as possible without overflowing.
  if (ShAmtMin <= MinLeadingZeros)
    ResMax = Max.shl(std::min(ShAmtMax, MinLeadingZeros));
  // Pick the largest number with all low bits sets that is both in the LHS
  // range and can be shifted without overflowing.
  if (MinLeadingZeros != MaxLeadingZeros && ShAmtMax >= MinLeadingZeros + 1) {
    unsigned Shift = std::max(ShAmtMin, MinLeadingZeros + 1);
    ResMax = APIntOps::umax(ResMax,
        APInt::getLowBitsSet(BitWidth, BitWidth - Shift).shl(Shift));
  }
  return getNonEmpty(ResMin, ResMax + 1);
}

It passes the exhaustive test without the correctness only flag.

Extending this to the nuw+nsw case should be fairly simple, as it basically just has one less usable bit. I haven't thought about the nsw case yet.

[nikic]

Variant with nuw+nsw support:

if (NoWrapKind & OBO::NoUnsignedWrap) {
  APInt Min = getUnsignedMin(), Max = getUnsignedMax();
  unsigned MinLeadingZeros = Max.countLeadingZeros();
  unsigned MaxLeadingZeros = Min.countLeadingZeros();

  // If additionally the nsw flag is set, we can treat this the same way as
  // nuw, but with one less bit available.
  unsigned EffectiveBitWidth = BitWidth;
  if (NoWrapKind & OBO::NoSignedWrap) {
    --EffectiveBitWidth;
    if (MinLeadingZeros != 0) --MinLeadingZeros;
    if (MaxLeadingZeros != 0) --MaxLeadingZeros;
  }

  unsigned ShAmtMax = ShAmt.getUnsignedMax().getLimitedValue();
  unsigned ShAmtMin = ShAmt.getUnsignedMin().getLimitedValue();
  if (ShAmtMin > MaxLeadingZeros)
    // All possible shifts will overflow.
    return getEmpty();

  APInt ResMin = Min.shl(ShAmtMin);
  // We'll pick the larger of two possible upper bounds, initialize to zero.
  APInt ResMax = APInt::getNullValue(BitWidth);
  // Shift the maximum as far as possible without overflowing.
  if (ShAmtMin <= MinLeadingZeros)
    ResMax = Max.shl(std::min(ShAmtMax, MinLeadingZeros));
  // Pick the largest number with all low bits sets that is both in the LHS
  // range and can be shifted without overflowing.
  if (MinLeadingZeros != MaxLeadingZeros && ShAmtMax >= MinLeadingZeros + 1) {
    unsigned Shift = std::max(ShAmtMin, MinLeadingZeros + 1);
    ResMax = APIntOps::umax(ResMax,
        APInt::getLowBitsSet(BitWidth, EffectiveBitWidth - Shift).shl(Shift));
  }
  return getNonEmpty(ResMin, ResMax + 1);
}

The nsw case is ... uh ... left as an exercise for the reader.

llvm/include/llvm/IR/ConstantRange.h

public:
   /// from a binary-or of a value in this range by a value in \p Other.
   ConstantRange binaryOr(const ConstantRange &Other) const;
 
   /// Return a new range representing the possible values resulting
   /// 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 left shift with wrap type \p NoWrapKind of a value in this
+  /// range and a value in \p Other.
+  /// If the result range is disjoint, the preferred range is determined by the
+  /// \p PreferredRangeType.
+  ConstantRange shlWithNoWrap(const ConstantRange &Other, unsigned NoWrapKind,
+                              PreferredRangeType RangeType = Smallest) 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;
 

llvm/lib/IR/ConstantRange.cpp

ConstantRange::shl(const ConstantRange &Other) const {
 
   APInt min = getUnsignedMin();
   min <<= Other.getUnsignedMin();
   max <<= Other_umax;
 
   return ConstantRange(std::move(min), std::move(max) + 1);
 }
 
+ConstantRange ConstantRange::shlWithNoWrap(const ConstantRange &Other,
+                                           unsigned NoWrapKind,
+                                           PreferredRangeType RangeType) const {
+  unsigned BitWidth = Other.getBitWidth();
+
+  // For given range of shift amounts, if we ignore all illegal shift amounts
+  // (that always produce poison), what shift amount range is left?
+  ConstantRange ShAmt = Other.intersectWith(
+      ConstantRange(APInt(BitWidth, 0), APInt(BitWidth, (BitWidth - 1) + 1)));
+
+  // Calculate the range for "X << Y" which is guaranteed not to wrap(overflow).
+  // (X is from this, and Y is from ShAmt)
+  if (isEmptySet() || ShAmt.isEmptySet())
+    return getEmpty();
+  if (isFullSet() && ShAmt.isFullSet())
+    return getFull();
+
+  using OBO = OverflowingBinaryOperator;
+  ConstantRange Result = shl(ShAmt);
+
+  if (NoWrapKind & OBO::NoUnsignedWrap) {
+    bool Overflow;
+
+    (void)getUnsignedMin().ushl_ov(ShAmt.getUnsignedMin(), Overflow);
+    if (Overflow)
+      return getEmpty(); // Always overflows.
+
+    if (NoWrapKind & OBO::NoSignedWrap) {
+      (void)getUnsignedMin().sshl_ov(ShAmt.getUnsignedMin(), Overflow);
+      if (Overflow)
+        return getEmpty(); // Always overflows.
+    }
+
+    Result = Result.intersectWith(ushl_sat(ShAmt), RangeType);
+  }
+
+  if (NoWrapKind & OBO::NoSignedWrap) {
+    // Since we know we have a non-empty set, there may be either one or two
+    // values closest to signed zero. If there is only a single such value,
+    // then there is no question which one we should use. However if there are
+    // two values, they have different polarities, and we are *only* interested
+    // in the one that's negative, which is returned second.
+    APInt ClosestToZero = getClosestToZero().back();
+    bool Overflow;
+    (void)ClosestToZero.sshl_ov(ShAmt.getUnsignedMin(), Overflow);
+    if (Overflow)
+      return getEmpty(); // Always overflows.
+
+    Result = Result.intersectWith(sshl_sat(ShAmt), RangeType);
+  }
+
+  return Result;
+}
+
 ConstantRange
 ConstantRange::lshr(const ConstantRange &Other) const {
   if (isEmptySet() || Other.isEmptySet())
     return getEmpty();
 
   APInt max = getUnsignedMax().lshr(Other.getUnsignedMin()) + 1;
   APInt min = getUnsignedMin().lshr(Other.getUnsignedMax());
   return getNonEmpty(std::move(min), std::move(max));

llvm/unittests/IR/ConstantRangeTest.cpp

TEST_F(ConstantRangeTest, Shl) {
   EXPECT_EQ(Some.shl(Wrap), Full);   // TODO: [0xa, 0x7ff << 0x5 + 1)
   EXPECT_EQ(Wrap.shl(Wrap), Full);
   EXPECT_EQ(
       Some2.shl(ConstantRange(APInt(16, 0x1))),
       ConstantRange(APInt(16, 0xfff << 0x1), APInt(16, 0x7fff << 0x1) + 1));
   EXPECT_EQ(One.shl(WrapNullMax), Full);
 }
 
+TEST_F(ConstantRangeTest, ShlWithNoWrap) {
+  typedef OverflowingBinaryOperator OBO;
+
+  TestAddWithNoSignedWrapExhaustive(
+      [](const ConstantRange &CR1, const ConstantRange &CR2) {
+        return CR1.shlWithNoWrap(CR2, OBO::NoSignedWrap);
+      },
+      [](bool &IsOverflow, const APInt &N1, const APInt &N2) {
+        return N1.sshl_ov(N2, IsOverflow);
+      },
+      /*CorrectnessOnly=*/true);
+
+  TestAddWithNoUnsignedWrapExhaustive(
+      [](const ConstantRange &CR1, const ConstantRange &CR2) {
+        return CR1.shlWithNoWrap(CR2, OBO::NoUnsignedWrap);
+      },
+      [](bool &IsOverflow, const APInt &N1, const APInt &N2) {
+        return N1.ushl_ov(N2, IsOverflow);
+      },
+      /*CorrectnessOnly=*/true);
+
+  TestAddWithNoSignedUnsignedWrapExhaustive(
+      [](const ConstantRange &CR1, const ConstantRange &CR2) {
+        return CR1.shlWithNoWrap(CR2, OBO::NoUnsignedWrap | OBO::NoSignedWrap);
+      },
+      [](bool &IsOverflow, const APInt &N1, const APInt &N2) {
+        return N1.sshl_ov(N2, IsOverflow);
+      },
+      [](bool &IsOverflow, const APInt &N1, const APInt &N2) {
+        return N1.ushl_ov(N2, IsOverflow);
+      },
+      /*CorrectnessOnly=*/true);
+}
+
 TEST_F(ConstantRangeTest, Lshr) {
   EXPECT_EQ(Full.lshr(Full), Full);
   EXPECT_EQ(Full.lshr(Empty), Empty);
   EXPECT_EQ(Full.lshr(One), ConstantRange(APInt(16, 0),
                                           APInt(16, (0xffff >> 0xa) + 1)));
   EXPECT_EQ(Full.lshr(Some), ConstantRange(APInt(16, 0),
                                            APInt(16, (0xffff >> 0xa) + 1)));
   EXPECT_EQ(Full.lshr(Wrap), Full);