SmallBitVector: Rework find_first/find_next and tweak test to test them (at least on 64 bit platforms).

llvm-svn: 102712

2 files changed, 15 insertions, 16 deletions

diff --git a/llvm/include/llvm/ADT/SmallBitVector.h b/llvm/include/llvm/ADT/SmallBitVector.h
index 884e631c88b..3441d0a90c9 100644
--- a/llvm/include/llvm/ADT/SmallBitVector.h
+++ b/llvm/include/llvm/ADT/SmallBitVector.h
@@ -199,13 +199,12 @@ public:
   int find_first() const {
     if (isSmall()) {
       uintptr_t Bits = getSmallBits();
-      if (sizeof(uintptr_t) * CHAR_BIT == 32) {
-        size_t FirstBit = CountTrailingZeros_32(Bits);
-        return FirstBit == 32 ? -1 : FirstBit;
-      } else if (sizeof(uintptr_t) * CHAR_BIT == 64) {
-        size_t FirstBit = CountTrailingZeros_64(Bits);
-        return FirstBit == 64 ? -1 : FirstBit;
-      }
+      if (Bits == 0)
+        return -1;
+      if (sizeof(uintptr_t) * CHAR_BIT == 32)
+        return CountTrailingZeros_32(Bits);
+      if (sizeof(uintptr_t) * CHAR_BIT == 64)
+        return CountTrailingZeros_64(Bits);
       assert(0 && "Unsupported!");
     }
     return getPointer()->find_first();
@@ -218,13 +217,12 @@ public:
       uintptr_t Bits = getSmallBits();
       // Mask off previous bits.
       Bits &= ~uintptr_t(0) << (Prev + 1);
-      if (sizeof(uintptr_t) * CHAR_BIT == 32) {
-        size_t FirstBit = CountTrailingZeros_32(Bits);
-        return FirstBit == 32 ? -1 : FirstBit;
-      } else if (sizeof(uintptr_t) * CHAR_BIT == 64) {
-        size_t FirstBit = CountTrailingZeros_64(Bits);
-        return FirstBit == 64 ? -1 : FirstBit;
-      }
+      if (Bits == 0 || Prev + 1 >= getSmallSize())
+        return -1;
+      if (sizeof(uintptr_t) * CHAR_BIT == 32)
+        return CountTrailingZeros_32(Bits);
+      if (sizeof(uintptr_t) * CHAR_BIT == 64)
+        return CountTrailingZeros_64(Bits);
       assert(0 && "Unsupported!");
     }
     return getPointer()->find_next(Prev);
diff --git a/llvm/unittests/ADT/SmallBitVectorTest.cpp b/llvm/unittests/ADT/SmallBitVectorTest.cpp
index a0c079d874a..e12abfe20da 100644
--- a/llvm/unittests/ADT/SmallBitVectorTest.cpp
+++ b/llvm/unittests/ADT/SmallBitVectorTest.cpp
@@ -55,7 +55,7 @@ TEST(SmallBitVectorTest, TrivialOperation) {
   Vec.resize(26, true);
   Vec.resize(29, false);
   Vec.resize(33, true);
-  Vec.resize(61, false);
+  Vec.resize(57, false);
   unsigned Count = 0;
   for (unsigned i = Vec.find_first(); i != -1u; i = Vec.find_next(i)) {
     ++Count;
@@ -66,7 +66,8 @@ TEST(SmallBitVectorTest, TrivialOperation) {
   EXPECT_EQ(Count, 23u);
   EXPECT_FALSE(Vec[0]);
   EXPECT_TRUE(Vec[32]);
-  EXPECT_FALSE(Vec[60]);
+  EXPECT_FALSE(Vec[56]);
+  Vec.resize(61, false);
 
   SmallBitVector Copy = Vec;
   SmallBitVector Alt(3, false);