geometry_rtree/
query_iter.rs

1//! Lazy spatial-query iteration.
2//!
3//! [`QueryIter`] preserves the compact hot path for built-in
4//! [`Predicate`] queries. [`QueryWithIter`] runs the same pruning walk
5//! for logical and user-defined [`QueryPredicate`] values. Both mirror
6//! Boost's `visitors/spatial_query.hpp` and pause without traversing
7//! beyond the last yielded value.
8
9use alloc::vec::Vec;
10use core::iter::FusedIterator;
11
12use crate::bounds::Bounds;
13use crate::indexable::Indexable;
14use crate::node::Node;
15use crate::predicate::{Predicate, QueryPredicate};
16
17enum LeafMode {
18    Filtered,
19    DumpAll,
20}
21
22fn stack_capacity(height: usize, max_fanout: usize) -> usize {
23    height.saturating_sub(1) * max_fanout.saturating_sub(1) + max_fanout
24}
25
26/// A lazy iterator over values whose bounds satisfy a built-in
27/// [`Predicate`].
28///
29/// Created by [`Rtree::query_iter`](crate::Rtree::query_iter). A
30/// subtree that cannot match is skipped, while one wholly covered by
31/// the predicate is drained without further relation tests.
32pub struct QueryIter<'a, T> {
33    inner: QueryIterInner<'a, T>,
34}
35
36enum QueryIterInner<'a, T> {
37    Intersects(QueryWithIter<'a, T, IntersectsPredicate>),
38    Other(QueryWithIter<'a, T, Predicate>),
39}
40
41struct IntersectsPredicate(Bounds);
42
43impl<T: Indexable> QueryPredicate<T> for IntersectsPredicate {
44    #[inline]
45    fn matches(&self, value: &T) -> bool {
46        self.0.intersects(&value.bounds())
47    }
48
49    #[inline]
50    fn could_match(&self, node: &Bounds) -> bool {
51        self.0.intersects(node)
52    }
53
54    #[inline]
55    fn covers_all(&self, node: &Bounds) -> bool {
56        self.0.contains(node)
57    }
58}
59
60impl<'a, T> QueryIter<'a, T> {
61    pub(crate) fn new(
62        root: &'a Node<T>,
63        predicate: Predicate,
64        height: usize,
65        max_fanout: usize,
66    ) -> Self {
67        let inner = match predicate {
68            Predicate::Intersects(bounds) => QueryIterInner::Intersects(QueryWithIter::new(
69                root,
70                IntersectsPredicate(bounds),
71                height,
72                max_fanout,
73            )),
74            other => QueryIterInner::Other(QueryWithIter::new(root, other, height, max_fanout)),
75        };
76        Self { inner }
77    }
78}
79
80impl<'a, T: Indexable> Iterator for QueryIter<'a, T> {
81    type Item = &'a T;
82
83    fn next(&mut self) -> Option<Self::Item> {
84        match &mut self.inner {
85            QueryIterInner::Intersects(iter) => iter.next(),
86            QueryIterInner::Other(iter) => iter.next(),
87        }
88    }
89
90    fn size_hint(&self) -> (usize, Option<usize>) {
91        match &self.inner {
92            QueryIterInner::Intersects(iter) => iter.size_hint(),
93            QueryIterInner::Other(iter) => iter.size_hint(),
94        }
95    }
96}
97
98impl<T: Indexable> FusedIterator for QueryIter<'_, T> {}
99
100/// A lazy iterator over values accepted by a logical or user-defined
101/// [`QueryPredicate`].
102///
103/// Created by [`Rtree::query_iter_with`](crate::Rtree::query_iter_with).
104pub struct QueryWithIter<'a, T, P> {
105    predicate: P,
106    stack: Vec<(&'a Node<T>, bool)>,
107    leaf: core::slice::Iter<'a, T>,
108    leaf_mode: LeafMode,
109}
110
111impl<'a, T, P> QueryWithIter<'a, T, P> {
112    pub(crate) fn new(root: &'a Node<T>, predicate: P, height: usize, max_fanout: usize) -> Self {
113        let mut stack = Vec::with_capacity(stack_capacity(height, max_fanout));
114        stack.push((root, false));
115        Self {
116            predicate,
117            stack,
118            leaf: [].iter(),
119            leaf_mode: LeafMode::Filtered,
120        }
121    }
122}
123
124impl<'a, T, P> Iterator for QueryWithIter<'a, T, P>
125where
126    T: Indexable,
127    P: QueryPredicate<T>,
128{
129    type Item = &'a T;
130
131    fn next(&mut self) -> Option<Self::Item> {
132        loop {
133            match self.leaf_mode {
134                LeafMode::Filtered => {
135                    for value in self.leaf.by_ref() {
136                        if self.predicate.matches(value) {
137                            return Some(value);
138                        }
139                    }
140                }
141                LeafMode::DumpAll => {
142                    if let Some(value) = self.leaf.next() {
143                        return Some(value);
144                    }
145                }
146            }
147            let (node, covered) = self.stack.pop()?;
148            match node {
149                Node::Leaf(values) => {
150                    self.leaf = values.iter();
151                    self.leaf_mode = if covered {
152                        LeafMode::DumpAll
153                    } else {
154                        LeafMode::Filtered
155                    };
156                }
157                Node::Branch(children) => {
158                    if covered {
159                        for (_, child) in children.iter().rev() {
160                            self.stack.push((child, true));
161                        }
162                    } else {
163                        for (bounds, child) in children.iter().rev() {
164                            if self.predicate.could_match(bounds) {
165                                self.stack.push((child, self.predicate.covers_all(bounds)));
166                            }
167                        }
168                    }
169                }
170            }
171        }
172    }
173
174    fn size_hint(&self) -> (usize, Option<usize>) {
175        (0, None)
176    }
177}
178
179impl<T, P> FusedIterator for QueryWithIter<'_, T, P>
180where
181    T: Indexable,
182    P: QueryPredicate<T>,
183{
184}
geometry_rtree/query_iter.rs

geometry_rtree/
query_iter.rs
