[KnownBits] Check functions that return zero for poison results

[nikic]

LGTM

llvm/unittests/Support/KnownBitsTest.cpp

 
 using BinaryBitsFn =
     llvm::function_ref<KnownBits(const KnownBits &, const KnownBits &)>;
 using BinaryIntFn =
     llvm::function_ref<std::optional<APInt>(const APInt &, const APInt &)>;
 using BinaryCheckFn =
     llvm::function_ref<bool(const KnownBits &, const KnownBits &)>;
 
+static bool checkOptimalityUnary(const KnownBits &) { return true; }
 static bool checkCorrectnessOnlyUnary(const KnownBits &) { return false; }
+static bool checkOptimalityBinary(const KnownBits &, const KnownBits &) {
+  return true;
+}
 static bool checkCorrectnessOnlyBinary(const KnownBits &, const KnownBits &) {
   return false;
 }
 
 static testing::AssertionResult isCorrect(const KnownBits &Exact,
                                           const KnownBits &Computed,
                                           ArrayRef<KnownBits> Inputs) {
   if (Computed.Zero.isSubsetOf(Exact.Zero) &&
static testing::AssertionResult isOptimal(const KnownBits &Exact,
   testing::AssertionResult Result = testing::AssertionFailure();
   Result << "Inputs = ";
   for (const KnownBits &Input : Inputs)
     Result << Input << ", ";
   Result << "Computed = " << Computed << ", Exact = " << Exact;
   return Result;
 }
 
-static void testUnaryOpExhaustive(
-    UnaryBitsFn BitsFn, UnaryIntFn IntFn,
-    UnaryCheckFn CheckOptimalityFn = [](const KnownBits &) { return true; }) {
+static void
+testUnaryOpExhaustive(UnaryBitsFn BitsFn, UnaryIntFn IntFn,
+                      UnaryCheckFn CheckOptimalityFn = checkOptimalityUnary) {
   unsigned Bits = 4;
   ForeachKnownBits(Bits, [&](const KnownBits &Known) {
     KnownBits Computed = BitsFn(Known);
     KnownBits Exact(Bits);
     Exact.Zero.setAllBits();
     Exact.One.setAllBits();
 
     ForeachNumInKnownBits(Known, [&](const APInt &N) {
       if (std::optional<APInt> Res = IntFn(N)) {
         Exact.One &= *Res;
         Exact.Zero &= ~*Res;
       }
     });
 
     EXPECT_TRUE(!Computed.hasConflict());
     EXPECT_TRUE(isCorrect(Exact, Computed, Known));
     // We generally don't want to return conflicting known bits, even if it is
     // legal for always poison results.
     if (CheckOptimalityFn(Known) && !Exact.hasConflict()) {
       EXPECT_TRUE(isOptimal(Exact, Computed, Known));
     }
   });
 }
 
-static void testBinaryOpExhaustive(
-    BinaryBitsFn BitsFn, BinaryIntFn IntFn,
-    BinaryCheckFn CheckOptimalityFn = [](const KnownBits &, const KnownBits &) {
-      return true;
-    }) {
+static void
+testBinaryOpExhaustive(BinaryBitsFn BitsFn, BinaryIntFn IntFn,
+                       BinaryCheckFn CheckOptimalityFn = checkOptimalityBinary,
+                       bool RefinePoisonToZero = false) {
   unsigned Bits = 4;
   ForeachKnownBits(Bits, [&](const KnownBits &Known1) {
     ForeachKnownBits(Bits, [&](const KnownBits &Known2) {
       KnownBits Computed = BitsFn(Known1, Known2);
       KnownBits Exact(Bits);
       Exact.Zero.setAllBits();
       Exact.One.setAllBits();
 
       ForeachNumInKnownBits(Known1, [&](const APInt &N1) {
         ForeachNumInKnownBits(Known2, [&](const APInt &N2) {
           if (std::optional<APInt> Res = IntFn(N1, N2)) {
             Exact.One &= *Res;
             Exact.Zero &= ~*Res;
           }
         });
       });
 
       EXPECT_TRUE(!Computed.hasConflict());
       EXPECT_TRUE(isCorrect(Exact, Computed, {Known1, Known2}));
       // We generally don't want to return conflicting known bits, even if it is
       // legal for always poison results.
       if (CheckOptimalityFn(Known1, Known2) && !Exact.hasConflict()) {
         EXPECT_TRUE(isOptimal(Exact, Computed, {Known1, Known2}));
       }
+      // In some cases we choose to return zero if the result is always poison.
+      if (RefinePoisonToZero && Exact.hasConflict()) {
+        EXPECT_TRUE(Computed.isZero());
+      }
     });
   });
 }
 
 namespace {
 
 TEST(KnownBitsTest, AddCarryExhaustive) {
   unsigned Bits = 4;
TEST(KnownBitsTest, BinaryExhaustive) {
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::shl(Known1, Known2);
       },
       [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
         if (N2.uge(N2.getBitWidth()))
           return std::nullopt;
         return N1.shl(N2);
-      });
+      },
+      checkOptimalityBinary, /* RefinePoisonToZero */ true);
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::shl(Known1, Known2, /* NUW */ true);
       },
       [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
         bool Overflow;
         APInt Res = N1.ushl_ov(N2, Overflow);
         if (Overflow)
           return std::nullopt;
         return Res;
-      });
+      },
+      checkOptimalityBinary, /* RefinePoisonToZero */ true);
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::shl(Known1, Known2, /* NUW */ false, /* NSW */ true);
       },
       [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
         bool Overflow;
         APInt Res = N1.sshl_ov(N2, Overflow);
         if (Overflow)
           return std::nullopt;
         return Res;
-      });
+      },
+      checkOptimalityBinary, /* RefinePoisonToZero */ true);
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::shl(Known1, Known2, /* NUW */ true, /* NSW */ true);
       },
       [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
         bool OverflowUnsigned, OverflowSigned;
         APInt Res = N1.ushl_ov(N2, OverflowUnsigned);
         (void)N1.sshl_ov(N2, OverflowSigned);
         if (OverflowUnsigned || OverflowSigned)
           return std::nullopt;
         return Res;
-      });
+      },
+      checkOptimalityBinary, /* RefinePoisonToZero */ true);
 
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::lshr(Known1, Known2);
       },
       [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
         if (N2.uge(N2.getBitWidth()))
           return std::nullopt;
         return N1.lshr(N2);
-      });
+      },
+      checkOptimalityBinary, /* RefinePoisonToZero */ true);
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::ashr(Known1, Known2);
       },
       [](const APInt &N1, const APInt &N2) -> std::optional<APInt> {
         if (N2.uge(N2.getBitWidth()))
           return std::nullopt;
         return N1.ashr(N2);
-      });
+      },
+      checkOptimalityBinary, /* RefinePoisonToZero */ true);
 
   testBinaryOpExhaustive(
       [](const KnownBits &Known1, const KnownBits &Known2) {
         return KnownBits::mul(Known1, Known2);
       },
       [](const APInt &N1, const APInt &N2) { return N1 * N2; },
       checkCorrectnessOnlyBinary);
   testBinaryOpExhaustive(