[UnitTests] Add exhaustive tests for KnownBits shl/lshr/ashr

Thanks to @foad for the suggestion.

diff --git a/llvm/unittests/Support/KnownBitsTest.cpp b/llvm/unittests/Support/KnownBitsTest.cpp
index d2f766f..57124cb 100644 (file)
--- a/llvm/unittests/Support/KnownBitsTest.cpp
+++ b/llvm/unittests/Support/KnownBitsTest.cpp
@@ -113,6 +113,9 @@ TEST(KnownBitsTest, BinaryExhaustive) {
       KnownBits KnownSMax(KnownAnd);
       KnownBits KnownSMin(KnownAnd);
       KnownBits KnownMul(KnownAnd);
+      KnownBits KnownShl(KnownAnd);
+      KnownBits KnownLShr(KnownAnd);
+      KnownBits KnownAShr(KnownAnd);
 
       ForeachNumInKnownBits(Known1, [&](const APInt &N1) {
         ForeachNumInKnownBits(Known2, [&](const APInt &N2) {
@@ -149,6 +152,24 @@ TEST(KnownBitsTest, BinaryExhaustive) {
           Res = N1 * N2;
           KnownMul.One &= Res;
           KnownMul.Zero &= ~Res;
+
+          if (N2.ult(1ULL << N1.getBitWidth())) {
+            Res = N1.shl(N2);
+            KnownShl.One &= Res;
+            KnownShl.Zero &= ~Res;
+
+            Res = N1.lshr(N2);
+            KnownLShr.One &= Res;
+            KnownLShr.Zero &= ~Res;
+
+            Res = N1.ashr(N2);
+            KnownAShr.One &= Res;
+            KnownAShr.Zero &= ~Res;
+          } else {
+            KnownShl.resetAll();
+            KnownLShr.resetAll();
+            KnownAShr.resetAll();
+          }
         });
       });
 
@@ -185,6 +206,18 @@ TEST(KnownBitsTest, BinaryExhaustive) {
       KnownBits ComputedMul = KnownBits::computeForMul(Known1, Known2);
       EXPECT_TRUE(ComputedMul.Zero.isSubsetOf(KnownMul.Zero));
       EXPECT_TRUE(ComputedMul.One.isSubsetOf(KnownMul.One));
+
+      KnownBits ComputedShl = KnownBits::shl(Known1, Known2);
+      EXPECT_TRUE(ComputedShl.Zero.isSubsetOf(KnownShl.Zero));
+      EXPECT_TRUE(ComputedShl.One.isSubsetOf(KnownShl.One));
+
+      KnownBits ComputedLShr = KnownBits::lshr(Known1, Known2);
+      EXPECT_TRUE(ComputedLShr.Zero.isSubsetOf(KnownLShr.Zero));
+      EXPECT_TRUE(ComputedLShr.One.isSubsetOf(KnownLShr.One));
+
+      KnownBits ComputedAShr = KnownBits::ashr(Known1, Known2);
+      EXPECT_TRUE(ComputedAShr.Zero.isSubsetOf(KnownAShr.Zero));
+      EXPECT_TRUE(ComputedAShr.One.isSubsetOf(KnownAShr.One));
     });
   });
 }