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