[ConstantRange] Support checking optimality for subset of inputs (NFC)

We always want to check correctness, but for some operations we
can only guarantee optimality for a subset of inputs. Accept an
additional predicate that determines whether optimality for a
given pair of ranges should be checked.

llvm/unittests/IR/ConstantRangeTest.cpp

@@ -111,14 +111,18 @@ testing::AssertionResult rangeContains(const ConstantRange &CR, const APInt &N,
 // Elems under the given PreferenceFn. The preference function should return
 // true if the first range argument is strictly preferred to the second one.
 static void TestRange(const ConstantRange &CR, const SmallBitVector &Elems,
-                      PreferFn PreferenceFn, ArrayRef<ConstantRange> Inputs) {
+                      PreferFn PreferenceFn, ArrayRef<ConstantRange> Inputs,
+                      bool CheckOptimality = true) {
   unsigned BitWidth = CR.getBitWidth();
 
   // Check conservative correctness.
   for (unsigned Elem : Elems.set_bits()) {
     EXPECT_TRUE(rangeContains(CR, APInt(BitWidth, Elem), Inputs));
   }
 
+  if (!CheckOptimality)
+    return;
+
   // Make sure we have at least one element for the code below.
   if (Elems.none()) {
     EXPECT_TRUE(CR.isEmptySet());
@@ -182,9 +186,23 @@ using BinaryRangeFn = llvm::function_ref<ConstantRange(const ConstantRange &,
                                                        const ConstantRange &)>;
 using BinaryIntFn = llvm::function_ref<Optional<APInt>(const APInt &,
                                                        const APInt &)>;
+using BinaryCheckFn = llvm::function_ref<bool(const ConstantRange &,
+                                              const ConstantRange &)>;
+
+static bool CheckAll(const ConstantRange &, const ConstantRange &) {
+  return true;
+}
+
+static bool CheckSingleElementsOnly(const ConstantRange &CR1,
+                                    const ConstantRange &CR2) {
+  return CR1.isSingleElement() && CR2.isSingleElement();
+}
 
+// CheckFn determines whether optimality is checked for a given range pair.
+// Correctness is always checked.
 static void TestBinaryOpExhaustive(BinaryRangeFn RangeFn, BinaryIntFn IntFn,
-                                   PreferFn PreferenceFn = PreferSmallest) {
+                                   PreferFn PreferenceFn = PreferSmallest,
+                                   BinaryCheckFn CheckFn = CheckAll) {
   unsigned Bits = 4;
   EnumerateTwoConstantRanges(
       Bits, [&](const ConstantRange &CR1, const ConstantRange &CR2) {
@@ -195,23 +213,8 @@ static void TestBinaryOpExhaustive(BinaryRangeFn RangeFn, BinaryIntFn IntFn,
               Elems.set(ResultN->getZExtValue());
           });
         });
-        TestRange(RangeFn(CR1, CR2), Elems, PreferenceFn, {CR1, CR2});
-      });
-}
-
-static void TestBinaryOpExhaustiveCorrectnessOnly(BinaryRangeFn RangeFn,
-                                                  BinaryIntFn IntFn) {
-  unsigned Bits = 4;
-  EnumerateTwoConstantRanges(
-      Bits, [&](const ConstantRange &CR1, const ConstantRange &CR2) {
-        ConstantRange ResultCR = RangeFn(CR1, CR2);
-        ForeachNumInConstantRange(CR1, [&](const APInt &N1) {
-          ForeachNumInConstantRange(CR2, [&](const APInt &N2) {
-            if (Optional<APInt> ResultN = IntFn(N1, N2)) {
-              EXPECT_TRUE(rangeContains(ResultCR, *ResultN, {CR1, CR2}));
-            }
-          });
-        });
+        TestRange(RangeFn(CR1, CR2), Elems, PreferenceFn, {CR1, CR2},
+                  CheckFn(CR1, CR2));
       });
 }
 
@@ -1303,15 +1306,17 @@ TEST_F(ConstantRangeTest, URem) {
                 .urem(ConstantRange(APInt(16, 10))),
             ConstantRange(APInt(16, 0), APInt(16, 10)));
 
-  TestBinaryOpExhaustiveCorrectnessOnly(
+  TestBinaryOpExhaustive(
       [](const ConstantRange &CR1, const ConstantRange &CR2) {
         return CR1.urem(CR2);
       },
       [](const APInt &N1, const APInt &N2) -> Optional<APInt> {
         if (N2.isZero())
           return None;
         return N1.urem(N2);
-      });
+      },
+      PreferSmallest,
+      CheckSingleElementsOnly);
 }
 
 TEST_F(ConstantRangeTest, SRem) {
@@ -1377,15 +1382,17 @@ TEST_F(ConstantRangeTest, SRem) {
                 .srem(ConstantRange(APInt(16, 10))),
             ConstantRange(APInt(16, 0), APInt(16, 10)));
 
-  TestBinaryOpExhaustiveCorrectnessOnly(
+  TestBinaryOpExhaustive(
       [](const ConstantRange &CR1, const ConstantRange &CR2) {
         return CR1.srem(CR2);
       },
       [](const APInt &N1, const APInt &N2) -> Optional<APInt> {
         if (N2.isZero())
           return None;
         return N1.srem(N2);
-      });
+      },
+      PreferSmallest,
+      CheckSingleElementsOnly);
 }
 
 TEST_F(ConstantRangeTest, Shl) {