rstared/rstar/algorithm/
removal.rs

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

rstared/rstar/algorithm/
removal.rs
