1use crate::{Bounds, Half};
2use std::{
3 cmp,
4 fmt::Debug,
5 ops::{Add, Sub},
6};
7
8#[derive(Debug)]
9pub(crate) struct BoundsTree<U>
10where
11 U: Clone + Debug + Default + PartialEq,
12{
13 root: Option<usize>,
14 nodes: Vec<Node<U>>,
15 stack: Vec<usize>,
16}
17
18impl<U> BoundsTree<U>
19where
20 U: Clone
21 + Debug
22 + PartialEq
23 + PartialOrd
24 + Add<U, Output = U>
25 + Sub<Output = U>
26 + Half
27 + Default,
28{
29 pub fn clear(&mut self) {
30 self.root = None;
31 self.nodes.clear();
32 self.stack.clear();
33 }
34
35 pub fn insert(&mut self, new_bounds: Bounds<U>) -> u32 {
36 // If the tree is empty, make the root the new leaf.
37 let Some(mut index) = self.root else {
38 let new_node = self.push_leaf(new_bounds, 1);
39 self.root = Some(new_node);
40 return 1;
41 };
42
43 // Search for the best place to add the new leaf based on heuristics.
44 let mut max_intersecting_ordering = 0;
45 while let Node::Internal {
46 left,
47 right,
48 bounds: node_bounds,
49 ..
50 } = &mut self.nodes[index]
51 {
52 let left = *left;
53 let right = *right;
54 *node_bounds = node_bounds.union(&new_bounds);
55 self.stack.push(index);
56
57 // Descend to the best-fit child, based on which one would increase
58 // the surface area the least. This attempts to keep the tree balanced
59 // in terms of surface area. If there is an intersection with the other child,
60 // add its keys to the intersections vector.
61 let left_cost = new_bounds.union(self.nodes[left].bounds()).half_perimeter();
62 let right_cost = new_bounds
63 .union(self.nodes[right].bounds())
64 .half_perimeter();
65 if left_cost < right_cost {
66 max_intersecting_ordering =
67 self.find_max_ordering(right, &new_bounds, max_intersecting_ordering);
68 index = left;
69 } else {
70 max_intersecting_ordering =
71 self.find_max_ordering(left, &new_bounds, max_intersecting_ordering);
72 index = right;
73 }
74 }
75
76 // We've found a leaf ('index' now refers to a leaf node).
77 // We'll insert a new parent node above the leaf and attach our new leaf to it.
78 let sibling = index;
79
80 // Check for collision with the located leaf node
81 let Node::Leaf {
82 bounds: sibling_bounds,
83 order: sibling_ordering,
84 ..
85 } = &self.nodes[index]
86 else {
87 unreachable!();
88 };
89 if sibling_bounds.intersects(&new_bounds) {
90 max_intersecting_ordering = cmp::max(max_intersecting_ordering, *sibling_ordering);
91 }
92
93 let ordering = max_intersecting_ordering + 1;
94 let new_node = self.push_leaf(new_bounds, ordering);
95 let new_parent = self.push_internal(sibling, new_node);
96
97 // If there was an old parent, we need to update its children indices.
98 if let Some(old_parent) = self.stack.last().copied() {
99 let Node::Internal { left, right, .. } = &mut self.nodes[old_parent] else {
100 unreachable!();
101 };
102
103 if *left == sibling {
104 *left = new_parent;
105 } else {
106 *right = new_parent;
107 }
108 } else {
109 // If the old parent was the root, the new parent is the new root.
110 self.root = Some(new_parent);
111 }
112
113 for node_index in self.stack.drain(..).rev() {
114 let Node::Internal {
115 max_order: max_ordering,
116 ..
117 } = &mut self.nodes[node_index]
118 else {
119 unreachable!()
120 };
121 if *max_ordering >= ordering {
122 break;
123 }
124 *max_ordering = ordering;
125 }
126
127 ordering
128 }
129
130 fn find_max_ordering(&self, index: usize, bounds: &Bounds<U>, mut max_ordering: u32) -> u32 {
131 match &self.nodes[index] {
132 Node::Leaf {
133 bounds: node_bounds,
134 order: ordering,
135 ..
136 } => {
137 if bounds.intersects(node_bounds) {
138 max_ordering = cmp::max(*ordering, max_ordering);
139 }
140 }
141 Node::Internal {
142 left,
143 right,
144 bounds: node_bounds,
145 max_order: node_max_ordering,
146 ..
147 } => {
148 if bounds.intersects(node_bounds) && max_ordering < *node_max_ordering {
149 let left_max_ordering = self.nodes[*left].max_ordering();
150 let right_max_ordering = self.nodes[*right].max_ordering();
151 if left_max_ordering > right_max_ordering {
152 max_ordering = self.find_max_ordering(*left, bounds, max_ordering);
153 max_ordering = self.find_max_ordering(*right, bounds, max_ordering);
154 } else {
155 max_ordering = self.find_max_ordering(*right, bounds, max_ordering);
156 max_ordering = self.find_max_ordering(*left, bounds, max_ordering);
157 }
158 }
159 }
160 }
161 max_ordering
162 }
163
164 fn push_leaf(&mut self, bounds: Bounds<U>, order: u32) -> usize {
165 self.nodes.push(Node::Leaf { bounds, order });
166 self.nodes.len() - 1
167 }
168
169 fn push_internal(&mut self, left: usize, right: usize) -> usize {
170 let left_node = &self.nodes[left];
171 let right_node = &self.nodes[right];
172 let new_bounds = left_node.bounds().union(right_node.bounds());
173 let max_ordering = cmp::max(left_node.max_ordering(), right_node.max_ordering());
174 self.nodes.push(Node::Internal {
175 bounds: new_bounds,
176 left,
177 right,
178 max_order: max_ordering,
179 });
180 self.nodes.len() - 1
181 }
182}
183
184impl<U> Default for BoundsTree<U>
185where
186 U: Clone + Debug + Default + PartialEq,
187{
188 fn default() -> Self {
189 BoundsTree {
190 root: None,
191 nodes: Vec::new(),
192 stack: Vec::new(),
193 }
194 }
195}
196
197#[derive(Debug, Clone)]
198enum Node<U>
199where
200 U: Clone + Debug + Default + PartialEq,
201{
202 Leaf {
203 bounds: Bounds<U>,
204 order: u32,
205 },
206 Internal {
207 left: usize,
208 right: usize,
209 bounds: Bounds<U>,
210 max_order: u32,
211 },
212}
213
214impl<U> Node<U>
215where
216 U: Clone + Debug + Default + PartialEq,
217{
218 fn bounds(&self) -> &Bounds<U> {
219 match self {
220 Node::Leaf { bounds, .. } => bounds,
221 Node::Internal { bounds, .. } => bounds,
222 }
223 }
224
225 fn max_ordering(&self) -> u32 {
226 match self {
227 Node::Leaf {
228 order: ordering, ..
229 } => *ordering,
230 Node::Internal {
231 max_order: max_ordering,
232 ..
233 } => *max_ordering,
234 }
235 }
236}
237
238#[cfg(test)]
239mod tests {
240 use super::*;
241 use crate::{Bounds, Point, Size};
242 use rand::{Rng, SeedableRng};
243
244 #[test]
245 fn test_insert() {
246 let mut tree = BoundsTree::<f32>::default();
247 let bounds1 = Bounds {
248 origin: Point { x: 0.0, y: 0.0 },
249 size: Size {
250 width: 10.0,
251 height: 10.0,
252 },
253 };
254 let bounds2 = Bounds {
255 origin: Point { x: 5.0, y: 5.0 },
256 size: Size {
257 width: 10.0,
258 height: 10.0,
259 },
260 };
261 let bounds3 = Bounds {
262 origin: Point { x: 10.0, y: 10.0 },
263 size: Size {
264 width: 10.0,
265 height: 10.0,
266 },
267 };
268
269 // Insert the bounds into the tree and verify the order is correct
270 assert_eq!(tree.insert(bounds1), 1);
271 assert_eq!(tree.insert(bounds2), 2);
272 assert_eq!(tree.insert(bounds3), 3);
273
274 // Insert non-overlapping bounds and verify they can reuse orders
275 let bounds4 = Bounds {
276 origin: Point { x: 20.0, y: 20.0 },
277 size: Size {
278 width: 10.0,
279 height: 10.0,
280 },
281 };
282 let bounds5 = Bounds {
283 origin: Point { x: 40.0, y: 40.0 },
284 size: Size {
285 width: 10.0,
286 height: 10.0,
287 },
288 };
289 let bounds6 = Bounds {
290 origin: Point { x: 25.0, y: 25.0 },
291 size: Size {
292 width: 10.0,
293 height: 10.0,
294 },
295 };
296 assert_eq!(tree.insert(bounds4), 1); // bounds4 does not overlap with bounds1, bounds2, or bounds3
297 assert_eq!(tree.insert(bounds5), 1); // bounds5 does not overlap with any other bounds
298 assert_eq!(tree.insert(bounds6), 2); // bounds6 overlaps with bounds4, so it should have a different order
299 }
300
301 #[test]
302 fn test_random_iterations() {
303 let max_bounds = 100;
304 for seed in 1..=1000 {
305 // let seed = 44;
306 let mut tree = BoundsTree::default();
307 let mut rng = rand::rngs::StdRng::seed_from_u64(seed as u64);
308 let mut expected_quads: Vec<(Bounds<f32>, u32)> = Vec::new();
309
310 // Insert a random number of random AABBs into the tree.
311 let num_bounds = rng.random_range(1..=max_bounds);
312 for _ in 0..num_bounds {
313 let min_x: f32 = rng.random_range(-100.0..100.0);
314 let min_y: f32 = rng.random_range(-100.0..100.0);
315 let width: f32 = rng.random_range(0.0..50.0);
316 let height: f32 = rng.random_range(0.0..50.0);
317 let bounds = Bounds {
318 origin: Point { x: min_x, y: min_y },
319 size: Size { width, height },
320 };
321
322 let expected_ordering = expected_quads
323 .iter()
324 .filter_map(|quad| quad.0.intersects(&bounds).then_some(quad.1))
325 .max()
326 .unwrap_or(0)
327 + 1;
328 expected_quads.push((bounds, expected_ordering));
329
330 // Insert the AABB into the tree and collect intersections.
331 let actual_ordering = tree.insert(bounds);
332 assert_eq!(actual_ordering, expected_ordering);
333 }
334 }
335 }
336}