rstar/algorithm/
intersection_iterator.rs

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