Merge pull request #13563 from vespa-engine/toregge/add-foreach-key-range-method-to-btree-iterator

Add foreach_key_range method to btree iterator, to scan a range of

3 files changed, 213 insertions, 0 deletions

diff --git a/vespalib/src/tests/btree/btree_test.cpp b/vespalib/src/tests/btree/btree_test.cpp
index 848c8a37125..63afd8b770f 100644
--- a/vespalib/src/tests/btree/btree_test.cpp
+++ b/vespalib/src/tests/btree/btree_test.cpp
@@ -36,6 +36,54 @@ toStr(const T & v)
     return ss.str();
 }
 
+class SequenceValidator
+{
+    int _wanted_count;
+    int _prev_key;
+    int _count;
+    bool _failed;
+
+public:
+    SequenceValidator(int start, int wanted_count)
+        : _wanted_count(wanted_count),
+          _prev_key(start - 1),
+          _count(0),
+          _failed(false)
+    {
+    }
+
+    bool failed() const {
+        return _failed || _wanted_count != _count;
+    }
+    
+    void operator()(int key) {
+        if (key != _prev_key + 1) {
+            _failed = true;
+        }
+        _prev_key = key;
+        ++_count;
+    }
+};
+
+class ForeachKeyValidator
+{
+    SequenceValidator & _validator;
+public:
+    ForeachKeyValidator(SequenceValidator &validator)
+        : _validator(validator)
+    {
+    }
+    void operator()(int key) {
+        _validator(key);
+    }
+};
+
+template <typename Iterator>
+void validate_subrange(Iterator &start, Iterator &end, SequenceValidator &validator) {
+    start.foreach_key_range(end, ForeachKeyValidator(validator));
+    EXPECT_FALSE(validator.failed());
+}
+
 }
 
 typedef BTreeTraits<4, 4, 31, false> MyTraits;
@@ -210,6 +258,8 @@ private:
 
     void
     requireThatIteratorDistanceWorks();
+
+    void requireThatForeachKeyWorks();
 public:
     int Main() override;
 };
@@ -1489,6 +1539,32 @@ Test::requireThatIteratorDistanceWorks()
     requireThatIteratorDistanceWorks(400);
 }
 
+void
+Test::requireThatForeachKeyWorks()
+{
+    using Tree = BTree<int, int, btree::NoAggregated, MyComp, MyTraits>;
+    using Iterator = typename Tree::ConstIterator;
+    Tree t;
+    populateTree(t, 256, 1);
+
+    {
+        // Whole range
+        SequenceValidator validator(1, 256);
+        t.foreach_key(ForeachKeyValidator(validator));
+        EXPECT_FALSE(validator.failed());
+    }
+    {
+        // Subranges
+        for (int startval = 1; startval < 259; ++startval) {
+            for (int endval = 1; endval < 259; ++endval) {
+                SequenceValidator validator(startval, std::max(0, std::min(endval,257) - std::min(startval, 257)));
+                Iterator start = t.lowerBound(startval);
+                Iterator end = t.lowerBound(endval);
+                validate_subrange(start, end, validator);
+            }
+        }
+    }
+};
 
 int
 Test::Main()
@@ -1515,6 +1591,7 @@ Test::Main()
     requireThatSmallNodesWorks();
     requireThatApplyWorks();
     requireThatIteratorDistanceWorks();
+    requireThatForeachKeyWorks();
 
     TEST_DONE();
 }
diff --git a/vespalib/src/vespa/vespalib/btree/btreeiterator.h b/vespalib/src/vespa/vespalib/btree/btreeiterator.h
index 55ab37759ad..6933fc1c2d0 100644
--- a/vespalib/src/vespa/vespalib/btree/btreeiterator.h
+++ b/vespalib/src/vespa/vespalib/btree/btreeiterator.h
@@ -303,6 +303,47 @@ protected:
      * @param pathSize     New tree height (number of levels of internal nodes)
      */
     VESPA_DLL_LOCAL void clearPath(uint32_t pathSize);
+
+    /**
+     * Call func with leaf entry key value as argument for all leaf entries in subtree
+     * from this iterator position to end of subtree.
+     */
+    template <typename FunctionType>
+    void
+    foreach_key_range_start(uint32_t level, FunctionType func) const
+    {
+        if (level > 0u) {
+            --level;
+            foreach_key_range_start(level, func);
+            auto &store = _allocator->getNodeStore();
+            auto node = _path[level].getNode();
+            uint32_t idx = _path[level].getIdx();
+            node->foreach_key_range(store, idx + 1, node->validSlots(), func);
+        } else {
+            _leaf.getNode()->foreach_key_range(_leaf.getIdx(), _leaf.getNode()->validSlots(), func);
+        }
+    }
+
+    /**
+     * Call func with leaf entry key value as argument for all leaf entries in subtree
+     * from start of subtree until this iterator position is reached (i.e. entries in
+     * subtree before this iterator position).
+     */
+    template <typename FunctionType>
+    void
+    foreach_key_range_end(uint32_t level, FunctionType func) const
+    {
+        if (level > 0u) {
+            --level;
+            auto &store = _allocator->getNodeStore();
+            auto node = _path[level].getNode();
+            uint32_t eidx = _path[level].getIdx();
+            node->foreach_key_range(store, 0, eidx, func);
+            foreach_key_range_end(level, func);
+        } else {
+            _leaf.getNode()->foreach_key_range(0, _leaf.getIdx(), func);
+        }
+    }
 public:
 
     bool
