[BitVector] Implement find_[first/last]_[set/unset]_in

[zturner]

Summary

This effort started when I wanted a method that did the following:

Given a BitVector and indices N and M, find the largest index less than M but greater than or equal to N which contains a set bit.

At first I just used find_prev(M) which I had added in an earlier patch, but I quickly discovered the performance of this was poor. My BitVector had over 500,000 bits in it, and in the worst case this could cause an O(n) scan all the way to the beginning. But I already had a lower bound on the range I wanted to search in, so I figured I could just make a function find_last_in(LowerBound, PriorTo).

But it quickly became apparent that if we already have this function, then find_last() is just find_last_in(0, Size) and find_prev(PriorTo) is just find_last_in(0, PriorTo).

In fact, all of the methods are generalizable this way. For example, find_first and find_next reduce to find_first_in(0, Size) and find_first_in(After, Size). The generalizations to the find_xxx_unset family of methods is obvious.

So I went all the way. There is now a find_in method for all combinations of [first, last] x [set, unset], and all find methods delegate to these methods.

[zturner]

And.... this apparently breaks everything. Looking into why, all my unittests passed, so it's something more subtle.

[zturner]

Fix the breakage. Given a BitVector of size N where the last bit (e.g. index N-1) was set, the failure occurred when you called find_first_in(N-1, Size). It should have been returning N-1 but it was returning -1. A test was added for this case, and edge boundary tests were added for all other functions as well.

[zturner]

Remove the "old" versions of the functions. I tracked down the problem by copy/pasting the original versions back in and asserting any time they didn't return the same value. Forgot to remove them in my previous upload though.

[zturner]

ping

[chandlerc]

Very nice generally. Comments only on one, but clearly apply to each variation.

Bit of a shame there isn't a clean way to share all the logic, but all the ideas to do that are *awful*. =/

llvm/include/llvm/ADT/BitVector.h

151

Given that we return an int, I would prefer that the arguments also be int (and the intermediate variables within).

152–153

Why isn't this an assert?

[zturner]

llvm/include/llvm/ADT/BitVector.h

151

The only reason we return an int is because we have to use a sentinel value to indicate "no bits found". The methods that these originated from find_next, find_prev, etc have always taken unsigned arguments even before I started tinkering on BitVector. Do you still think it's a good idea to change all of them?

152–153

I suppose I should assert if Begin > End, but Begin == End should probably still just return -1 I think.

[chandlerc]

LGTM with the assert changes.

llvm/include/llvm/ADT/BitVector.h

151

I mean, my view is that we're doing a lot of arithmetic on these things including potentially subtraction. So generally, unless there is a problem with using a signed integer, I would prefer that.

If we return int, then we've given up on having more than 2^31 - 1 bits in one of these containers.

That said, a bitvector is one of the few datastructures where it seems entirely conceivable that we'll actually end up with 500mb vector of bits and an int index won't work. But an unsigned won't either, so I suspect that isn't the deciding factor and if we care we need a 64-bit type.

So I guess my feeling is that we should do one of two things here:
a) Switch the return types to an unsigned type and use something like npos, or
b) Switch the argument types to a signed type and continue to use -1.

(Whether we use 32-bits or 64-bits should be orthogonal as I see no reason why for a *bit* count 2gb vs 4gb would be "enough".)

However, looking at the code, I agree that the patch as-is remains consistent, and you shouldn't do either (a) or (b) in *this* patch. But I'd somewhat like to consider doing one of them in a subsequent patch.

My personal, strong preference is to use (b) because I find asserts easier to write for signed types in the face of arithmetic and we have strong usage of UBSan to ensure we don't develop bugs here. I also have a somewhat powerful dislike of npos. But I think either (a) or (b) would be an improvement. Anyways, sorry for the digression.

152–153

Ah, yes, that makes sense.

llvm/include/llvm/ADT/BitVector.h

bool all() const {
     return true;
   }
 
   /// none - Returns true if none of the bits are set.
   bool none() const {
     return !any();
   }
 
