From 5eb064ea696cffca0a8ae5e9a3a4db1cc616d583 Mon Sep 17 00:00:00 2001 From: Tor Egge Date: Mon, 20 Nov 2023 16:16:48 +0100 Subject: Add step_forward to btree iterator. --- vespalib/src/tests/btree/btree_test.cpp | 38 +++++++++++ vespalib/src/vespa/vespalib/btree/btreeiterator.h | 19 ++++++ .../src/vespa/vespalib/btree/btreeiterator.hpp | 79 ++++++++++++++++++++++ 3 files changed, 136 insertions(+) (limited to 'vespalib/src') 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 void requireThatUpperBoundWorksT(); void requireThatIteratorDistanceWorks(int numEntries); + void test_step_forward(int num_entries); }; template @@ -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; 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(*_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 +void +BTreeIteratorBase:: +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 void -- cgit v1.2.3