Merge pull request #29391 from vespa-engine/toregge/add-step-forward-to-btree-iterator

Add step_forward to btree iterator.

3 files changed, 136 insertions, 0 deletions

diff --git a/vespalib/src/tests/btree/btree_test.cpp b/vespalib/src/tests/btree/btree_test.cpp
index f6fa962fa9d..21838676906 100644
--- a/vespalib/src/tests/btree/btree_test.cpp
+++ b/vespalib/src/tests/btree/btree_test.cpp
@@ -231,6 +231,7 @@ protected:
     template <typename TreeType>
     void requireThatUpperBoundWorksT();
     void requireThatIteratorDistanceWorks(int numEntries);
+    void test_step_forward(int num_entries);
 };
 
 template <typename LeafNodeType>
@@ -1519,6 +1520,33 @@ BTreeTest::requireThatIteratorDistanceWorks(int numEntries)
     }
 }
 
+void
+BTreeTest::test_step_forward(int num_entries)
+{
+    GenerationHandler g;
+    MyTree tree;
+    for (int i = 0; i < num_entries; ++i) {
+        tree.insert(i, toStr(i));
+    }
+    auto it = tree.begin();
+    auto ite = it;
+    ite.end();
+    for (int i = 0; i <= num_entries; ++i) {
+        auto iit = tree.lowerBound(i);
+        auto iit2 = iit;
+        iit2 += (num_entries - i);
+        EXPECT_TRUE(iit2.identical(ite));
+        iit2 = iit;
+        iit2 += (1000000 + num_entries);
+        EXPECT_TRUE(iit2.identical(ite));
+        for (int j = i; j <= num_entries; ++j) {
+            auto jit = tree.lowerBound(j);
+            auto iit3 = iit;
+            iit3 += (j - i);
+            EXPECT_TRUE(iit3.identical(jit));
+        }
+    }
+}
 
 TEST_F(BTreeTest, require_that_iterator_distance_works)
 {
@@ -1530,6 +1558,16 @@ TEST_F(BTreeTest, require_that_iterator_distance_works)
     requireThatIteratorDistanceWorks(400);
 }
 
+TEST_F(BTreeTest, require_that_step_forward_works)
+{
+    test_step_forward(1);
+    test_step_forward(3);
+    test_step_forward(8);
+    test_step_forward(20);
+    test_step_forward(100);
+    test_step_forward(400);
+}
+
 TEST_F(BTreeTest, require_that_foreach_key_works)
 {
     using Tree = BTree<int, int, btree::NoAggregated, MyComp, MyTraits>;
diff --git a/vespalib/src/vespa/vespalib/btree/btreeiterator.h b/vespalib/src/vespa/vespalib/btree/btreeiterator.h
index 53e4d004981..9519951f9e2 100644
--- a/vespalib/src/vespa/vespalib/btree/btreeiterator.h
+++ b/vespalib/src/vespa/vespalib/btree/btreeiterator.h
@@ -232,6 +232,11 @@ protected:
         return *this;
     }
 
+    /*
+     * Step iterator forwards the given number of steps.
+     */
+    void step_forward(size_t steps);
+
     ~BTreeIteratorBase();
     BTreeIteratorBase(const BTreeIteratorBase &other);
     BTreeIteratorBase &operator=(const BTreeIteratorBase &other);
@@ -502,6 +507,7 @@ protected:
     using ParentType::_compatLeafNode;
     using ParentType::clearPath;
     using ParentType::setupEmpty;
+    using ParentType::step_forward;
 public:
     using ParentType::end;
 
@@ -557,6 +563,13 @@ public:
         return *this;
     }
 
