rstar/algorithm/
iterators.rs

1use crate::algorithm::selection_functions::*;
2use crate::node::{ParentNode, RTreeNode};
3use crate::object::RTreeObject;
4use core::ops::ControlFlow;
5
6#[cfg(doc)]
7use crate::RTree;
8
9use smallvec::SmallVec;
10
11pub use super::intersection_iterator::IntersectionIterator;
12pub use super::removal::{DrainIterator, IntoIter};
13
14/// Iterator returned by [`RTree::locate_all_at_point`].
15pub type LocateAllAtPoint<'a, T> = SelectionIterator<'a, T, SelectAtPointFunction<T>>;
16/// Iterator returned by [`RTree::locate_all_at_point_mut`].
17pub type LocateAllAtPointMut<'a, T> = SelectionIteratorMut<'a, T, SelectAtPointFunction<T>>;
18
19/// Iterator returned by [`RTree::locate_in_envelope`].
20pub type LocateInEnvelope<'a, T> = SelectionIterator<'a, T, SelectInEnvelopeFunction<T>>;
21/// Iterator returned by [`RTree::locate_in_envelope_mut`].
22pub type LocateInEnvelopeMut<'a, T> = SelectionIteratorMut<'a, T, SelectInEnvelopeFunction<T>>;
23
24/// Iterator returned by [`RTree::locate_in_envelope_intersecting`].
25pub type LocateInEnvelopeIntersecting<'a, T> =
26    SelectionIterator<'a, T, SelectInEnvelopeFuncIntersecting<T>>;
27/// Iterator returned by [`RTree::locate_in_envelope_intersecting_mut`].
28pub type LocateInEnvelopeIntersectingMut<'a, T> =
29    SelectionIteratorMut<'a, T, SelectInEnvelopeFuncIntersecting<T>>;
30
31/// Iterator returned by [`RTree::iter`].
32pub type RTreeIterator<'a, T> = SelectionIterator<'a, T, SelectAllFunc>;
33/// Iterator returned by [`RTree::iter_mut`].
34pub type RTreeIteratorMut<'a, T> = SelectionIteratorMut<'a, T, SelectAllFunc>;
35
36/// Iterator returned by [`RTree::locate_within_distance`].
37pub type LocateWithinDistanceIterator<'a, T> =
38    SelectionIterator<'a, T, SelectWithinDistanceFunction<T>>;
39
40/// Iterator returned by `RTree::locate_*` methods.
41pub struct SelectionIterator<'a, T, Func>
42where
43    T: RTreeObject + 'a,
44    Func: SelectionFunction<T>,
45{
46    func: Func,
47    current_nodes: SmallVec<[&'a RTreeNode<T>; 24]>,
48}
49
50impl<'a, T, Func> SelectionIterator<'a, T, Func>
51where
52    T: RTreeObject,
53    Func: SelectionFunction<T>,
54{
55    pub(crate) fn new(root: &'a ParentNode<T>, func: Func) -> Self {
56        let current_nodes =
57            if !root.children.is_empty() && func.should_unpack_parent(&root.envelope()) {
58                root.children.iter().collect()
59            } else {
60                SmallVec::new()
61            };
62
63        SelectionIterator {
64            func,
65            current_nodes,
66        }
67    }
68}
69
70impl<'a, T, Func> Iterator for SelectionIterator<'a, T, Func>
71where
72    T: RTreeObject,
73    Func: SelectionFunction<T>,
74{
75    type Item = &'a T;
76
77    fn next(&mut self) -> Option<&'a T> {
78        while let Some(next) = self.current_nodes.pop() {
79            match next {
80                RTreeNode::Leaf(ref t) => {
81                    if self.func.should_unpack_leaf(t) {
82                        return Some(t);
83                    }
84                }
85                RTreeNode::Parent(ref data) => {
86                    if self.func.should_unpack_parent(&data.envelope) {
87                        self.current_nodes.extend(&data.children);
88                    }
89                }
90            }
91        }
92        None
93    }
94}
95
96/// Internal iteration variant of [`SelectionIterator`]
97pub fn select_nodes<'a, T, Func, V, B>(
98    root: &'a ParentNode<T>,
99    func: &Func,
100    visitor: &mut V,
101) -> ControlFlow<B>
102where
103    T: RTreeObject,
104    Func: SelectionFunction<T>,
105    V: FnMut(&'a T) -> ControlFlow<B>,
106{
107    struct Args<'a, Func, V> {
108        func: &'a Func,
109        visitor: &'a mut V,
110    }
111
112    fn inner<'a, T, Func, V, B>(
113        parent: &'a ParentNode<T>,
114        args: &mut Args<'_, Func, V>,
115    ) -> ControlFlow<B>
116    where
117        T: RTreeObject,
118        Func: SelectionFunction<T>,
119        V: FnMut(&'a T) -> ControlFlow<B>,
120    {
121        for node in parent.children.iter() {
122            match node {
123                RTreeNode::Leaf(ref t) => {
124                    if args.func.should_unpack_leaf(t) {
125                        (args.visitor)(t)?;
126                    }
127                }
128                RTreeNode::Parent(ref data) => {
129                    if args.func.should_unpack_parent(&data.envelope()) {
130                        inner(data, args)?;
131                    }
132                }
133            }
134        }
135
136        ControlFlow::Continue(())
137    }
138
139    if !root.children.is_empty() && func.should_unpack_parent(&root.envelope()) {
140        inner(root, &mut Args { func, visitor })?;
141    }
142
143    ControlFlow::Continue(())
144}
145
146/// Iterator type returned by `RTree::locate_*_mut` methods.
147pub struct SelectionIteratorMut<'a, T, Func>
148where
149    T: RTreeObject + 'a,
150    Func: SelectionFunction<T>,
151{
152    func: Func,
153    current_nodes: SmallVec<[&'a mut RTreeNode<T>; 32]>,
154}
155
156impl<'a, T, Func> SelectionIteratorMut<'a, T, Func>
157where
158    T: RTreeObject,
159    Func: SelectionFunction<T>,
160{
161    pub(crate) fn new(root: &'a mut ParentNode<T>, func: Func) -> Self {
162        let current_nodes =
163            if !root.children.is_empty() && func.should_unpack_parent(&root.envelope()) {
164                root.children.iter_mut().collect()
165            } else {
166                SmallVec::new()
167            };
168
169        SelectionIteratorMut {
170            func,
171            current_nodes,
172        }
173    }
174}
175
176impl<'a, T, Func> Iterator for SelectionIteratorMut<'a, T, Func>
177where
178    T: RTreeObject,
179    Func: SelectionFunction<T>,
180{
181    type Item = &'a mut T;
182
183    fn next(&mut self) -> Option<&'a mut T> {
184        while let Some(next) = self.current_nodes.pop() {
185            match next {
186                RTreeNode::Leaf(ref mut t) => {
187                    if self.func.should_unpack_leaf(t) {
188                        return Some(t);
189                    }
190                }
191                RTreeNode::Parent(ref mut data) => {
192                    if self.func.should_unpack_parent(&data.envelope) {
193                        self.current_nodes.extend(&mut data.children);
194                    }
195                }
196            }
197        }
198        None
199    }
200}
201
202/// Internal iteration variant of [`SelectionIteratorMut`]
203pub fn select_nodes_mut<'a, T, Func, V, B>(
204    root: &'a mut ParentNode<T>,
205    func: &Func,
206    visitor: &mut V,
207) -> ControlFlow<B>
208where
209    T: RTreeObject,
210    Func: SelectionFunction<T>,
211    V: FnMut(&'a mut T) -> ControlFlow<B>,
212{
213    struct Args<'a, Func, V> {
214        func: &'a Func,
215        visitor: &'a mut V,
216    }
217
218    fn inner<'a, T, Func, V, B>(
219        parent: &'a mut ParentNode<T>,
220        args: &mut Args<'_, Func, V>,
221    ) -> ControlFlow<B>
222    where
223        T: RTreeObject,
224        Func: SelectionFunction<T>,
225        V: FnMut(&'a mut T) -> ControlFlow<B>,
226    {
227        for node in parent.children.iter_mut() {
228            match node {
229                RTreeNode::Leaf(ref mut t) => {
230                    if args.func.should_unpack_leaf(t) {
231                        (args.visitor)(t)?;
232                    }
233                }
234                RTreeNode::Parent(ref mut data) => {
235                    if args.func.should_unpack_parent(&data.envelope()) {
236                        inner(data, args)?;
237                    }
238                }
239            }
240        }
241
242        ControlFlow::Continue(())
243    }
244
245    if !root.children.is_empty() && func.should_unpack_parent(&root.envelope()) {
246        inner(root, &mut Args { func, visitor })?;
247    }
248
249    ControlFlow::Continue(())
250}
251
252#[cfg(test)]
253mod test {
254    use crate::aabb::AABB;
255    use crate::envelope::Envelope;
256    use crate::object::RTreeObject;
257    use crate::rtree::RTree;
258    use crate::test_utilities::{create_random_points, create_random_rectangles, SEED_1};
259    use crate::SelectionFunction;
260
261    #[test]
262    fn test_root_node_is_not_always_unpacked() {
263        struct SelectNoneFunc {}
264
265        impl SelectionFunction<[i32; 2]> for SelectNoneFunc {
266            fn should_unpack_parent(&self, _: &AABB<[i32; 2]>) -> bool {
267                false
268            }
269        }
270
271        let mut tree = RTree::bulk_load(vec![[0i32, 0]]);
272
273        let mut elements = tree.locate_with_selection_function(SelectNoneFunc {});
274        assert!(elements.next().is_none());
275        drop(elements);
276
277        let mut elements = tree.locate_with_selection_function_mut(SelectNoneFunc {});
278        assert!(elements.next().is_none());
279    }
280
281    #[test]
282    fn test_locate_all() {
283        const NUM_RECTANGLES: usize = 400;
284        let rectangles = create_random_rectangles(NUM_RECTANGLES, SEED_1);
285        let tree = RTree::bulk_load(rectangles.clone());
286
287        let query_points = create_random_points(20, SEED_1);
288
289        for p in &query_points {
290            let contained_sequential: Vec<_> = rectangles
291                .iter()
292                .filter(|rectangle| rectangle.envelope().contains_point(p))
293                .cloned()
294                .collect();
295
296            let contained_rtree: Vec<_> = tree.locate_all_at_point(p).cloned().collect();
297
298            contained_sequential
299                .iter()
300                .all(|r| contained_rtree.contains(r));
301            contained_rtree
302                .iter()
303                .all(|r| contained_sequential.contains(r));
304        }
305    }
306
307    #[test]
308    fn test_locate_in_envelope() {
309        let points = create_random_points(100, SEED_1);
310        let tree = RTree::bulk_load(points.clone());
311        let envelope = AABB::from_corners([0.5, 0.5], [1.0, 1.0]);
312        let contained_in_envelope: Vec<_> = points
313            .iter()
314            .filter(|point| envelope.contains_point(point))
315            .cloned()
316            .collect();
317        let len = contained_in_envelope.len();
318        assert!(10 < len && len < 90, "unexpected point distribution");
319        let located: Vec<_> = tree.locate_in_envelope(&envelope).cloned().collect();
320        assert_eq!(len, located.len());
321        for point in &contained_in_envelope {
322            assert!(located.contains(point));
323        }
324    }
325
326    #[test]
327    fn test_locate_with_selection_func() {
328        use crate::SelectionFunction;
329
330        struct SelectLeftOfZeroPointFiveFunc;
331
332        impl SelectionFunction<[f64; 2]> for SelectLeftOfZeroPointFiveFunc {
333            fn should_unpack_parent(&self, parent_envelope: &AABB<[f64; 2]>) -> bool {
334                parent_envelope.lower()[0] < 0.5 || parent_envelope.upper()[0] < 0.5
335            }
336
337            fn should_unpack_leaf(&self, child: &[f64; 2]) -> bool {
338                child[0] < 0.5
339            }
340        }
341
342        let func = SelectLeftOfZeroPointFiveFunc;
343
344        let points = create_random_points(100, SEED_1);
345        let tree = RTree::bulk_load(points.clone());
346        let iterative_count = points
347            .iter()
348            .filter(|leaf| func.should_unpack_leaf(leaf))
349            .count();
350        let selected = tree
351            .locate_with_selection_function(func)
352            .collect::<Vec<_>>();
353
354        assert_eq!(iterative_count, selected.len());
355        assert!(iterative_count > 20); // Make sure that we do test something interesting
356        for point in &selected {
357            assert!(point[0] < 0.5);
358        }
359    }
360
361    #[test]
362    fn test_iteration() {
363        const NUM_POINTS: usize = 1000;
364        let points = create_random_points(NUM_POINTS, SEED_1);
365        let mut tree = RTree::new();
366        for p in &points {
367            tree.insert(*p);
368        }
369        let mut count = 0usize;
370        for p in tree.iter() {
371            assert!(points.iter().any(|q| q == p));
372            count += 1;
373        }
374        assert_eq!(count, NUM_POINTS);
375        count = 0;
376        for p in tree.iter_mut() {
377            assert!(points.iter().any(|q| q == p));
378            count += 1;
379        }
380        assert_eq!(count, NUM_POINTS);
381        for p in &points {
382            assert!(tree.iter().any(|q| q == p));
383            assert!(tree.iter_mut().any(|q| q == p));
384        }
385    }
386
387    #[test]
388    fn test_locate_within_distance() {
389        use crate::primitives::Line;
390
391        let points = create_random_points(100, SEED_1);
392        let tree = RTree::bulk_load(points.clone());
393        let circle_radius_2 = 0.3;
394        let circle_origin = [0.2, 0.6];
395        let contained_in_circle: Vec<_> = points
396            .iter()
397            .filter(|point| Line::new(circle_origin, **point).length_2() <= circle_radius_2)
398            .cloned()
399            .collect();
400        let located: Vec<_> = tree
401            .locate_within_distance(circle_origin, circle_radius_2)
402            .cloned()
403            .collect();
404
405        assert_eq!(located.len(), contained_in_circle.len());
406        for point in &contained_in_circle {
407            assert!(located.contains(point));
408        }
409    }
410
411    #[test]
412    fn test_locate_within_distance_on_empty_tree() {
413        let tree: RTree<[f64; 3]> = RTree::new();
414        tree.locate_within_distance([0.0, 0.0, 0.0], 10.0);
415
416        let tree: RTree<[i64; 3]> = RTree::new();
417        tree.locate_within_distance([0, 0, 0], 10);
418    }
419}
rstar/algorithm/iterators.rs

rstar/algorithm/
iterators.rs
