rstar/algorithm/
intersection_iterator.rs

1use crate::node::ParentNode;
2use crate::Envelope;
3use crate::RTreeNode;
4use crate::RTreeNode::*;
5use crate::RTreeObject;
6
7#[cfg(not(test))]
8use alloc::vec::Vec;
9use core::mem::take;
10
11#[cfg(doc)]
12use crate::RTree;
13
14/// Iterator returned by [`RTree::intersection_candidates_with_other_tree`].
15pub struct IntersectionIterator<'a, T, U = T>
16where
17    T: RTreeObject,
18    U: RTreeObject,
19{
20    todo_list: Vec<(&'a RTreeNode<T>, &'a RTreeNode<U>)>,
21    candidates: Vec<&'a RTreeNode<U>>,
22}
23
24impl<'a, T, U> IntersectionIterator<'a, T, U>
25where
26    T: RTreeObject,
27    U: RTreeObject<Envelope = T::Envelope>,
28{
29    pub(crate) fn new(root1: &'a ParentNode<T>, root2: &'a ParentNode<U>) -> Self {
30        let mut intersections = IntersectionIterator {
31            todo_list: Vec::new(),
32            candidates: Vec::new(),
33        };
34        intersections.add_intersecting_children(root1, root2);
35        intersections
36    }
37
38    fn push_if_intersecting(&mut self, node1: &'a RTreeNode<T>, node2: &'a RTreeNode<U>) {
39        if node1.envelope().intersects(&node2.envelope()) {
40            self.todo_list.push((node1, node2));
41        }
42    }
43
44    fn add_intersecting_children(
45        &mut self,
46        parent1: &'a ParentNode<T>,
47        parent2: &'a ParentNode<U>,
48    ) {
49        if !parent1.envelope().intersects(&parent2.envelope()) {
50            return;
51        }
52        let children1 = parent1
53            .children()
54            .iter()
55            .filter(|c1| c1.envelope().intersects(&parent2.envelope()));
56
57        let mut children2 = take(&mut self.candidates);
58        children2.extend(
59            parent2
60                .children()
61                .iter()
62                .filter(|c2| c2.envelope().intersects(&parent1.envelope())),
63        );
64
65        for child1 in children1 {
66            for child2 in &children2 {
67                self.push_if_intersecting(child1, child2);
68            }
69        }
70
71        children2.clear();
72        self.candidates = children2;
73    }
74}
75
76impl<'a, T, U> Iterator for IntersectionIterator<'a, T, U>
77where
78    T: RTreeObject,
79    U: RTreeObject<Envelope = T::Envelope>,
80{
81    type Item = (&'a T, &'a U);
82
83    fn next(&mut self) -> Option<Self::Item> {
84        while let Some(next) = self.todo_list.pop() {
85            match next {
86                (Leaf(t1), Leaf(t2)) => return Some((t1, t2)),
87                (leaf @ Leaf(_), Parent(p)) => {
88                    p.children()
89                        .iter()
90                        .for_each(|c| self.push_if_intersecting(leaf, c));
91                }
92                (Parent(p), leaf @ Leaf(_)) => {
93                    p.children()
94                        .iter()
95                        .for_each(|c| self.push_if_intersecting(c, leaf));
96                }
97                (Parent(p1), Parent(p2)) => {
98                    self.add_intersecting_children(p1, p2);
99                }
100            }
101        }
102        None
103    }
104}
105
106#[cfg(test)]
107mod test {
108    use crate::test_utilities::*;
109    use crate::{Envelope, RTree, RTreeObject};
110
111    #[test]
112    fn test_intersection_between_trees() {
113        let rectangles1 = create_random_rectangles(100, SEED_1);
114        let rectangles2 = create_random_rectangles(42, SEED_2);
115
116        let mut intersections_brute_force = Vec::new();
117        for rectangle1 in &rectangles1 {
118            for rectangle2 in &rectangles2 {
119                if rectangle1.envelope().intersects(&rectangle2.envelope()) {
120                    intersections_brute_force.push((rectangle1, rectangle2));
121                }
122            }
123        }
124
125        let tree1 = RTree::bulk_load(rectangles1.clone());
126        let tree2 = RTree::bulk_load(rectangles2.clone());
127        let mut intersections_from_trees = tree1
128            .intersection_candidates_with_other_tree(&tree2)
129            .collect::<Vec<_>>();
130
131        intersections_brute_force.sort_by(|a, b| a.partial_cmp(b).unwrap());
132        intersections_from_trees.sort_by(|a, b| a.partial_cmp(b).unwrap());
133        assert_eq!(intersections_brute_force, intersections_from_trees);
134    }
135
136    #[test]
137    fn test_trivial_intersections() {
138        let points1 = create_random_points(1000, SEED_1);
139        let points2 = create_random_points(2000, SEED_2);
140        let tree1 = RTree::bulk_load(points1);
141        let tree2 = RTree::bulk_load(points2);
142
143        assert_eq!(
144            tree1
145                .intersection_candidates_with_other_tree(&tree2)
146                .count(),
147            0
148        );
149        assert_eq!(
150            tree1
151                .intersection_candidates_with_other_tree(&tree1)
152                .count(),
153            tree1.size()
154        );
155    }
156}