rstar/algorithm/
removal.rs

1use core::mem::replace;
2
3use crate::algorithm::selection_functions::SelectionFunction;
4use crate::node::{ParentNode, RTreeNode};
5use crate::object::RTreeObject;
6use crate::params::RTreeParams;
7use crate::{Envelope, RTree};
8
9#[cfg(not(test))]
10use alloc::{vec, vec::Vec};
11
12#[allow(unused_imports)] // Import is required when building without std
13use num_traits::Float;
14
15/// Iterator returned by `impl IntoIter for RTree`.
16///
17/// Consumes the whole tree and yields all leaf objects.
18pub struct IntoIter<T>
19where
20    T: RTreeObject,
21{
22    node_stack: Vec<RTreeNode<T>>,
23}
24
25impl<T> IntoIter<T>
26where
27    T: RTreeObject,
28{
29    pub(crate) fn new(root: ParentNode<T>) -> Self {
30        Self {
31            node_stack: vec![RTreeNode::Parent(root)],
32        }
33    }
34}
35
36impl<T> Iterator for IntoIter<T>
37where
38    T: RTreeObject,
39{
40    type Item = T;
41
42    fn next(&mut self) -> Option<Self::Item> {
43        while let Some(node) = self.node_stack.pop() {
44            match node {
45                RTreeNode::Leaf(object) => return Some(object),
46                RTreeNode::Parent(parent) => self.node_stack.extend(parent.children),
47            }
48        }
49
50        None
51    }
52}
53
54/// Iterator returned by `RTree::drain_*` methods.
55///
56/// Draining iterator that removes elements of the tree selected by a
57/// [`SelectionFunction`]. Returned by
58/// [`RTree::drain_with_selection_function`] and related methods.
59///
60/// # Remarks
61///
62/// This iterator is similar to the one returned by `Vec::drain` or
63/// `Vec::drain_filter`. Dropping the iterator at any point removes only
64/// the yielded values (this behaviour is unlike `Vec::drain_*`). Leaking
65/// this iterator leads to a leak amplification where all elements of the
66/// tree are leaked.
67pub struct DrainIterator<'a, T, R, Params>
68where
69    T: RTreeObject,
70    Params: RTreeParams,
71    R: SelectionFunction<T>,
72{
73    node_stack: Vec<(ParentNode<T>, usize, usize)>,
74    removal_function: R,
75    rtree: &'a mut RTree<T, Params>,
76    original_size: usize,
77}
78
79impl<'a, T, R, Params> DrainIterator<'a, T, R, Params>
80where
81    T: RTreeObject,
82    Params: RTreeParams,
83    R: SelectionFunction<T>,
84{
85    pub(crate) fn new(rtree: &'a mut RTree<T, Params>, removal_function: R) -> Self {
86        // We replace with a root as a brand new RTree in case the iterator is
87        // `mem::forgot`ten.
88
89        // Instead of using `new_with_params`, we avoid an allocation for
90        // the normal usage and replace root with an empty `Vec`.
91        let root = replace(
92            rtree.root_mut(),
93            ParentNode {
94                children: vec![],
95                envelope: Envelope::new_empty(),
96            },
97        );
98        let original_size = replace(rtree.size_mut(), 0);
99
100        let m = Params::MIN_SIZE;
101        let max_depth = (original_size as f32).log(m.max(2) as f32).ceil() as usize;
102        let mut node_stack = Vec::with_capacity(max_depth);
103        node_stack.push((root, 0, 0));
104
105        DrainIterator {
106            node_stack,
107            original_size,
108            removal_function,
109            rtree,
110        }
111    }
112
113    fn pop_node(&mut self, increment_idx: bool) -> Option<(ParentNode<T>, usize)> {
114        debug_assert!(!self.node_stack.is_empty());
115
116        let (mut node, _, num_removed) = self.node_stack.pop().unwrap();
117
118        // We only compute envelope for the current node as the parent
119        // is taken care of when it is popped.
120
121        // TODO: May be make this a method on `ParentNode`
122        if num_removed > 0 {
123            node.envelope = crate::node::envelope_for_children(&node.children);
124        }
125
126        // If there is no parent, this is the new root node to set back in the rtree
127        // O/w, get the new top in stack
128        let (parent_node, parent_idx, parent_removed) = match self.node_stack.last_mut() {
129            Some(pn) => (&mut pn.0, &mut pn.1, &mut pn.2),
130            None => return Some((node, num_removed)),
131        };
132
133        // Update the remove count on parent
134        *parent_removed += num_removed;
135
136        // If the node has no children, we don't need to add it back to the parent
137        if node.children.is_empty() {
138            return None;
139        }
140
141        // Put the child back (but re-arranged)
142        parent_node.children.push(RTreeNode::Parent(node));
143
144        // Swap it with the current item and increment idx.
145
146        // A minor optimization is to avoid the swap in the destructor,
147        // where we aren't going to be iterating any more.
148        if !increment_idx {
149            return None;
150        }
151
152        // Note that during iteration, parent_idx may be equal to
153        // (previous) children.len(), but this is okay as the swap will be
154        // a no-op.
155        let parent_len = parent_node.children.len();
156        parent_node.children.swap(*parent_idx, parent_len - 1);
157        *parent_idx += 1;
158
159        None
160    }
161}
162
163impl<'a, T, R, Params> Iterator for DrainIterator<'a, T, R, Params>
164where
165    T: RTreeObject,
166    Params: RTreeParams,
167    R: SelectionFunction<T>,
168{
169    type Item = T;
170
171    fn next(&mut self) -> Option<Self::Item> {
172        'attempt_loop: loop {
173            // Get reference to top node or return None.
174            let (node, idx, remove_count) = match self.node_stack.last_mut() {
175                Some(node) => (&mut node.0, &mut node.1, &mut node.2),
176                None => return None,
177            };
178
179            // Try to find a selected item to return.
180            if *idx > 0 || self.removal_function.should_unpack_parent(&node.envelope) {
181                while *idx < node.children.len() {
182                    match &mut node.children[*idx] {
183                        RTreeNode::Parent(_) => {
184                            // Swap node with last, remove and return the value.
185                            // No need to increment idx as something else has replaced it;
186                            // or idx == new len, and we'll handle it in the next iteration.
187                            let child = match node.children.swap_remove(*idx) {
188                                RTreeNode::Leaf(_) => unreachable!("DrainIterator bug!"),
189                                RTreeNode::Parent(node) => node,
190                            };
191                            self.node_stack.push((child, 0, 0));
192                            continue 'attempt_loop;
193                        }
194                        RTreeNode::Leaf(ref leaf) => {
195                            if self.removal_function.should_unpack_leaf(leaf) {
196                                // Swap node with last, remove and return the value.
197                                // No need to increment idx as something else has replaced it;
198                                // or idx == new len, and we'll handle it in the next iteration.
199                                *remove_count += 1;
200                                return match node.children.swap_remove(*idx) {
201                                    RTreeNode::Leaf(data) => Some(data),
202                                    _ => unreachable!("RemovalIterator bug!"),
203                                };
204                            }
205                            *idx += 1;
206                        }
207                    }
208                }
209            }
210
211            // Pop top node and clean-up if done
212            if let Some((new_root, total_removed)) = self.pop_node(true) {
213                // This happens if we are done with the iteration.
214                // Set the root back in rtree and return None
215                *self.rtree.root_mut() = new_root;
216                *self.rtree.size_mut() = self.original_size - total_removed;
217                return None;
218            }
219        }
220    }
221}
222
223impl<'a, T, R, Params> Drop for DrainIterator<'a, T, R, Params>
224where
225    T: RTreeObject,
226    Params: RTreeParams,
227    R: SelectionFunction<T>,
228{
229    fn drop(&mut self) {
230        // Re-assemble back the original rtree and update envelope as we
231        // re-assemble.
232        if self.node_stack.is_empty() {
233            // The iteration handled everything, nothing to do.
234            return;
235        }
236
237        loop {
238            debug_assert!(!self.node_stack.is_empty());
239            if let Some((new_root, total_removed)) = self.pop_node(false) {
240                *self.rtree.root_mut() = new_root;
241                *self.rtree.size_mut() = self.original_size - total_removed;
242                break;
243            }
244        }
245    }
246}
247
248#[cfg(test)]
249mod test {
250    use std::mem::forget;
251
252    use crate::algorithm::selection_functions::{SelectAllFunc, SelectInEnvelopeFuncIntersecting};
253    use crate::point::PointExt;
254    use crate::primitives::Line;
255    use crate::test_utilities::{create_random_points, create_random_rectangles, SEED_1, SEED_2};
256    use crate::AABB;
257
258    use super::*;
259
260    #[test]
261    fn test_remove_and_insert() {
262        const SIZE: usize = 1000;
263        let points = create_random_points(SIZE, SEED_1);
264        let later_insertions = create_random_points(SIZE, SEED_2);
265        let mut tree = RTree::bulk_load(points.clone());
266        for (point_to_remove, point_to_add) in points.iter().zip(later_insertions.iter()) {
267            assert!(tree.remove_at_point(point_to_remove).is_some());
268            tree.insert(*point_to_add);
269        }
270        assert_eq!(tree.size(), SIZE);
271        assert!(points.iter().all(|p| !tree.contains(p)));
272        assert!(later_insertions.iter().all(|p| tree.contains(p)));
273        for point in &later_insertions {
274            assert!(tree.remove_at_point(point).is_some());
275        }
276        assert_eq!(tree.size(), 0);
277    }
278
279    #[test]
280    fn test_remove_and_insert_rectangles() {
281        const SIZE: usize = 1000;
282        let initial_rectangles = create_random_rectangles(SIZE, SEED_1);
283        let new_rectangles = create_random_rectangles(SIZE, SEED_2);
284        let mut tree = RTree::bulk_load(initial_rectangles.clone());
285
286        for (rectangle_to_remove, rectangle_to_add) in
287            initial_rectangles.iter().zip(new_rectangles.iter())
288        {
289            assert!(tree.remove(rectangle_to_remove).is_some());
290            tree.insert(*rectangle_to_add);
291        }
292        assert_eq!(tree.size(), SIZE);
293        assert!(initial_rectangles.iter().all(|p| !tree.contains(p)));
294        assert!(new_rectangles.iter().all(|p| tree.contains(p)));
295        for rectangle in &new_rectangles {
296            assert!(tree.contains(rectangle));
297        }
298        for rectangle in &initial_rectangles {
299            assert!(!tree.contains(rectangle));
300        }
301        for rectangle in &new_rectangles {
302            assert!(tree.remove(rectangle).is_some());
303        }
304        assert_eq!(tree.size(), 0);
305    }
306
307    #[test]
308    fn test_remove_at_point() {
309        let points = create_random_points(1000, SEED_1);
310        let mut tree = RTree::bulk_load(points.clone());
311        for point in &points {
312            let size_before_removal = tree.size();
313            assert!(tree.remove_at_point(point).is_some());
314            assert!(tree.remove_at_point(&[1000.0, 1000.0]).is_none());
315            assert_eq!(size_before_removal - 1, tree.size());
316        }
317    }
318
319    #[test]
320    fn test_remove() {
321        let points = create_random_points(1000, SEED_1);
322        let offsets = create_random_points(1000, SEED_2);
323        let scaled = offsets.iter().map(|p| p.mul(0.05));
324        let edges: Vec<_> = points
325            .iter()
326            .zip(scaled)
327            .map(|(from, offset)| Line::new(*from, from.add(&offset)))
328            .collect();
329        let mut tree = RTree::bulk_load(edges.clone());
330        for edge in &edges {
331            let size_before_removal = tree.size();
332            assert!(tree.remove(edge).is_some());
333            assert!(tree.remove(edge).is_none());
334            assert_eq!(size_before_removal - 1, tree.size());
335        }
336    }
337
338    #[test]
339    fn test_drain_iterator() {
340        const SIZE: usize = 1000;
341        let points = create_random_points(SIZE, SEED_1);
342        let mut tree = RTree::bulk_load(points);
343
344        let drain_count = DrainIterator::new(&mut tree, SelectAllFunc)
345            .take(250)
346            .count();
347        assert_eq!(drain_count, 250);
348        assert_eq!(tree.size(), 750);
349
350        let drain_count = DrainIterator::new(&mut tree, SelectAllFunc)
351            .take(250)
352            .count();
353        assert_eq!(drain_count, 250);
354        assert_eq!(tree.size(), 500);
355
356        // Test Drain forget soundness
357        forget(DrainIterator::new(&mut tree, SelectAllFunc));
358        // Check tree has no nodes
359        // Tests below will check the same tree can be used again
360        assert_eq!(tree.size(), 0);
361
362        let points = create_random_points(1000, SEED_1);
363        points.into_iter().for_each(|pt| tree.insert(pt));
364
365        // The total for this is 406 (for SEED_1)
366        let env = AABB::from_corners([-2., -0.6], [0.5, 0.85]);
367
368        let sel = SelectInEnvelopeFuncIntersecting::new(env);
369        let drain_count = DrainIterator::new(&mut tree, sel).take(80).count();
370        assert_eq!(drain_count, 80);
371
372        let sel = SelectInEnvelopeFuncIntersecting::new(env);
373        let drain_count = DrainIterator::new(&mut tree, sel).count();
374        assert_eq!(drain_count, 326);
375
376        let sel = SelectInEnvelopeFuncIntersecting::new(env);
377        let sel_count = tree.locate_with_selection_function(sel).count();
378        assert_eq!(sel_count, 0);
379        assert_eq!(tree.size(), 1000 - 80 - 326);
380    }
381
382    #[test]
383    fn test_into_iter() {
384        const SIZE: usize = 100;
385        let mut points = create_random_points(SIZE, SEED_1);
386        let tree = RTree::bulk_load(points.clone());
387
388        let mut vec = tree.into_iter().collect::<Vec<_>>();
389
390        assert_eq!(vec.len(), points.len());
391
392        points.sort_unstable_by(|lhs, rhs| lhs.partial_cmp(rhs).unwrap());
393        vec.sort_unstable_by(|lhs, rhs| lhs.partial_cmp(rhs).unwrap());
394
395        assert_eq!(points, vec);
396    }
397}
rstar/algorithm/removal.rs

rstar/algorithm/
removal.rs
