[ConstantRange] Add PreferredResultType support for unionWith()

[nikic]

Summary

This extends D59959 to unionWith(), allowing to specify that a non-wrapping unsigned/signed range is preferred. This is somewhat less useful than the intersect case, because union operations are rarer. An example use would the the phi union computed in SCEV.

The implementation is mostly a straightforward use of getPreferredRange(), but I also had to adjust some <=/< checks to make sure that no ranges with lower==upper get constructed before they're passed to getPreferredRange(), as these have additional constraints.

[lebedev.ri]

llvm/lib/IR/ConstantRange.cpp

470

Please do this as preparatory NFC commit.

608–610

And here we have

    //  L---U            : this
    //        L---U      : CR
or
    //        L---U      : this
    //  L---U            : CR

=>

    //  L---------U      : this
or
    //        L--------- : this
    // -----U            : CR

?

634–636

So we will produce

    // ----U       L---- : this
    //       L---U       : CR
=>
    // ----------U L----
or
    // ----U L----------

and in either case we bridge either of the gaps?

llvm/unittests/IR/ConstantRangeTest.cpp

367

Also, could you please do this as NFC commit

[nikic]

Rebase over committed changes.

[nikic]

Add extra diagrams.

[nikic]

llvm/lib/IR/ConstantRange.cpp

470

I've included this change while committing D59959.

608–610

Right. I've added a comment.

634–636

Right. I've added the two possible results to the existing comment.

llvm/unittests/IR/ConstantRangeTest.cpp

367

Done in rL357874.

[lebedev.ri]

Okay, looks good to me, thank you!

llvm/include/llvm/IR/ConstantRange.h

public:
   ConstantRange intersectWith(const ConstantRange &CR,
                               PreferredRangeType Type = Smallest) const;
 
   /// Return the range that results from the union of this range
   /// with another range.  The resultant range is guaranteed to include the
   /// elements of both sets, but may contain more.  For example, [3, 9) union
   /// [12,15) is [3, 15), which includes 9, 10, and 11, which were not included
   /// in either set before.
