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