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

rstar/algorithm/
removal.rs
