rstar/algorithm/
intersection_iterator.rs1use 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
14pub 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}