rstar/algorithm/
iterators.rs

1use crate::algorithm::selection_functions::*;
2use crate::node::{ParentNode, RTreeNode};
3use crate::object::RTreeObject;
4
5#[cfg(doc)]
6use crate::RTree;
7
8use smallvec::SmallVec;
9
10pub use super::intersection_iterator::IntersectionIterator;
11pub use super::removal::DrainIterator;
12
13/// Iterator returned by [`RTree::locate_all_at_point`].
14pub type LocateAllAtPoint<'a, T> = SelectionIterator<'a, T, SelectAtPointFunction<T>>;
15/// Iterator returned by [`RTree::locate_all_at_point_mut`].
16pub type LocateAllAtPointMut<'a, T> = SelectionIteratorMut<'a, T, SelectAtPointFunction<T>>;
17
18/// Iterator returned by [`RTree::locate_in_envelope`].
19pub type LocateInEnvelope<'a, T> = SelectionIterator<'a, T, SelectInEnvelopeFunction<T>>;
20/// Iterator returned by [`RTree::locate_in_envelope_mut`].
21pub type LocateInEnvelopeMut<'a, T> = SelectionIteratorMut<'a, T, SelectInEnvelopeFunction<T>>;
22
23/// Iterator returned by [`RTree::locate_in_envelope_intersecting`].
24pub type LocateInEnvelopeIntersecting<'a, T> =
25    SelectionIterator<'a, T, SelectInEnvelopeFuncIntersecting<T>>;
26/// Iterator returned by [`RTree::locate_in_envelope_intersecting_mut`].
27pub type LocateInEnvelopeIntersectingMut<'a, T> =
28    SelectionIteratorMut<'a, T, SelectInEnvelopeFuncIntersecting<T>>;
29
30/// Iterator returned by [`RTree::iter`].
31pub type RTreeIterator<'a, T> = SelectionIterator<'a, T, SelectAllFunc>;
32/// Iterator returned by [`RTree::iter_mut`].
33pub type RTreeIteratorMut<'a, T> = SelectionIteratorMut<'a, T, SelectAllFunc>;
34
35/// Iterator returned by [`RTree::locate_within_distance`].
36pub type LocateWithinDistanceIterator<'a, T> =
37    SelectionIterator<'a, T, SelectWithinDistanceFunction<T>>;
38
39/// Iterator returned by `RTree::locate_*` methods.
40pub struct SelectionIterator<'a, T, Func>
41where
42    T: RTreeObject + 'a,
43    Func: SelectionFunction<T>,
44{
45    func: Func,
46    current_nodes: SmallVec<[&'a RTreeNode<T>; 24]>,
47}
48
49impl<'a, T, Func> SelectionIterator<'a, T, Func>
50where
51    T: RTreeObject,
52    Func: SelectionFunction<T>,
53{
54    pub(crate) fn new(root: &'a ParentNode<T>, func: Func) -> Self {
55        let current_nodes = if func.should_unpack_parent(&root.envelope()) {
56            root.children.iter().collect()
57        } else {
58            SmallVec::new()
59        };
60
61        SelectionIterator {
62            func,
63            current_nodes,
64        }
65    }
66}
67
68impl<'a, T, Func> Iterator for SelectionIterator<'a, T, Func>
69where
70    T: RTreeObject,
71    Func: SelectionFunction<T>,
72{
73    type Item = &'a T;
74
75    fn next(&mut self) -> Option<&'a T> {
76        while let Some(next) = self.current_nodes.pop() {
77            match next {
78                RTreeNode::Leaf(ref t) => {
79                    if self.func.should_unpack_leaf(t) {
80                        return Some(t);
81                    }
82                }
83                RTreeNode::Parent(ref data) => {
84                    if self.func.should_unpack_parent(&data.envelope) {
85                        self.current_nodes.extend(&data.children);
86                    }
87                }
88            }
89        }
90        None
91    }
92}
93
94/// Iterator type returned by `RTree::locate_*_mut` methods.
95pub struct SelectionIteratorMut<'a, T, Func>
96where
97    T: RTreeObject + 'a,
98    Func: SelectionFunction<T>,
99{
100    func: Func,
101    current_nodes: SmallVec<[&'a mut RTreeNode<T>; 32]>,
102}
103
104impl<'a, T, Func> SelectionIteratorMut<'a, T, Func>
105where
106    T: RTreeObject,
107    Func: SelectionFunction<T>,
108{
109    pub(crate) fn new(root: &'a mut ParentNode<T>, func: Func) -> Self {
110        let current_nodes = if func.should_unpack_parent(&root.envelope()) {
111            root.children.iter_mut().collect()
112        } else {
113            SmallVec::new()
114        };
115
116        SelectionIteratorMut {
117            func,
118            current_nodes,
119        }
120    }
121}
122
123impl<'a, T, Func> Iterator for SelectionIteratorMut<'a, T, Func>
124where
125    T: RTreeObject,
126    Func: SelectionFunction<T>,
127{
128    type Item = &'a mut T;
129
130    fn next(&mut self) -> Option<&'a mut T> {
131        while let Some(next) = self.current_nodes.pop() {
132            match next {
133                RTreeNode::Leaf(ref mut t) => {
134                    if self.func.should_unpack_leaf(t) {
135                        return Some(t);
136                    }
137                }
138                RTreeNode::Parent(ref mut data) => {
139                    if self.func.should_unpack_parent(&data.envelope) {
140                        self.current_nodes.extend(&mut data.children);
141                    }
142                }
143            }
144        }
145        None
146    }
147}
148
149#[cfg(test)]
150mod test {
151    use crate::aabb::AABB;
152    use crate::envelope::Envelope;
153    use crate::object::RTreeObject;
154    use crate::rtree::RTree;
155    use crate::test_utilities::{create_random_points, create_random_rectangles, SEED_1};
156    use crate::SelectionFunction;
157
158    #[test]
159    fn test_root_node_is_not_always_unpacked() {
160        struct SelectNoneFunc {}
161
162        impl SelectionFunction<[i32; 2]> for SelectNoneFunc {
163            fn should_unpack_parent(&self, _: &AABB<[i32; 2]>) -> bool {
164                false
165            }
166        }
167
168        let mut tree = RTree::bulk_load(vec![[0i32, 0]]);
169
170        let mut elements = tree.locate_with_selection_function(SelectNoneFunc {});
171        assert!(elements.next().is_none());
172        drop(elements);
173
174        let mut elements = tree.locate_with_selection_function_mut(SelectNoneFunc {});
175        assert!(elements.next().is_none());
176    }
177
178    #[test]
179    fn test_locate_all() {
180        const NUM_RECTANGLES: usize = 400;
181        let rectangles = create_random_rectangles(NUM_RECTANGLES, SEED_1);
182        let tree = RTree::bulk_load(rectangles.clone());
183
184        let query_points = create_random_points(20, SEED_1);
185
186        for p in &query_points {
187            let contained_sequential: Vec<_> = rectangles
188                .iter()
189                .filter(|rectangle| rectangle.envelope().contains_point(p))
190                .cloned()
191                .collect();
192
193            let contained_rtree: Vec<_> = tree.locate_all_at_point(p).cloned().collect();
194
195            contained_sequential
196                .iter()
197                .all(|r| contained_rtree.contains(r));
198            contained_rtree
199                .iter()
200                .all(|r| contained_sequential.contains(r));
201        }
202    }
203
204    #[test]
205    fn test_locate_in_envelope() {
206        let points = create_random_points(100, SEED_1);
207        let tree = RTree::bulk_load(points.clone());
208        let envelope = AABB::from_corners([0.5, 0.5], [1.0, 1.0]);
209        let contained_in_envelope: Vec<_> = points
210            .iter()
211            .filter(|point| envelope.contains_point(point))
212            .cloned()
213            .collect();
214        let len = contained_in_envelope.len();
215        assert!(10 < len && len < 90, "unexpected point distribution");
216        let located: Vec<_> = tree.locate_in_envelope(&envelope).cloned().collect();
217        assert_eq!(len, located.len());
218        for point in &contained_in_envelope {
219            assert!(located.contains(point));
220        }
221    }
222
223    #[test]
224    fn test_locate_with_selection_func() {
225        use crate::SelectionFunction;
226
227        struct SelectLeftOfZeroPointFiveFunc;
228
229        impl SelectionFunction<[f64; 2]> for SelectLeftOfZeroPointFiveFunc {
230            fn should_unpack_parent(&self, parent_envelope: &AABB<[f64; 2]>) -> bool {
231                parent_envelope.lower()[0] < 0.5 || parent_envelope.upper()[0] < 0.5
232            }
233
234            fn should_unpack_leaf(&self, child: &[f64; 2]) -> bool {
235                child[0] < 0.5
236            }
237        }
238
239        let func = SelectLeftOfZeroPointFiveFunc;
240
241        let points = create_random_points(100, SEED_1);
242        let tree = RTree::bulk_load(points.clone());
243        let iterative_count = points
244            .iter()
245            .filter(|leaf| func.should_unpack_leaf(leaf))
246            .count();
247        let selected = tree
248            .locate_with_selection_function(func)
249            .collect::<Vec<_>>();
250
251        assert_eq!(iterative_count, selected.len());
252        assert!(iterative_count > 20); // Make sure that we do test something interesting
253        for point in &selected {
254            assert!(point[0] < 0.5);
255        }
256    }
257
258    #[test]
259    fn test_iteration() {
260        const NUM_POINTS: usize = 1000;
261        let points = create_random_points(NUM_POINTS, SEED_1);
262        let mut tree = RTree::new();
263        for p in &points {
264            tree.insert(p.clone());
265        }
266        let mut count = 0usize;
267        for p in tree.iter() {
268            assert!(points.iter().any(|q| q == p));
269            count += 1;
270        }
271        assert_eq!(count, NUM_POINTS);
272        count = 0;
273        for p in tree.iter_mut() {
274            assert!(points.iter().any(|q| q == p));
275            count += 1;
276        }
277        assert_eq!(count, NUM_POINTS);
278        for p in &points {
279            assert!(tree.iter().any(|q| q == p));
280            assert!(tree.iter_mut().any(|q| q == p));
281        }
282    }
283
284    #[test]
285    fn test_locate_within_distance() {
286        use crate::primitives::Line;
287
288        let points = create_random_points(100, SEED_1);
289        let tree = RTree::bulk_load(points.clone());
290        let circle_radius_2 = 0.3;
291        let circle_origin = [0.2, 0.6];
292        let contained_in_circle: Vec<_> = points
293            .iter()
294            .filter(|point| Line::new(circle_origin, **point).length_2() <= circle_radius_2)
295            .cloned()
296            .collect();
297        let located: Vec<_> = tree
298            .locate_within_distance(circle_origin, circle_radius_2)
299            .cloned()
300            .collect();
301
302        assert_eq!(located.len(), contained_in_circle.len());
303        for point in &contained_in_circle {
304            assert!(located.contains(point));
305        }
306    }
307}
rstar/algorithm/iterators.rs

rstar/algorithm/
iterators.rs
