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