[ADT] CoalescingBitVector: Add advanceToLowerBound iterator operation

[vsk]

Summary

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

[vsk]

Allow advanceToLowerBound to be called on end(), as this happens in practice.

[aprantl]

llvm/include/llvm/ADT/CoalescingBitVector.h

321

The doxygen comment of find() doesn't actually say what its complexity is. Might be nice to add that information.

[aprantl]

llvm/include/llvm/ADT/CoalescingBitVector.h

321

It would also be good to explain when to use this function over advance_to().

[jmorse]

LGTM otherwise

[vsk]

Add additional documentation, as suggested.

llvm/include/llvm/ADT/CoalescingBitVector.h

void resetCache() {
         CachedStart = MapIterator.start();
         CachedStop = MapIterator.stop();
       } else {
         setToEnd();
       }
     }
 
     /// 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.
         return;
       OffsetIntoMapIterator = Index - CachedStart;
     }
 
     const_iterator(UnderlyingIterator MapIt) : MapIterator(MapIt) {
const_iterator &operator++() { // Pre-increment (++It).
       return *this;
     }
 
     const_iterator operator++(int) { // Post-increment (It++).
       const_iterator tmp = *this;
       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()); }
 
   const_iterator end() const { return const_iterator(); }
 
   /// 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())
       return end();
     auto It = const_iterator(UnderlyingIt);
     It.advanceTo(Index);
     return It;
   }

llvm/unittests/ADT/CoalescingBitVectorTest.cpp

TEST(CoalescingBitVectorTest, FindLowerBound) {
   EXPECT_EQ(*Find4, BigNum2);
 
   BV.clear();
   BV.set({1, 2, 3});
   EXPECT_EQ(*BV.find(2), 2u);
   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;
   {
     raw_string_ostream OS(S);
     UBitVec::Allocator Alloc;
     UBitVec BV(Alloc);
     BV.set({1});
     BV.print(OS);