-  ConstantRange unionWith(const ConstantRange &CR) const;
+  ConstantRange unionWith(const ConstantRange &CR,
+                          PreferredRangeType Type = Smallest) const;
 
   /// Return a new range representing the possible values resulting
   /// from an application of the specified cast operator to this range. \p
   /// BitWidth is the target bitwidth of the cast.  For casts which don't
   /// change bitwidth, it must be the same as the source bitwidth.  For casts
   /// which do change bitwidth, the bitwidth must be consistent with the
   /// requested cast and source bitwidth.
   ConstantRange castOp(Instruction::CastOps CastOp,

llvm/lib/IR/ConstantRange.cpp

if (Lower == Upper)
     return *this;
   return ConstantRange(Lower - Val, Upper - Val);
 }
 
 ConstantRange ConstantRange::difference(const ConstantRange &CR) const {
   return intersectWith(CR.inverse());
 }
 
 static ConstantRange getPreferredRange(
     const ConstantRange &CR1, const ConstantRange &CR2,
     ConstantRange::PreferredRangeType Type) {
   if (Type == ConstantRange::Unsigned) {
     if (!CR1.isWrappedSet() && CR2.isWrappedSet())
       return CR1;
     if (CR1.isWrappedSet() && !CR2.isWrappedSet())
       return CR2;
   } else if (Type == ConstantRange::Signed) {
if (CR.Upper.ule(Lower)) {
     return ConstantRange(CR.Lower, Upper);
   }
 
   // --U L------ : this
   // ------U L-- : CR
   return getPreferredRange(*this, CR, Type);
 }
 
-ConstantRange ConstantRange::unionWith(const ConstantRange &CR) const {
+ConstantRange ConstantRange::unionWith(const ConstantRange &CR,
+                                       PreferredRangeType Type) const {
   assert(getBitWidth() == CR.getBitWidth() &&
          "ConstantRange types don't agree!");
 
   if (   isFullSet() || CR.isEmptySet()) return *this;
   if (CR.isFullSet() ||    isEmptySet()) return CR;
 
-  if (!isUpperWrapped() && CR.isUpperWrapped()) return CR.unionWith(*this);
+  if (!isUpperWrapped() && CR.isUpperWrapped())
+    return CR.unionWith(*this, Type);
 
   if (!isUpperWrapped() && !CR.isUpperWrapped()) {
-    if (CR.Upper.ult(Lower) || Upper.ult(CR.Lower)) {
-      // If the two ranges are disjoint, find the smaller gap and bridge it.
-      APInt d1 = CR.Lower - Upper, d2 = Lower - CR.Upper;
-      if (d1.ult(d2))
-        return ConstantRange(Lower, CR.Upper);
-      return ConstantRange(CR.Lower, Upper);
-    }
+    if (CR.Upper.ult(Lower) || Upper.ult(CR.Lower))
+      return getPreferredRange(
+          ConstantRange(Lower, CR.Upper), ConstantRange(CR.Lower, Upper), Type);
 
     APInt L = CR.Lower.ult(Lower) ? CR.Lower : Lower;
     APInt U = (CR.Upper - 1).ugt(Upper - 1) ? CR.Upper : Upper;
 
     if (L.isNullValue() && U.isNullValue())
       return getFull();
 
     return ConstantRange(std::move(L), std::move(U));
   }
 
   if (!CR.isUpperWrapped()) {
     // ------U   L-----  and  ------U   L----- : this
     //   L--U                            L--U  : CR
     if (CR.Upper.ule(Upper) || CR.Lower.uge(Lower))
       return *this;
 
     // ------U   L----- : this
     //    L---------U   : CR
     if (CR.Lower.ule(Upper) && Lower.ule(CR.Upper))
       return getFull();
 
     // ----U       L---- : this
     //       L---U       : CR
-    //    <d1>  <d2>
-    if (Upper.ule(CR.Lower) && CR.Upper.ule(Lower)) {
-      APInt d1 = CR.Lower - Upper, d2 = Lower - CR.Upper;
-      if (d1.ult(d2))
-        return ConstantRange(Lower, CR.Upper);
-      return ConstantRange(CR.Lower, Upper);
-    }
+    if (Upper.ult(CR.Lower) && CR.Upper.ult(Lower))
+      return getPreferredRange(
+          ConstantRange(Lower, CR.Upper), ConstantRange(CR.Lower, Upper), Type);
 
     // ----U     L----- : this
     //        L----U    : CR
-    if (Upper.ult(CR.Lower) && Lower.ult(CR.Upper))
+    if (Upper.ult(CR.Lower) && Lower.ule(CR.Upper))
       return ConstantRange(CR.Lower, Upper);
 
     // ------U    L---- : this
     //    L-----U       : CR
-    assert(CR.Lower.ult(Upper) && CR.Upper.ult(Lower) &&
+    assert(CR.Lower.ule(Upper) && CR.Upper.ult(Lower) &&
            "ConstantRange::unionWith missed a case with one range wrapped");
     return ConstantRange(Lower, CR.Upper);
   }
 
   // ------U    L----  and  ------U    L---- : this
   // -U  L-----------  and  ------------U  L : CR
   if (CR.Lower.ule(Upper) || Lower.ule(CR.Upper))
     return getFull();

llvm/unittests/IR/ConstantRangeTest.cpp

TEST_F(ConstantRangeTest, IntersectWith) {
  
   // [15, 0) /\ [7, 6) = [15, 0)
   LHS = ConstantRange(APInt(32, 15), APInt(32, 0));
   RHS = ConstantRange(APInt(32, 7), APInt(32, 6));
   EXPECT_EQ(LHS.intersectWith(RHS), ConstantRange(APInt(32, 15), APInt(32, 0)));
 }
 
 template<typename Fn1, typename Fn2>
 void testBinarySetOperationExhaustive(Fn1 OpFn, Fn2 InResultFn) {
   unsigned Bits = 4;
   EnumerateTwoConstantRanges(Bits,
       [=](const ConstantRange &CR1, const ConstantRange &CR2) {
         // Collect up to three contiguous unsigned ranges. The HaveInterrupt
         // variables are used determine when we have to switch to the next
         // range because the previous one ended.
         APInt Lower1(Bits, 0), Upper1(Bits, 0);
         APInt Lower2(Bits, 0), Upper2(Bits, 0);
testBinarySetOperationExhaustive(
          ConstantRange::PreferredRangeType Type) {
         return CR1.intersectWith(CR2, Type);
       },
       [](const ConstantRange &CR1, const ConstantRange &CR2, const APInt &N) {
         return CR1.contains(N) && CR2.contains(N);
       });
 }
 
+TEST_F(ConstantRangeTest, UnionWithExhaustive) {
+  testBinarySetOperationExhaustive(
+      [](const ConstantRange &CR1, const ConstantRange &CR2,
+         ConstantRange::PreferredRangeType Type) {
+        return CR1.unionWith(CR2, Type);
+      },
+      [](const ConstantRange &CR1, const ConstantRange &CR2, const APInt &N) {
+        return CR1.contains(N) || CR2.contains(N);
+      });
+}
+
 TEST_F(ConstantRangeTest, UnionWith) {
   EXPECT_EQ(Wrap.unionWith(One),
             ConstantRange(APInt(16, 0xaaa), APInt(16, 0xb)));
   EXPECT_EQ(One.unionWith(Wrap), Wrap.unionWith(One));
   EXPECT_EQ(Empty.unionWith(Empty), Empty);
   EXPECT_EQ(Full.unionWith(Full), Full);
   EXPECT_EQ(Some.unionWith(Wrap), Full);