-  /// find_first - Returns the index of the first set bit, -1 if none
-  /// of the bits are set.
-  int find_first() const {
-    for (unsigned i = 0; i < NumBitWords(size()); ++i)
-      if (Bits[i] != 0)
-        return i * BITWORD_SIZE + countTrailingZeros(Bits[i]);
+  /// find_first_in - Returns the index of the first set bit in the range
+  /// [Begin, End).  Returns -1 if all bits in the range are unset.
+  int find_first_in(unsigned Begin, unsigned End) const {
+    if (Begin >= End)
       return -1;
-  }
 
-  /// find_last - Returns the index of the last set bit, -1 if none of the bits
-  /// are set.
-  int find_last() const {
-    if (Size == 0)
-      return -1;
+    assert(End <= size());
 
-    unsigned N = NumBitWords(size());
-    assert(N > 0);
+    unsigned FirstWord = Begin / BITWORD_SIZE;
+    unsigned LastWord = (End - 1) / BITWORD_SIZE;
 
-    unsigned i = N - 1;
-    while (i > 0 && Bits[i] == BitWord(0))
-      --i;
+    // Check subsequent words.
+    for (unsigned i = FirstWord; i <= LastWord; ++i) {
+      BitWord Copy = Bits[i];
 
-    return int((i + 1) * BITWORD_SIZE - countLeadingZeros(Bits[i])) - 1;
+      if (i == FirstWord) {
+        unsigned FirstBit = Begin % BITWORD_SIZE;
+        Copy &= maskTrailingZeros<BitWord>(FirstBit);
       }
 
-  /// find_first_unset - Returns the index of the first unset bit, -1 if all
-  /// of the bits are set.
-  int find_first_unset() const {
-    for (unsigned i = 0; i < NumBitWords(size()); ++i)
-      if (Bits[i] != ~0UL) {
-        unsigned Result = i * BITWORD_SIZE + countTrailingOnes(Bits[i]);
-        return Result < size() ? Result : -1;
-      }
+      if (i == LastWord) {
+        unsigned LastBit = (End - 1) % BITWORD_SIZE;
+        Copy &= maskTrailingOnes<BitWord>(LastBit);
+      }
+      if (Copy != 0)
+        return i * BITWORD_SIZE + countTrailingZeros(Copy);
     return -1;
   }
 
-  /// find_last_unset - Returns the index of the last unset bit, -1 if all of
-  /// the bits are set.
-  int find_last_unset() const {
-    if (Size == 0)
+  /// find_last_in - Returns the index of the last set bit in the range
+  /// [Begin, End).  Returns -1 if all bits in the range are unset.
+  int find_last_in(unsigned Begin, unsigned End) const {
+    if (End == 0 || Begin >= End)
       return -1;
 
-    const unsigned N = NumBitWords(size());
-    assert(N > 0);
+    assert(End <= Size);
 
-    unsigned i = N - 1;
-    BitWord W = Bits[i];
+    unsigned LastWord = (End - 1) / BITWORD_SIZE;
+    unsigned FirstWord = Begin / BITWORD_SIZE;
 
-    // The last word in the BitVector has some unused bits, so we need to set
-    // them all to 1 first.  Set them all to 1 so they don't get treated as
-    // valid unset bits.
-    unsigned UnusedCount = BITWORD_SIZE - size() % BITWORD_SIZE;
-    W |= maskLeadingOnes<BitWord>(UnusedCount);
+    for (unsigned i = LastWord + 1; i >= FirstWord + 1; --i) {
+      unsigned CurrentWord = i - 1;
 
-    while (W == ~BitWord(0) && --i > 0)
-      W = Bits[i];
+      BitWord Copy = Bits[CurrentWord];
+      if (CurrentWord == LastWord) {
+        unsigned LastBit = (End - 1) % BITWORD_SIZE;
+        Copy &= maskTrailingOnes<BitWord>(LastBit + 1);
+      }
 
-    return int((i + 1) * BITWORD_SIZE - countLeadingOnes(W)) - 1;
+      if (CurrentWord == FirstWord) {
+        unsigned FirstBit = Begin % BITWORD_SIZE;
+        Copy &= maskTrailingZeros<BitWord>(FirstBit);
+      }
+
+      if (Copy != 0)
+        return (CurrentWord + 1) * BITWORD_SIZE - countLeadingZeros(Copy) - 1;
     }
 
-  /// find_next - Returns the index of the next set bit following the
-  /// "Prev" bit. Returns -1 if the next set bit is not found.
-  int find_next(unsigned Prev) const {
-    ++Prev;
-    if (Prev >= Size)
-      return -1;
+  }
 
-    unsigned WordPos = Prev / BITWORD_SIZE;
-    unsigned BitPos = Prev % BITWORD_SIZE;
-    BitWord Copy = Bits[WordPos];
-    // Mask off previous bits.
-    Copy &= maskTrailingZeros<BitWord>(BitPos);
+  /// find_first_unset_in - Returns the index of the first unset bit in the
+  /// range [Begin, End).  Returns -1 if all bits in the range are set.
+  int find_first_unset_in(unsigned Begin, unsigned End) const {
+    if (Begin >= End)
+      return -1;
 
-    if (Copy != 0)
-      return WordPos * BITWORD_SIZE + countTrailingZeros(Copy);
+    assert(End <= size());
+
+    unsigned FirstWord = Begin / BITWORD_SIZE;
+    unsigned LastWord = (End - 1) / BITWORD_SIZE;
 
     // Check subsequent words.
-    for (unsigned i = WordPos+1; i < NumBitWords(size()); ++i)
-      if (Bits[i] != 0)
-        return i * BITWORD_SIZE + countTrailingZeros(Bits[i]);
+    for (unsigned i = FirstWord; i <= LastWord; ++i) {
+      BitWord Copy = Bits[i];
+
+      if (i == FirstWord) {
+        unsigned FirstBit = Begin % BITWORD_SIZE;
+        Copy |= maskTrailingOnes<BitWord>(FirstBit);
+      }
+
+      if (i == LastWord) {
+        unsigned LastBit = (End - 1) % BITWORD_SIZE;
+        Copy |= maskTrailingZeros<BitWord>(LastBit + 1);
+      }
+      if (Copy != ~0UL) {
+        unsigned Result = i * BITWORD_SIZE + countTrailingOnes(Copy);
+        return Result < size() ? Result : -1;
+      }
+    }
     return -1;
   }
 
-  /// find_next_unset - Returns the index of the next unset bit following the
-  /// "Prev" bit.  Returns -1 if all remaining bits are set.
-  int find_next_unset(unsigned Prev) const {
-    ++Prev;
-    if (Prev >= Size)
+  /// find_last_unset_in - Returns the index of the last unset bit in the
+  /// range [Begin, End).  Returns -1 if all bits in the range are set.
+  int find_last_unset_in(unsigned Begin, unsigned End) const {
+    if (End == 0 || Begin >= End)
       return -1;
 
-    unsigned WordPos = Prev / BITWORD_SIZE;
-    unsigned BitPos = Prev % BITWORD_SIZE;
-    BitWord Copy = Bits[WordPos];
-    // Mask in previous bits.
-    BitWord Mask = (1 << BitPos) - 1;
-    Copy |= Mask;
+    assert(End <= Size);
 
-    if (Copy != ~0UL)
-      return next_unset_in_word(WordPos, Copy);
+    unsigned LastWord = (End - 1) / BITWORD_SIZE;
+    unsigned FirstWord = Begin / BITWORD_SIZE;
 
-    // Check subsequent words.
-    for (unsigned i = WordPos + 1; i < NumBitWords(size()); ++i)
-      if (Bits[i] != ~0UL)
-        return next_unset_in_word(i, Bits[i]);
+    for (unsigned i = LastWord + 1; i >= FirstWord + 1; --i) {
+      unsigned CurrentWord = i - 1;
+
+      BitWord Copy = Bits[CurrentWord];
+      if (CurrentWord == LastWord) {
+        unsigned LastBit = (End - 1) % BITWORD_SIZE;
+        Copy |= maskTrailingZeros<BitWord>(LastBit + 1);
+      }
+
+      if (CurrentWord == FirstWord) {
+        unsigned FirstBit = Begin % BITWORD_SIZE;
+        Copy |= maskTrailingOnes<BitWord>(FirstBit);
+      }
+
+      if (Copy != ~0UL) {
+        unsigned Result =
+            (CurrentWord + 1) * BITWORD_SIZE - countLeadingOnes(Copy) - 1;
+        return Result < Size ? Result : -1;
+      }
+    }
     return -1;
   }
 
+  /// find_first - Returns the index of the first set bit, -1 if none
+  /// of the bits are set.
+  int find_first() const { return find_first_in(0, Size); }
+
+  /// find_last - Returns the index of the last set bit, -1 if none of the bits
+  /// are set.
+  int find_last() const { return find_last_in(0, Size); }
+
+  /// find_next - Returns the index of the next set bit following the
+  /// "Prev" bit. Returns -1 if the next set bit is not found.
+  int find_next(unsigned Prev) const { return find_first_in(Prev + 1, Size); }
+
   /// find_prev - Returns the index of the first set bit that precedes the
   /// the bit at \p PriorTo.  Returns -1 if all previous bits are unset.
-  int find_prev(unsigned PriorTo) const {
-    if (PriorTo == 0)
-      return -1;
+  int find_prev(unsigned PriorTo) const { return find_last_in(0, PriorTo); }
 
-    --PriorTo;
+  /// find_first_unset - Returns the index of the first unset bit, -1 if all
+  /// of the bits are set.
+  int find_first_unset() const { return find_first_unset_in(0, Size); }
 
-    unsigned WordPos = PriorTo / BITWORD_SIZE;
-    unsigned BitPos = PriorTo % BITWORD_SIZE;
-    BitWord Copy = Bits[WordPos];
-    // Mask off next bits.
-    Copy &= maskTrailingOnes<BitWord>(BitPos + 1);
+  /// find_next_unset - Returns the index of the next unset bit following the
+  /// "Prev" bit.  Returns -1 if all remaining bits are set.
+  int find_next_unset(unsigned Prev) const {
+    return find_first_unset_in(Prev + 1, Size);
+  }
 
-    if (Copy != 0)
-      return (WordPos + 1) * BITWORD_SIZE - countLeadingZeros(Copy) - 1;
+  /// find_last_unset - Returns the index of the last unset bit, -1 if all of
+  /// the bits are set.
+  int find_last_unset() const { return find_last_unset_in(0, Size); }
 
-    // Check previous words.
-    for (unsigned i = 1; i <= WordPos; ++i) {
-      unsigned Index = WordPos - i;
-      if (Bits[Index] == 0)
-        continue;
-      return (Index + 1) * BITWORD_SIZE - countLeadingZeros(Bits[Index]) - 1;
-    }
-    return -1;
+  /// find_prev_unset - Returns the index of the first unset bit that precedes
+  /// the bit at \p PriorTo.  Returns -1 if all previous bits are set.
+  int find_prev_unset(unsigned PriorTo) {
+    return find_last_unset_in(0, PriorTo);
   }
 
   /// clear - Removes all bits from the bitvector. Does not change capacity.
   void clear() {
     Size = 0;
   }
 
   /// resize - Grow or shrink the bitvector.

llvm/unittests/ADT/BitVectorTest.cpp

TYPED_TEST(BitVectorTest, TrivialOperation) {
   EXPECT_EQ(0U, Vec.count());
   EXPECT_EQ(0U, Vec.size());
   EXPECT_FALSE(Vec.any());
   EXPECT_TRUE(Vec.all());
   EXPECT_TRUE(Vec.none());
   EXPECT_TRUE(Vec.empty());
 }
 
-TYPED_TEST(BitVectorTest, FindOperations) {
+TYPED_TEST(BitVectorTest, SimpleFindOps) {
   // Test finding in an empty BitVector.
   TypeParam A;
   EXPECT_EQ(-1, A.find_first());
   EXPECT_EQ(-1, A.find_last());
   EXPECT_EQ(-1, A.find_first_unset());
   EXPECT_EQ(-1, A.find_last_unset());
   EXPECT_EQ(-1, A.find_next(0));
   EXPECT_EQ(-1, A.find_next_unset(0));
TYPED_TEST(BitVectorTest, SimpleFindOps) {
   EXPECT_EQ(0, A.find_first_unset());
   EXPECT_EQ(19, A.find_last_unset());
   EXPECT_EQ(5, A.find_next_unset(3));
   EXPECT_EQ(5, A.find_next_unset(4));
   EXPECT_EQ(13, A.find_next_unset(12));
   EXPECT_EQ(17, A.find_next_unset(15));
 }
 
+TEST(BitVectorTest, FindInRangeMultiWord) {
+  BitVector Vec;
+
+  Vec.resize(200);
+  Vec.set(3, 7);
+  Vec.set(24, 35);
+  Vec.set(50, 70);
+  Vec.set(150);
+  Vec.set(152);
+  Vec.set(154);
+
+  // find first
+  EXPECT_EQ(-1, Vec.find_first_in(0, 0));
+  EXPECT_EQ(-1, Vec.find_first_in(24, 24));
+  EXPECT_EQ(-1, Vec.find_first_in(7, 24));
+
+  EXPECT_EQ(3, Vec.find_first_in(0, 10));
+  EXPECT_EQ(4, Vec.find_first_in(4, 10));
+  EXPECT_EQ(150, Vec.find_first_in(100, 200));
+  EXPECT_EQ(152, Vec.find_first_in(151, 200));
+  EXPECT_EQ(154, Vec.find_first_in(153, 200));
+  EXPECT_EQ(-1, Vec.find_first_in(155, 200));
+
+  // find last
+  EXPECT_EQ(-1, Vec.find_last_in(0, 0));
+  EXPECT_EQ(-1, Vec.find_last_in(24, 24));
+  EXPECT_EQ(-1, Vec.find_last_in(7, 24));
+
+  EXPECT_EQ(6, Vec.find_last_in(0, 10));
+  EXPECT_EQ(5, Vec.find_last_in(0, 6));
+  EXPECT_EQ(154, Vec.find_last_in(100, 200));
+  EXPECT_EQ(152, Vec.find_last_in(100, 154));
+  EXPECT_EQ(150, Vec.find_last_in(100, 152));
+  EXPECT_EQ(-1, Vec.find_last_in(100, 150));
+
+  // find first unset
+  EXPECT_EQ(-1, Vec.find_first_unset_in(0, 0));
+  EXPECT_EQ(-1, Vec.find_first_unset_in(23, 23));
+  EXPECT_EQ(-1, Vec.find_first_unset_in(24, 35));
+
+  EXPECT_EQ(0, Vec.find_first_unset_in(0, 10));
+  EXPECT_EQ(1, Vec.find_first_unset_in(1, 10));
+  EXPECT_EQ(7, Vec.find_first_unset_in(5, 25));
+  EXPECT_EQ(151, Vec.find_first_unset_in(150, 200));
+  EXPECT_EQ(151, Vec.find_first_unset_in(151, 200));
+  EXPECT_EQ(153, Vec.find_first_unset_in(152, 200));
+  EXPECT_EQ(153, Vec.find_first_unset_in(153, 200));
+  EXPECT_EQ(155, Vec.find_first_unset_in(154, 200));
+  EXPECT_EQ(155, Vec.find_first_unset_in(155, 200));
+
+  // find last unset
+  EXPECT_EQ(-1, Vec.find_last_unset_in(0, 0));
+  EXPECT_EQ(-1, Vec.find_last_unset_in(23, 23));
+  EXPECT_EQ(-1, Vec.find_last_unset_in(24, 35));
+
+  EXPECT_EQ(9, Vec.find_last_unset_in(0, 10));
+  EXPECT_EQ(8, Vec.find_last_unset_in(0, 9));
+  EXPECT_EQ(2, Vec.find_last_unset_in(0, 7));
+  EXPECT_EQ(149, Vec.find_last_unset_in(100, 151));
+  EXPECT_EQ(151, Vec.find_last_unset_in(100, 152));
+  EXPECT_EQ(151, Vec.find_last_unset_in(100, 153));
+  EXPECT_EQ(153, Vec.find_last_unset_in(100, 154));
+  EXPECT_EQ(153, Vec.find_last_unset_in(100, 155));
+  EXPECT_EQ(155, Vec.find_last_unset_in(100, 156));
+}
+
+TEST(BitVectorTest, FindInRangeSingleWord) {
+  // When the bit vector contains only a single word, this is slightly different
+  // than when the bit vector contains multiple words, because masks are applied
+  // to the front and back of the same word.  So make sure this works.
+  BitVector Vec;
+
+  Vec.resize(25);
+  Vec.set(2, 4);
+  Vec.set(6, 9);
+  Vec.set(12, 15);
+  Vec.set(19);
+  Vec.set(21);
+  Vec.set(23);
+
+  // find first
+  EXPECT_EQ(-1, Vec.find_first_in(0, 0));
+  EXPECT_EQ(-1, Vec.find_first_in(24, 24));
+  EXPECT_EQ(-1, Vec.find_first_in(9, 12));
+
+  EXPECT_EQ(2, Vec.find_first_in(0, 10));
+  EXPECT_EQ(6, Vec.find_first_in(4, 10));
+  EXPECT_EQ(19, Vec.find_first_in(18, 25));
+  EXPECT_EQ(21, Vec.find_first_in(20, 25));
+  EXPECT_EQ(23, Vec.find_first_in(22, 25));
+  EXPECT_EQ(-1, Vec.find_first_in(24, 25));
+
+  // find last
+  EXPECT_EQ(-1, Vec.find_last_in(0, 0));
+  EXPECT_EQ(-1, Vec.find_last_in(24, 24));
+  EXPECT_EQ(-1, Vec.find_last_in(9, 12));
+
+  EXPECT_EQ(8, Vec.find_last_in(0, 10));
+  EXPECT_EQ(3, Vec.find_last_in(0, 6));
+  EXPECT_EQ(23, Vec.find_last_in(18, 25));
+  EXPECT_EQ(21, Vec.find_last_in(18, 23));
+  EXPECT_EQ(19, Vec.find_last_in(18, 21));
+  EXPECT_EQ(-1, Vec.find_last_in(18, 19));
+
+  // find first unset
+  EXPECT_EQ(-1, Vec.find_first_unset_in(0, 0));
+  EXPECT_EQ(-1, Vec.find_first_unset_in(23, 23));
+  EXPECT_EQ(-1, Vec.find_first_unset_in(6, 9));
+
+  EXPECT_EQ(0, Vec.find_first_unset_in(0, 6));
+  EXPECT_EQ(1, Vec.find_first_unset_in(1, 6));
+  EXPECT_EQ(9, Vec.find_first_unset_in(7, 13));
+  EXPECT_EQ(18, Vec.find_first_unset_in(18, 25));
+  EXPECT_EQ(20, Vec.find_first_unset_in(19, 25));
+  EXPECT_EQ(20, Vec.find_first_unset_in(20, 25));
+  EXPECT_EQ(22, Vec.find_first_unset_in(21, 25));
+  EXPECT_EQ(22, Vec.find_first_unset_in(22, 25));
+  EXPECT_EQ(24, Vec.find_first_unset_in(23, 25));
+  EXPECT_EQ(24, Vec.find_first_unset_in(24, 25));
+
+  // find last unset
+  EXPECT_EQ(-1, Vec.find_last_unset_in(0, 0));
+  EXPECT_EQ(-1, Vec.find_last_unset_in(23, 23));
+  EXPECT_EQ(-1, Vec.find_last_unset_in(6, 9));
+
+  EXPECT_EQ(5, Vec.find_last_unset_in(0, 6));
+  EXPECT_EQ(4, Vec.find_last_unset_in(0, 5));
+  EXPECT_EQ(1, Vec.find_last_unset_in(0, 4));
+  EXPECT_EQ(11, Vec.find_last_unset_in(7, 13));
+  EXPECT_EQ(24, Vec.find_last_unset_in(18, 25));
+  EXPECT_EQ(22, Vec.find_last_unset_in(18, 24));
+  EXPECT_EQ(22, Vec.find_last_unset_in(18, 23));
+  EXPECT_EQ(20, Vec.find_last_unset_in(18, 22));
+  EXPECT_EQ(20, Vec.find_last_unset_in(18, 21));
+  EXPECT_EQ(18, Vec.find_last_unset_in(18, 20));
+  EXPECT_EQ(18, Vec.find_last_unset_in(18, 19));
+}
+
 TYPED_TEST(BitVectorTest, CompoundAssignment) {
   TypeParam A;
   A.resize(10);
   A.set(4);
   A.set(7);
 
   TypeParam B;
   B.resize(50);