diff --git a/llvm/include/llvm/ADT/CoalescingBitVector.h b/llvm/include/llvm/ADT/CoalescingBitVector.h --- a/llvm/include/llvm/ADT/CoalescingBitVector.h +++ b/llvm/include/llvm/ADT/CoalescingBitVector.h @@ -266,9 +266,9 @@ } /// Advance the iterator to \p Index, if it is contained within the current - /// interval. + /// interval. The public-facing method which supports advancing past the + /// current interval is \ref advanceToLowerBound. void advanceTo(IndexT Index) { - assert(OffsetIntoMapIterator == 0 && "Not implemented"); assert(Index <= CachedStop && "Cannot advance to OOB index"); if (Index < CachedStart) // We're already past this index. @@ -314,6 +314,25 @@ operator++(); return tmp; } + + /// Advance the iterator to the first set bit AT, OR AFTER, \p Index. If + /// no such set bit exists, advance to end(). This is like std::lower_bound. + /// This is useful if \p Index is close to the current iterator position. + /// However, unlike \ref find(), this has worst-case O(n) performance. + void advanceToLowerBound(IndexT Index) { + if (OffsetIntoMapIterator == kIteratorAtTheEndOffset) + return; + + // Advance to the first interval containing (or past) Index, or to end(). + while (Index > CachedStop) { + ++MapIterator; + resetCache(); + if (OffsetIntoMapIterator == kIteratorAtTheEndOffset) + return; + } + + advanceTo(Index); + } }; const_iterator begin() const { return const_iterator(Intervals.begin()); } @@ -322,6 +341,8 @@ /// Return an iterator pointing to the first set bit AT, OR AFTER, \p Index. /// If no such set bit exists, return end(). This is like std::lower_bound. + /// This has worst-case logarithmic performance (roughly O(log(gaps between + /// contiguous ranges))). const_iterator find(IndexT Index) const { auto UnderlyingIt = Intervals.find(Index); if (UnderlyingIt == Intervals.end()) diff --git a/llvm/unittests/ADT/CoalescingBitVectorTest.cpp b/llvm/unittests/ADT/CoalescingBitVectorTest.cpp --- a/llvm/unittests/ADT/CoalescingBitVectorTest.cpp +++ b/llvm/unittests/ADT/CoalescingBitVectorTest.cpp @@ -453,6 +453,44 @@ EXPECT_EQ(*BV.find(3), 3u); } +TEST(CoalescingBitVectorTest, AdvanceToLowerBound) { + U64BitVec::Allocator Alloc; + U64BitVec BV(Alloc); + uint64_t BigNum1 = uint64_t(1) << 32; + uint64_t BigNum2 = (uint64_t(1) << 33) + 1; + + auto advFromBegin = [&](uint64_t To) -> U64BitVec::const_iterator { + auto It = BV.begin(); + It.advanceToLowerBound(To); + return It; + }; + + EXPECT_TRUE(advFromBegin(BigNum1) == BV.end()); + BV.set(BigNum1); + auto Find1 = advFromBegin(BigNum1); + EXPECT_EQ(*Find1, BigNum1); + BV.set(BigNum2); + auto Find2 = advFromBegin(BigNum1); + EXPECT_EQ(*Find2, BigNum1); + auto Find3 = advFromBegin(BigNum2); + EXPECT_EQ(*Find3, BigNum2); + BV.reset(BigNum1); + auto Find4 = advFromBegin(BigNum1); + EXPECT_EQ(*Find4, BigNum2); + + BV.clear(); + BV.set({1, 2, 3}); + EXPECT_EQ(*advFromBegin(2), 2u); + EXPECT_EQ(*advFromBegin(3), 3u); + auto It = BV.begin(); + It.advanceToLowerBound(0); + EXPECT_EQ(*It, 1u); + It.advanceToLowerBound(100); + EXPECT_TRUE(It == BV.end()); + It.advanceToLowerBound(100); + EXPECT_TRUE(It == BV.end()); +} + TEST(CoalescingBitVectorTest, Print) { std::string S; {