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(..) {
115 let Node::Internal {
116 max_order: max_ordering,
117 ..
118 } = &mut self.nodes[node_index]
119 else {
120 unreachable!()
121 };
122 *max_ordering = cmp::max(*max_ordering, ordering);
123 }
124
125 ordering
126 }
127
128 fn find_max_ordering(&self, index: usize, bounds: &Bounds<U>, mut max_ordering: u32) -> u32 {
129 match &self.nodes[index] {
130 Node::Leaf {
131 bounds: node_bounds,
132 order: ordering,
133 ..
134 } => {
135 if bounds.intersects(node_bounds) {
136 max_ordering = cmp::max(*ordering, max_ordering);
137 }
138 }
139 Node::Internal {
140 left,
141 right,
142 bounds: node_bounds,
143 max_order: node_max_ordering,
144 ..
145 } => {
146 if bounds.intersects(node_bounds) && max_ordering < *node_max_ordering {
147 let left_max_ordering = self.nodes[*left].max_ordering();
148 let right_max_ordering = self.nodes[*right].max_ordering();
149 if left_max_ordering > right_max_ordering {
150 max_ordering = self.find_max_ordering(*left, bounds, max_ordering);
151 max_ordering = self.find_max_ordering(*right, bounds, max_ordering);
152 } else {
153 max_ordering = self.find_max_ordering(*right, bounds, max_ordering);
154 max_ordering = self.find_max_ordering(*left, bounds, max_ordering);
155 }
156 }
157 }
158 }
159 max_ordering
160 }
161
162 fn push_leaf(&mut self, bounds: Bounds<U>, order: u32) -> usize {
163 self.nodes.push(Node::Leaf { bounds, order });
164 self.nodes.len() - 1
165 }
166
167 fn push_internal(&mut self, left: usize, right: usize) -> usize {
168 let left_node = &self.nodes[left];
169 let right_node = &self.nodes[right];
170 let new_bounds = left_node.bounds().union(right_node.bounds());
171 let max_ordering = cmp::max(left_node.max_ordering(), right_node.max_ordering());
172 self.nodes.push(Node::Internal {
173 bounds: new_bounds,
174 left,
175 right,
176 max_order: max_ordering,
177 });
178 self.nodes.len() - 1
179 }
180}
181
182impl<U> Default for BoundsTree<U>
183where
184 U: Clone + Debug + Default + PartialEq,
185{
186 fn default() -> Self {
187 BoundsTree {
188 root: None,
189 nodes: Vec::new(),
190 stack: Vec::new(),
191 }
192 }
193}
194
195#[derive(Debug, Clone)]
196enum Node<U>
197where
198 U: Clone + Debug + Default + PartialEq,
199{
200 Leaf {
201 bounds: Bounds<U>,
202 order: u32,
203 },
204 Internal {
205 left: usize,
206 right: usize,
207 bounds: Bounds<U>,
208 max_order: u32,
209 },
210}
211
212impl<U> Node<U>
213where
214 U: Clone + Debug + Default + PartialEq,
215{
216 fn bounds(&self) -> &Bounds<U> {
217 match self {
218 Node::Leaf { bounds, .. } => bounds,
219 Node::Internal { bounds, .. } => bounds,
220 }
221 }
222
223 fn max_ordering(&self) -> u32 {
224 match self {
225 Node::Leaf {
226 order: ordering, ..
227 } => *ordering,
228 Node::Internal {
229 max_order: max_ordering,
230 ..
231 } => *max_ordering,
232 }
233 }
234}
235
236#[cfg(test)]
237mod tests {
238 use super::*;
239 use crate::{Bounds, Point, Size};
240
241 #[test]
242 fn test_insert() {
243 let mut tree = BoundsTree::<f32>::default();
244 let bounds1 = Bounds {
245 origin: Point { x: 0.0, y: 0.0 },
246 size: Size {
247 width: 10.0,
248 height: 10.0,
249 },
250 };
251 let bounds2 = Bounds {
252 origin: Point { x: 5.0, y: 5.0 },
253 size: Size {
254 width: 10.0,
255 height: 10.0,
256 },
257 };
258 let bounds3 = Bounds {
259 origin: Point { x: 10.0, y: 10.0 },
260 size: Size {
261 width: 10.0,
262 height: 10.0,
263 },
264 };
265
266 // Insert the bounds into the tree and verify the order is correct
267 assert_eq!(tree.insert(bounds1), 1);
268 assert_eq!(tree.insert(bounds2), 2);
269 assert_eq!(tree.insert(bounds3), 3);
270
271 // Insert non-overlapping bounds and verify they can reuse orders
272 let bounds4 = Bounds {
273 origin: Point { x: 20.0, y: 20.0 },
274 size: Size {
275 width: 10.0,
276 height: 10.0,
277 },
278 };
279 let bounds5 = Bounds {
280 origin: Point { x: 40.0, y: 40.0 },
281 size: Size {
282 width: 10.0,
283 height: 10.0,
284 },
285 };
286 let bounds6 = Bounds {
287 origin: Point { x: 25.0, y: 25.0 },
288 size: Size {
289 width: 10.0,
290 height: 10.0,
291 },
292 };
293 assert_eq!(tree.insert(bounds4), 1); // bounds4 does not overlap with bounds1, bounds2, or bounds3
294 assert_eq!(tree.insert(bounds5), 1); // bounds5 does not overlap with any other bounds
295 assert_eq!(tree.insert(bounds6), 2); // bounds6 overlaps with bounds4, so it should have a different order
296 }
297}