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

rstar/algorithm/
iterators.rs
