geometry_rtree/
query_iter.rs

1//! Lazy spatial-query iteration — the walk behind
2//! [`Rtree::query_iter`](crate::rtree::Rtree::query_iter).
3//!
4//! The pruning walk of `visitors/spatial_query.hpp`, unrolled from
5//! recursion onto an explicit node stack so it can pause between
6//! yields: a consumer that stops early performs no traversal past the
7//! value it stopped at. [`Rtree::query`](crate::rtree::Rtree::query)
8//! collects this iterator, so the crate has exactly one query walk.
9
10use alloc::vec::Vec;
11use core::iter::FusedIterator;
12
13use crate::indexable::Indexable;
14use crate::node::Node;
15use crate::predicate::Predicate;
16
17/// How the cursor drains the current leaf: `Filtered` tests each value
18/// with [`Predicate::matches`]; `DumpAll` yields every value of a leaf
19/// whose subtree box the predicate [`covers_all`](Predicate::covers_all)
20/// — each one is a guaranteed match.
21enum LeafMode {
22    Filtered,
23    DumpAll,
24}
25
26/// A lazy iterator over the values whose bounds satisfy a
27/// [`Predicate`], in depth-first tree order.
28///
29/// Created by [`Rtree::query_iter`](crate::rtree::Rtree::query_iter).
30/// Holds the unvisited subtrees on an explicit stack plus a cursor into
31/// the current leaf; each [`next`](Iterator::next) drains the cursor,
32/// then pops the stack — a popped leaf installs a new cursor, a popped
33/// branch pushes the children that pass
34/// [`Predicate::could_match`] (reversed, so pops visit them
35/// first-to-last, the recursive walk's order). Each stack entry carries
36/// a `covered` flag — set once `Predicate::covers_all` holds for a
37/// subtree's box — under which children are pushed and leaves dumped
38/// without further predicate tests, still one value per `next` call.
39/// One `Vec` allocation per iterator, none per element.
40pub struct QueryIter<'a, T> {
41    predicate: Predicate,
42    stack: Vec<(&'a Node<T>, bool)>,
43    leaf: core::slice::Iter<'a, T>,
44    leaf_mode: LeafMode,
45}
46
47impl<'a, T> QueryIter<'a, T> {
48    /// `max_fanout` is the split strategy's branch maximum
49    /// (`Params::BRANCH_MAX`): the
50    /// stack's worst case is every level's unvisited siblings pending
51    /// at once — the root pop pushes up to `max_fanout` entries, each
52    /// deeper branch pop nets up to `max_fanout − 1` more, peaking at
53    /// `(height − 1)·(max_fanout − 1) + 1` — so sizing to `height`
54    /// alone under-allocates and reallocates on descent (a
55    /// tree-size-dependent number of times, since `height` itself grows
56    /// with tree size); the capacity below over-provisions that peak by
57    /// one level of slack and keeps the one allocation constant across
58    /// tree sizes at a fixed fanout (lazy R4).
59    pub(crate) fn new(
60        root: &'a Node<T>,
61        predicate: Predicate,
62        height: usize,
63        max_fanout: usize,
64    ) -> Self {
65        let capacity = height.saturating_sub(1) * max_fanout.saturating_sub(1) + max_fanout;
66        let mut stack = Vec::with_capacity(capacity);
67        stack.push((root, false));
68        Self {
69            predicate,
70            stack,
71            leaf: [].iter(),
72            leaf_mode: LeafMode::Filtered,
73        }
74    }
75}
76
77impl<'a, T: Indexable> Iterator for QueryIter<'a, T> {
78    type Item = &'a T;
79
80    fn next(&mut self) -> Option<&'a T> {
81        loop {
82            match self.leaf_mode {
83                LeafMode::Filtered => {
84                    for value in self.leaf.by_ref() {
85                        if self.predicate.matches(&value.bounds()) {
86                            return Some(value);
87                        }
88                    }
89                }
90                LeafMode::DumpAll => {
91                    if let Some(value) = self.leaf.next() {
92                        return Some(value);
93                    }
94                }
95            }
96            let (node, covered) = self.stack.pop()?;
97            match node {
98                Node::Leaf(values) => {
99                    self.leaf = values.iter();
100                    self.leaf_mode = if covered {
101                        LeafMode::DumpAll
102                    } else {
103                        LeafMode::Filtered
104                    };
105                }
106                Node::Branch(children) => {
107                    if covered {
108                        for (_, child) in children.iter().rev() {
109                            self.stack.push((child, true));
110                        }
111                    } else {
112                        for (bounds, child) in children.iter().rev() {
113                            if self.predicate.could_match(bounds) {
114                                self.stack.push((child, self.predicate.covers_all(bounds)));
115                            }
116                        }
117                    }
118                }
119            }
120        }
121    }
122
123    /// `(0, None)`: a predicate cannot bound its match count without
124    /// doing the walk.
125    fn size_hint(&self) -> (usize, Option<usize>) {
126        (0, None)
127    }
128}
129
130impl<T: Indexable> FusedIterator for QueryIter<'_, T> {}
geometry_rtree/query_iter.rs

geometry_rtree/
query_iter.rs
