Merge pull request #29405 from vespa-engine/toregge/add-step-backward-to-btree-iterator

Add step_backward to btree iterator.

3 files changed, 151 insertions, 27 deletions

diff --git a/vespalib/src/tests/btree/btree_test.cpp b/vespalib/src/tests/btree/btree_test.cpp
index 21838676906..30e3df7dd14 100644
--- a/vespalib/src/tests/btree/btree_test.cpp
+++ b/vespalib/src/tests/btree/btree_test.cpp
@@ -232,6 +232,7 @@ protected:
     void requireThatUpperBoundWorksT();
     void requireThatIteratorDistanceWorks(int numEntries);
     void test_step_forward(int num_entries);
+    void test_step_backward(int num_entries);
 };
 
 template <typename LeafNodeType>
@@ -1476,8 +1477,12 @@ BTreeTest::requireThatIteratorDistanceWorks(int numEntries)
             iitbs.binarySeek(i);
             ++it;
         }
-        iitlsp.linearSeekPast(i);
-        iitbsp.binarySeekPast(i);
+        if (iitlsp.valid()) {
+            iitlsp.linearSeekPast(i);
+        }
+        if (iitbsp.valid()) {
+            iitbsp.binarySeekPast(i);
+        }
         Iterator iitlsp2 = iitls;
         Iterator iitbsp2 = iitbs;
         Iterator iitnr = i < numEntries ? iitn : tree.begin();
@@ -1548,8 +1553,35 @@ BTreeTest::test_step_forward(int num_entries)
     }
 }
 
