From 2730f84bd06da08299d1921fdc348123307a6dba Mon Sep 17 00:00:00 2001
From: Simon Pilgrim <llvm-dev@redking.me.uk>
Date: Wed, 4 Nov 2020 18:19:26 +0000
Subject: [PATCH] [UnitTests] Add exhaustive tests for KnownBits shl/lshr/ashr

Thanks to @foad for the suggestion.
---
 unittests/Support/KnownBitsTest.cpp | 33 +++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/unittests/Support/KnownBitsTest.cpp b/unittests/Support/KnownBitsTest.cpp
index d2f766fb334..57124cb503b 100644
--- a/unittests/Support/KnownBitsTest.cpp
+++ b/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));
     });
   });
 }