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<T, R, Params> Iterator for DrainIterator<'_, 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) = self.node_stack.last_mut()?;
175
176            // Try to find a selected item to return.
177            if *idx > 0 || self.removal_function.should_unpack_parent(&node.envelope) {
178                while *idx < node.children.len() {
179                    match &mut node.children[*idx] {
180                        RTreeNode::Parent(_) => {
181                            // Swap node with last, remove and return the value.
182                            // No need to increment idx as something else has replaced it;
183                            // or idx == new len, and we'll handle it in the next iteration.
184                            let child = match node.children.swap_remove(*idx) {
185                                RTreeNode::Leaf(_) => unreachable!("DrainIterator bug!"),
186                                RTreeNode::Parent(node) => node,
187                            };
188                            self.node_stack.push((child, 0, 0));
189                            continue 'attempt_loop;
190                        }
191                        RTreeNode::Leaf(ref leaf) => {
192                            if self.removal_function.should_unpack_leaf(leaf) {
193                                // Swap node with last, remove and return the value.
194                                // No need to increment idx as something else has replaced it;
195                                // or idx == new len, and we'll handle it in the next iteration.
196                                *remove_count += 1;
197                                return match node.children.swap_remove(*idx) {
198                                    RTreeNode::Leaf(data) => Some(data),
199                                    RTreeNode::Parent(_) => unreachable!("RemovalIterator bug!"),
200                                };
201                            }
202                            *idx += 1;
203                        }
204                    }
205                }
206            }
207
208            // Pop top node and clean-up if done
209            if let Some((new_root, total_removed)) = self.pop_node(true) {
210                // This happens if we are done with the iteration.
211                // Set the root back in rtree and return None
212                *self.rtree.root_mut() = new_root;
213                *self.rtree.size_mut() = self.original_size - total_removed;
214                return None;
215            }
216        }
217    }
218}
219
220impl<T, R, Params> Drop for DrainIterator<'_, T, R, Params>
221where
222    T: RTreeObject,
223    Params: RTreeParams,
224    R: SelectionFunction<T>,
225{
226    fn drop(&mut self) {
227        // Re-assemble back the original rtree and update envelope as we
228        // re-assemble.
229        if self.node_stack.is_empty() {
230            // The iteration handled everything, nothing to do.
231            return;
232        }
233
234        loop {
235            debug_assert!(!self.node_stack.is_empty());
236            if let Some((new_root, total_removed)) = self.pop_node(false) {
237                *self.rtree.root_mut() = new_root;
238                *self.rtree.size_mut() = self.original_size - total_removed;
239                break;
240            }
241        }
242    }
243}
244
245#[cfg(test)]
246mod test {
247    use std::mem::forget;
248
249    use crate::algorithm::selection_functions::{SelectAllFunc, SelectInEnvelopeFuncIntersecting};
250    use crate::point::PointExt;
251    use crate::primitives::Line;
252    use crate::test_utilities::{create_random_points, create_random_rectangles, SEED_1, SEED_2};
253    use crate::AABB;
254
255    use super::*;
256
257    #[test]
258    fn test_remove_and_insert() {
259        const SIZE: usize = 1000;
260        let points = create_random_points(SIZE, SEED_1);
261        let later_insertions = create_random_points(SIZE, SEED_2);
262        let mut tree = RTree::bulk_load(points.clone());
263        for (point_to_remove, point_to_add) in points.iter().zip(later_insertions.iter()) {
264            assert!(tree.remove_at_point(*point_to_remove).is_some());
265            tree.insert(*point_to_add);
266        }
267        assert_eq!(tree.size(), SIZE);
268        assert!(points.iter().all(|p| !tree.contains(p)));
269        assert!(later_insertions.iter().all(|p| tree.contains(p)));
270        for point in &later_insertions {
271            assert!(tree.remove_at_point(*point).is_some());
272        }
273        assert_eq!(tree.size(), 0);
274    }
275
276    #[test]
277    fn test_remove_and_insert_rectangles() {
278        const SIZE: usize = 1000;
279        let initial_rectangles = create_random_rectangles(SIZE, SEED_1);
280        let new_rectangles = create_random_rectangles(SIZE, SEED_2);
281        let mut tree = RTree::bulk_load(initial_rectangles.clone());
282
283        for (rectangle_to_remove, rectangle_to_add) in
284            initial_rectangles.iter().zip(new_rectangles.iter())
285        {
286            assert!(tree.remove(rectangle_to_remove).is_some());
287            tree.insert(*rectangle_to_add);
288        }
289        assert_eq!(tree.size(), SIZE);
290        assert!(initial_rectangles.iter().all(|p| !tree.contains(p)));
291        assert!(new_rectangles.iter().all(|p| tree.contains(p)));
292        for rectangle in &new_rectangles {
293            assert!(tree.contains(rectangle));
294        }
295        for rectangle in &initial_rectangles {
296            assert!(!tree.contains(rectangle));
297        }
298        for rectangle in &new_rectangles {
299            assert!(tree.remove(rectangle).is_some());
300        }
301        assert_eq!(tree.size(), 0);
302    }
303
304    #[test]
305    fn test_remove_at_point() {
306        let points = create_random_points(1000, SEED_1);
307        let mut tree = RTree::bulk_load(points.clone());
308        for point in points {
309            let size_before_removal = tree.size();
310            assert!(tree.remove_at_point(point).is_some());
311            assert!(tree.remove_at_point([1000.0, 1000.0]).is_none());
312            assert_eq!(size_before_removal - 1, tree.size());
313        }
314    }
315
316    #[test]
317    fn test_remove() {
318        let points = create_random_points(1000, SEED_1);
319        let offsets = create_random_points(1000, SEED_2);
320        let scaled = offsets.iter().map(|p| p.mul(0.05));
321        let edges: Vec<_> = points
322            .iter()
323            .zip(scaled)
324            .map(|(from, offset)| Line::new(*from, from.add(&offset)))
325            .collect();
326        let mut tree = RTree::bulk_load(edges.clone());
327        for edge in &edges {
328            let size_before_removal = tree.size();
329            assert!(tree.remove(edge).is_some());
330            assert!(tree.remove(edge).is_none());
331            assert_eq!(size_before_removal - 1, tree.size());
332        }
333    }
334
335    #[test]
336    fn test_drain_iterator() {
337        const SIZE: usize = 1000;
338        let points = create_random_points(SIZE, SEED_1);
339        let mut tree = RTree::bulk_load(points);
340
341        let drain_count = DrainIterator::new(&mut tree, SelectAllFunc)
342            .take(250)
343            .count();
344        assert_eq!(drain_count, 250);
345        assert_eq!(tree.size(), 750);
346
347        let drain_count = DrainIterator::new(&mut tree, SelectAllFunc)
348            .take(250)
349            .count();
350        assert_eq!(drain_count, 250);
351        assert_eq!(tree.size(), 500);
352
353        // Test Drain forget soundness
354        forget(DrainIterator::new(&mut tree, SelectAllFunc));
355        // Check tree has no nodes
356        // Tests below will check the same tree can be used again
357        assert_eq!(tree.size(), 0);
358
359        let points = create_random_points(1000, SEED_1);
360        points.into_iter().for_each(|pt| tree.insert(pt));
361
362        // The total for this is 406 (for SEED_1)
363        let env = AABB::from_corners([-2., -0.6], [0.5, 0.85]);
364
365        let sel = SelectInEnvelopeFuncIntersecting::new(env);
366        let drain_count = DrainIterator::new(&mut tree, sel).take(80).count();
367        assert_eq!(drain_count, 80);
368
369        let sel = SelectInEnvelopeFuncIntersecting::new(env);
370        let drain_count = DrainIterator::new(&mut tree, sel).count();
371        assert_eq!(drain_count, 326);
372
373        let sel = SelectInEnvelopeFuncIntersecting::new(env);
374        let sel_count = tree.locate_with_selection_function(sel).count();
375        assert_eq!(sel_count, 0);
376        assert_eq!(tree.size(), 1000 - 80 - 326);
377    }
378
379    #[test]
380    fn test_into_iter() {
381        const SIZE: usize = 100;
382        let mut points = create_random_points(SIZE, SEED_1);
383        let tree = RTree::bulk_load(points.clone());
384
385        let mut vec = tree.into_iter().collect::<Vec<_>>();
386
387        assert_eq!(vec.len(), points.len());
388
389        points.sort_unstable_by(|lhs, rhs| lhs.partial_cmp(rhs).unwrap());
390        vec.sort_unstable_by(|lhs, rhs| lhs.partial_cmp(rhs).unwrap());
391
392        assert_eq!(points, vec);
393    }
394}
rstar/algorithm/removal.rs

rstar/algorithm/
removal.rs