+void
+BTreeTest::test_step_backward(int num_entries)
+{
+    GenerationHandler g;
+    MyTree tree;
+    for (int i = 0; i < num_entries; ++i) {
+        tree.insert(i, toStr(i));
+    }
+    auto it = tree.begin();
+    for (int i = 0; i <= num_entries; ++i) {
+        auto iit = tree.lowerBound(i);
+        auto iit2 = iit;
+        iit2 -= i;
+        EXPECT_TRUE(iit2.identical(it));
+        iit2 = iit;
+        iit2 -= (1000000 + i);
+        EXPECT_TRUE(iit2.identical(it));
+        for (int j = 0; j <= i; ++j) {
+            auto jit = tree.lowerBound(j);
+            auto iit3 = iit;
+            iit3 -= (i - j);
+            EXPECT_TRUE(iit3.identical(jit));
+        }
+    }
+}
+
 TEST_F(BTreeTest, require_that_iterator_distance_works)
 {
+    requireThatIteratorDistanceWorks(0);
     requireThatIteratorDistanceWorks(1);
     requireThatIteratorDistanceWorks(3);
     requireThatIteratorDistanceWorks(8);
@@ -1560,6 +1592,7 @@ TEST_F(BTreeTest, require_that_iterator_distance_works)
 
 TEST_F(BTreeTest, require_that_step_forward_works)
 {
+    test_step_forward(0);
     test_step_forward(1);
     test_step_forward(3);
     test_step_forward(8);
@@ -1568,6 +1601,17 @@ TEST_F(BTreeTest, require_that_step_forward_works)
     test_step_forward(400);
 }
 
+TEST_F(BTreeTest, require_that_step_backward_works)
+{
+    test_step_backward(0);
+    test_step_backward(1);
+    test_step_backward(3);
+    test_step_backward(8);
+    test_step_backward(20);
+    test_step_backward(100);
+    test_step_backward(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 9519951f9e2..cd63499a5ed 100644
--- a/vespalib/src/vespa/vespalib/btree/btreeiterator.h
+++ b/vespalib/src/vespa/vespalib/btree/btreeiterator.h
@@ -232,11 +232,18 @@ protected:
         return *this;
     }
 
+    void set_subtree_position(const InternalNodeType* node, uint32_t level, uint32_t idx, size_t position);
+
     /*
      * Step iterator forwards the given number of steps.
      */
     void step_forward(size_t steps);
 
+    /*
+     * Step iterator backwards the given number of steps.
+     */
+    void step_backward(size_t steps);
+
     ~BTreeIteratorBase();
     BTreeIteratorBase(const BTreeIteratorBase &other);
     BTreeIteratorBase &operator=(const BTreeIteratorBase &other);
@@ -507,6 +514,7 @@ protected:
     using ParentType::_compatLeafNode;
     using ParentType::clearPath;
     using ParentType::setupEmpty;
+    using ParentType::step_backward;
     using ParentType::step_forward;
 public:
     using ParentType::end;
@@ -570,6 +578,14 @@ public:
         step_forward(steps);
         return *this;
     }
+
+    /*
+     * Step iterator backward the given number of steps.
+     */
+    BTreeConstIterator & operator-=(size_t steps) {
+        step_backward(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
@@ -712,6 +728,7 @@ public:
     using ParentType::_leafRoot;
     using ParentType::_compatLeafNode;
     using ParentType::end;
+    using ParentType::step_backward;
     using ParentType::step_forward;
     using EntryRef = datastore::EntryRef;
 
@@ -746,6 +763,11 @@ public:
         return *this;
     }
 
+    BTreeIterator & operator-=(size_t steps) {
+        step_backward(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 3119d05cfd9..b9afce54f6b 100644
--- a/vespalib/src/vespa/vespalib/btree/btreeiterator.hpp
+++ b/vespalib/src/vespa/vespalib/btree/btreeiterator.hpp
@@ -530,6 +530,38 @@ 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>::
+set_subtree_position(const InternalNodeType* node, uint32_t level, uint32_t idx, size_t position)
+{
+    /*
+     * Walk down subtree adjusting iterator for new partial position.
+     */
+    _path[level].setIdx(idx);
+    size_t remaining_steps = position;
+    while (level > 0) {
+        --level;
+        node = _allocator->mapInternalRef(node->getChild(idx));
+        assert(remaining_steps < node->validLeaves());
+        idx = 0;
+        while (idx < node->validSlots()) {
+            auto valid_leaves = _allocator->validLeaves(node->getChild(idx));
+            if (remaining_steps < valid_leaves) {
+                break;
+            }
+            remaining_steps -= valid_leaves;
+            ++idx;
+        }
+        assert(idx < node->validSlots());
+        _path[level].setNodeAndIdx(node, idx);
+    }
+    auto lnode = _allocator->mapLeafRef(node->getChild(idx));
+    assert(remaining_steps < lnode->validSlots());
+    _leaf.setNodeAndIdx(lnode, remaining_steps);
+}
+
+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();
@@ -557,10 +589,7 @@ step_forward(size_t steps)
         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();
+            auto valid_leaves = _allocator->validLeaves(node->getChild(idx));
             if (remaining_steps < valid_leaves) {
                 break;
             }
@@ -577,32 +606,61 @@ step_forward(size_t steps)
             }
         }
     }
+    set_subtree_position(node, level, idx, remaining_steps);
+}
+
+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_backward(size_t steps)
+{
+    int64_t remaining_steps = steps;
+    if (remaining_steps == 0) {
+        return;
+    }
+    if (_leaf.getNode() == nullptr) {
+        rbegin();
+        if (_leaf.getNode() == nullptr) {
+            return;
+        }
+        --remaining_steps;
+    }
+    auto idx = _leaf.getIdx();
+    if (idx >= remaining_steps) {
+        _leaf.setIdx(idx - remaining_steps);
+        return;
+    }
+    if (_pathSize == 0) {
+        _leaf.setIdx(0);
+        return;
+    }
+    remaining_steps -= idx;
+    uint32_t level = 0;
+    uint32_t levels = _pathSize;
+    const InternalNodeType* node;
     /*
-     * Walk down subtree adjusting iterator for new position.
+     * Find intermediate node representing subtree containing old and 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;
+    for (;;) {
+        node = _path[level].getNode();
+        idx = _path[level].getIdx();
+        while (idx > 0 && remaining_steps > 0) {
+            --idx;
+            remaining_steps -= _allocator->validLeaves(node->getChild(idx));
+        }
+        if (remaining_steps <= 0) {
+            break;
+        } else {
+            ++level;
+            if (level == levels) {
+                begin();
+                return;
             }
-            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);
+    set_subtree_position(node, level, idx, -remaining_steps);
 }
 
 template <typename KeyT, typename DataT, typename AggrT, typename CompareT,