[ADT] CoalescingBitVector: Add advanceToLowerBound iterator operation

advanceToLowerBound moves an iterator to the first bit set at, or after,
the given index. This can be faster than doing IntervalMap::find.

rdar://60046261

Differential Revision: https://reviews.llvm.org/D76466

2 files changed, 61 insertions, 2 deletions

diff --git a/llvm/include/llvm/ADT/CoalescingBitVector.h b/llvm/include/llvm/ADT/CoalescingBitVector.h
index 34352da..8ef6f4f 100644
--- a/llvm/include/llvm/ADT/CoalescingBitVector.h
+++ b/llvm/include/llvm/ADT/CoalescingBitVector.h
@@ -266,9 +266,9 @@ public:
     }
 
     /// 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 @@ public:
       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 @@ public:
 
   /// 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
index 9536efe..06ac5df 100644
--- a/llvm/unittests/ADT/CoalescingBitVectorTest.cpp
+++ b/llvm/unittests/ADT/CoalescingBitVectorTest.cpp
@@ -453,6 +453,44 @@ TEST(CoalescingBitVectorTest, FindLowerBound) {
   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;
   {