@@ -451,6 +492,68 @@ public:
             _leafRoot->foreach_key(func);
         }
     }
+
+    /**
+     * Call func with leaf entry key value as argument for all leaf entries in tree from
+     * this iterator position until end_itr position is reached (i.e. entries in
+     * range [this iterator, end_itr)).
+     */
+    template <typename FunctionType>
+    void
+    foreach_key_range(const BTreeIteratorBase &end_itr, FunctionType func) const
+    {
+        if (!valid()) {
+            return;
+        }
+        if (!end_itr.valid()) {
+            foreach_key_range_start(_pathSize, func);
+            return;
+        }
+        assert(_pathSize == end_itr._pathSize);
+        assert(_allocator == end_itr._allocator);
+        uint32_t level = _pathSize;
+        if (level > 0u) {
+            /**
+             * Tree has intermediate nodes. Detect lowest shared tree node for this
+             * iterator and end_itr.
+             */
+            uint32_t idx;
+            uint32_t eidx;
+            do {
+                --level;
+                assert(_path[level].getNode() == end_itr._path[level].getNode());
+                idx = _path[level].getIdx();
+                eidx = end_itr._path[level].getIdx();
+                if (idx > eidx) {
+                    return;
+                }
+                if (idx != eidx) {
+                    ++level;
+                    break;
+                }
+            } while (level != 0);
+            if (level > 0u) {
+                // Lowest shared node is an intermediate node.
+                // Left subtree for child [idx], from this iterator position to end of subtree.
+                foreach_key_range_start(level - 1, func);
+                auto &store = _allocator->getNodeStore();
+                auto node = _path[level - 1].getNode();
+                // Any intermediate subtrees for children [idx + 1, eidx).
+                node->foreach_key_range(store, idx + 1, eidx, func);
+                // Right subtree for child [eidx], from start of subtree to end_itr position.
+                end_itr.foreach_key_range_end(level - 1, func);
+                return;
+            } else {
+                // Lowest shared node is a leaf node.
+                assert(_leaf.getNode() == end_itr._leaf.getNode());
+            }
+        }
+        uint32_t idx = _leaf.getIdx();
+        uint32_t eidx = end_itr._leaf.getIdx();
+        if (idx < eidx) {
+            _leaf.getNode()->foreach_key_range(idx, eidx, func);
+        }
+    }
 };
 
 
diff --git a/vespalib/src/vespa/vespalib/btree/btreenode.h b/vespalib/src/vespa/vespalib/btree/btreenode.h
index b34be33ccf5..0c70e70bc6a 100644
--- a/vespalib/src/vespa/vespalib/btree/btreenode.h
+++ b/vespalib/src/vespa/vespalib/btree/btreenode.h
@@ -370,6 +370,26 @@ public:
         }
     }
 
+    /**
+     * Call func with leaf entry key value as argument for all leaf entries in subtrees
+     * for children [start_idx, end_idx).
+     */
+    template <typename NodeStoreType, typename FunctionType>
+    void foreach_key_range(NodeStoreType &store, uint32_t start_idx, uint32_t end_idx, FunctionType func) const {
+        const BTreeNode::Ref *it = this->_data;
+        const BTreeNode::Ref *ite = it + end_idx;
+        it += start_idx;
+        if (this->getLevel() > 1u) {
+            for (; it != ite; ++it) {
+                store.mapInternalRef(*it)->foreach_key(store, func);
+            }
+        } else {
+            for (; it != ite; ++it) {
+                store.mapLeafRef(*it)->foreach_key(func);
+            }
+        }
+    }
+
     template <typename NodeStoreType, typename FunctionType>
     void foreach(NodeStoreType &store, FunctionType func) const {
         const BTreeNode::Ref *it = this->_data;
@@ -459,6 +479,19 @@ public:
         }
     }
 
+    /**
+     * Call func with leaf entry key value as argument for leaf entries [start_idx, end_idx).
+     */
+    template <typename FunctionType>
+    void foreach_key_range(uint32_t start_idx, uint32_t end_idx, FunctionType func) const {
+        const KeyT *it = _keys;
+        const KeyT *ite = it + end_idx;
+        it += start_idx;
+        for (; it != ite; ++it) {
+            func(*it);
+        }
+    }
+
     template <typename FunctionType>
     void foreach(FunctionType func) const {
         const KeyT *it = _keys;