+    /*
+     * Step iterator forwards the given number of steps.
+     */
+    BTreeConstIterator & operator+=(size_t steps) {
+        step_forward(steps);
+        return *this;
+    }
     /**
      * Position iterator at first position with a key that is greater
      * than or equal to the key argument.  The iterator must be set up
@@ -699,6 +712,7 @@ public:
     using ParentType::_leafRoot;
     using ParentType::_compatLeafNode;
     using ParentType::end;
+    using ParentType::step_forward;
     using EntryRef = datastore::EntryRef;
 
     BTreeIterator(BTreeNode::Ref root, const NodeAllocatorType &allocator)
@@ -727,6 +741,11 @@ public:
         return *this;
     }
 
+    BTreeIterator & operator+=(size_t steps) {
+        step_forward(steps);
+        return *this;
+    }
+
     NodeAllocatorType & getAllocator() const {
         return const_cast<NodeAllocatorType &>(*_allocator);
     }
diff --git a/vespalib/src/vespa/vespalib/btree/btreeiterator.hpp b/vespalib/src/vespa/vespalib/btree/btreeiterator.hpp
index fdab3a94a5b..3119d05cfd9 100644
--- a/vespalib/src/vespa/vespalib/btree/btreeiterator.hpp
+++ b/vespalib/src/vespa/vespalib/btree/btreeiterator.hpp
@@ -526,6 +526,85 @@ identical(const BTreeIteratorBase &rhs) const
 }
 
 
+template <typename KeyT, typename DataT, typename AggrT,
+          uint32_t INTERNAL_SLOTS, uint32_t LEAF_SLOTS, uint32_t PATH_SIZE>
+void
+BTreeIteratorBase<KeyT, DataT, AggrT, INTERNAL_SLOTS, LEAF_SLOTS, PATH_SIZE>::
+step_forward(size_t steps)
+{
+    auto lnode = _leaf.getNode();
+    if (lnode == nullptr) {
+        return;
+    }
+    auto idx = _leaf.getIdx();
+    if (idx + steps < lnode->validSlots()) {
+        _leaf.setIdx(idx + steps);
+        return;
+    }
+    if (_pathSize == 0) {
+        _leaf.invalidate();
+        return;
+    }
+    size_t remaining_steps = steps - (lnode->validSlots() - idx);
+    uint32_t level = 0;
+    uint32_t levels = _pathSize;
+    const InternalNodeType* node;
+    /*
+     * Find intermediate node representing subtree containing old and new
+     * position.
+     */
+    for (;;) {
+        node = _path[level].getNode();
+        idx = _path[level].getIdx() + 1;
+        while (idx < node->validSlots()) {
+            auto ref = node->getChild(idx);
+            auto valid_leaves = (level != 0) ?
+                                _allocator->mapInternalRef(ref)->validLeaves() :
+                                _allocator->mapLeafRef(ref)->validLeaves();
+            if (remaining_steps < valid_leaves) {
+                break;
+            }
+            remaining_steps -= valid_leaves;
+            ++idx;
+        }
+        if (idx < node->validSlots()) {
+            break;
+        } else {
+            ++level;
+            if (level == levels) {
+                end();
+                return;
+            }
+        }
+    }
+    /*
+     * Walk down subtree adjusting iterator for new position.
+     */
+    _path[level].setIdx(idx);
+    while (level > 0) {
+        --level;
+        node = _allocator->mapInternalRef(node->getChild(idx));
+        assert(remaining_steps < node->validLeaves());
+        idx = 0;
+        while (idx < node->validSlots()) {
+            auto ref = node->getChild(idx);
+            auto valid_leaves = (level != 0) ?
+                                _allocator->mapInternalRef(ref)->validLeaves() :
+                                _allocator->mapLeafRef(ref)->validLeaves();
+            if (remaining_steps < valid_leaves) {
+                break;
+            }
+            remaining_steps -= valid_leaves;
+            ++idx;
+        }
+        assert(idx < node->validSlots());
+        _path[level].setNodeAndIdx(node, idx);
+    }
+    lnode = _allocator->mapLeafRef(node->getChild(idx));
+    assert(remaining_steps < lnode->validSlots());
+    _leaf.setNodeAndIdx(lnode, remaining_steps);
+}
+
 template <typename KeyT, typename DataT, typename AggrT, typename CompareT,
           typename TraitsT>
 void