1use crate::{Bounds, Half};
2use std::{
3 cmp,
4 fmt::Debug,
5 ops::{Add, Sub},
6};
7
8/// Maximum children per internal node (R-tree style branching factor).
9/// Higher values = shorter tree = fewer cache misses, but more work per node.
10const MAX_CHILDREN: usize = 12;
11
12/// A spatial tree optimized for finding maximum ordering among intersecting bounds.
13///
14/// This is an R-tree variant specifically designed for the use case of assigning
15/// z-order to overlapping UI elements. Key optimizations:
16/// - Tracks the leaf with global max ordering for O(1) fast-path queries
17/// - Uses higher branching factor (4) for lower tree height
18/// - Aggressive pruning during search based on max_order metadata
19#[derive(Debug)]
20pub(crate) struct BoundsTree<U>
21where
22 U: Clone + Debug + Default + PartialEq,
23{
24 /// All nodes stored contiguously for cache efficiency.
25 nodes: Vec<Node<U>>,
26 /// Index of the root node, if any.
27 root: Option<usize>,
28 /// Index of the leaf with the highest ordering (for fast-path lookups).
29 max_leaf: Option<usize>,
30 /// Reusable stack for tree traversal during insertion.
31 insert_path: Vec<usize>,
32 /// Reusable stack for search operations.
33 search_stack: Vec<usize>,
34}
35
36/// A node in the bounds tree.
37#[derive(Debug, Clone)]
38struct Node<U>
39where
40 U: Clone + Debug + Default + PartialEq,
41{
42 /// Bounding box containing this node and all descendants.
43 bounds: Bounds<U>,
44 /// Maximum ordering value in this subtree.
45 max_order: u32,
46 /// Node-specific data.
47 kind: NodeKind,
48}
49
50#[derive(Debug, Clone)]
51enum NodeKind {
52 /// Leaf node containing actual bounds data.
53 Leaf {
54 /// The ordering assigned to this bounds.
55 order: u32,
56 },
57 /// Internal node with children.
58 Internal {
59 /// Indices of child nodes (2 to MAX_CHILDREN).
60 children: NodeChildren,
61 },
62}
63
64/// Fixed-size array for child indices, avoiding heap allocation.
65#[derive(Debug, Clone)]
66struct NodeChildren {
67 // Keeps an invariant where the max order child is always at the end
68 indices: [usize; MAX_CHILDREN],
69 len: u8,
70}
71
72impl NodeChildren {
73 fn new() -> Self {
74 Self {
75 indices: [0; MAX_CHILDREN],
76 len: 0,
77 }
78 }
79
80 fn push(&mut self, index: usize) {
81 debug_assert!((self.len as usize) < MAX_CHILDREN);
82 self.indices[self.len as usize] = index;
83 self.len += 1;
84 }
85
86 fn len(&self) -> usize {
87 self.len as usize
88 }
89
90 fn as_slice(&self) -> &[usize] {
91 &self.indices[..self.len as usize]
92 }
93}
94
95impl<U> BoundsTree<U>
96where
97 U: Clone
98 + Debug
99 + PartialEq
100 + PartialOrd
101 + Add<U, Output = U>
102 + Sub<Output = U>
103 + Half
104 + Default,
105{
106 /// Clears all nodes from the tree.
107 pub fn clear(&mut self) {
108 self.nodes.clear();
109 self.root = None;
110 self.max_leaf = None;
111 self.insert_path.clear();
112 self.search_stack.clear();
113 }
114
115 /// Inserts bounds into the tree and returns its assigned ordering.
116 ///
117 /// The ordering is one greater than the maximum ordering of any
118 /// existing bounds that intersect with the new bounds.
119 pub fn insert(&mut self, new_bounds: Bounds<U>) -> u32 {
120 // Find maximum ordering among intersecting bounds
121 let max_intersecting = self.find_max_ordering(&new_bounds);
122 let ordering = max_intersecting + 1;
123
124 // Insert the new leaf
125 let new_leaf_idx = self.insert_leaf(new_bounds, ordering);
126
127 // Update max_leaf tracking
128 self.max_leaf = match self.max_leaf {
129 None => Some(new_leaf_idx),
130 Some(old_idx) if self.nodes[old_idx].max_order < ordering => Some(new_leaf_idx),
131 some => some,
132 };
133
134 ordering
135 }
136
137 /// Finds the maximum ordering among all bounds that intersect with the query.
138 fn find_max_ordering(&mut self, query: &Bounds<U>) -> u32 {
139 let Some(root_idx) = self.root else {
140 return 0;
141 };
142
143 // Fast path: check if the max-ordering leaf intersects
144 if let Some(max_idx) = self.max_leaf {
145 let max_node = &self.nodes[max_idx];
146 if query.intersects(&max_node.bounds) {
147 return max_node.max_order;
148 }
149 }
150
151 // Slow path: search the tree
152 self.search_stack.clear();
153 self.search_stack.push(root_idx);
154
155 let mut max_found = 0u32;
156
157 while let Some(node_idx) = self.search_stack.pop() {
158 let node = &self.nodes[node_idx];
159
160 // Pruning: skip if this subtree can't improve our result
161 if node.max_order <= max_found {
162 continue;
163 }
164
165 // Spatial pruning: skip if bounds don't intersect
166 if !query.intersects(&node.bounds) {
167 continue;
168 }
169
170 match &node.kind {
171 NodeKind::Leaf { order } => {
172 max_found = cmp::max(max_found, *order);
173 }
174 NodeKind::Internal { children } => {
175 // Children are maintained with highest max_order at the end.
176 // Push in forward order to highest (last) is popped first.
177 for &child_idx in children.as_slice() {
178 if self.nodes[child_idx].max_order > max_found {
179 self.search_stack.push(child_idx);
180 }
181 }
182 }
183 }
184 }
185
186 max_found
187 }
188
189 /// Inserts a leaf node with the given bounds and ordering.
190 /// Returns the index of the new leaf.
191 fn insert_leaf(&mut self, bounds: Bounds<U>, order: u32) -> usize {
192 let new_leaf_idx = self.nodes.len();
193 self.nodes.push(Node {
194 bounds: bounds.clone(),
195 max_order: order,
196 kind: NodeKind::Leaf { order },
197 });
198
199 let Some(root_idx) = self.root else {
200 // Tree is empty, new leaf becomes root
201 self.root = Some(new_leaf_idx);
202 return new_leaf_idx;
203 };
204
205 // If root is a leaf, create internal node with both
206 if matches!(self.nodes[root_idx].kind, NodeKind::Leaf { .. }) {
207 let root_bounds = self.nodes[root_idx].bounds.clone();
208 let root_order = self.nodes[root_idx].max_order;
209
210 let mut children = NodeChildren::new();
211 // Max end invariant
212 if order > root_order {
213 children.push(root_idx);
214 children.push(new_leaf_idx);
215 } else {
216 children.push(new_leaf_idx);
217 children.push(root_idx);
218 }
219
220 let new_root_idx = self.nodes.len();
221 self.nodes.push(Node {
222 bounds: root_bounds.union(&bounds),
223 max_order: cmp::max(root_order, order),
224 kind: NodeKind::Internal { children },
225 });
226 self.root = Some(new_root_idx);
227 return new_leaf_idx;
228 }
229
230 // Descend to find the best internal node to insert into
231 self.insert_path.clear();
232 let mut current_idx = root_idx;
233
234 loop {
235 let current = &self.nodes[current_idx];
236 let NodeKind::Internal { children } = ¤t.kind else {
237 unreachable!("Should only traverse internal nodes");
238 };
239
240 self.insert_path.push(current_idx);
241
242 // Find the best child to descend into
243 let mut best_child_idx = children.as_slice()[0];
244 let mut best_child_pos = 0;
245 let mut best_cost = bounds
246 .union(&self.nodes[best_child_idx].bounds)
247 .half_perimeter();
248
249 for (pos, &child_idx) in children.as_slice().iter().enumerate().skip(1) {
250 let cost = bounds.union(&self.nodes[child_idx].bounds).half_perimeter();
251 if cost < best_cost {
252 best_cost = cost;
253 best_child_idx = child_idx;
254 best_child_pos = pos;
255 }
256 }
257
258 // Check if best child is a leaf or internal
259 if matches!(self.nodes[best_child_idx].kind, NodeKind::Leaf { .. }) {
260 // Best child is a leaf. Check if current node has room for another child.
261 if children.len() < MAX_CHILDREN {
262 // Add new leaf directly to this node
263 let node = &mut self.nodes[current_idx];
264
265 if let NodeKind::Internal { children } = &mut node.kind {
266 children.push(new_leaf_idx);
267 // Swap new leaf only if it has the highest max_order
268 if order <= node.max_order {
269 let last = children.len() - 1;
270 children.indices.swap(last - 1, last);
271 }
272 }
273
274 node.bounds = node.bounds.union(&bounds);
275 node.max_order = cmp::max(node.max_order, order);
276 break;
277 } else {
278 // Node is full, create new internal with [best_leaf, new_leaf]
279 let sibling_bounds = self.nodes[best_child_idx].bounds.clone();
280 let sibling_order = self.nodes[best_child_idx].max_order;
281
282 let mut new_children = NodeChildren::new();
283 // Max end invariant
284 if order > sibling_order {
285 new_children.push(best_child_idx);
286 new_children.push(new_leaf_idx);
287 } else {
288 new_children.push(new_leaf_idx);
289 new_children.push(best_child_idx);
290 }
291
292 let new_internal_idx = self.nodes.len();
293 let new_internal_max = cmp::max(sibling_order, order);
294 self.nodes.push(Node {
295 bounds: sibling_bounds.union(&bounds),
296 max_order: new_internal_max,
297 kind: NodeKind::Internal {
298 children: new_children,
299 },
300 });
301
302 // Replace the leaf with the new internal in parent
303 let parent = &mut self.nodes[current_idx];
304 if let NodeKind::Internal { children } = &mut parent.kind {
305 let children_len = children.len();
306
307 children.indices[best_child_pos] = new_internal_idx;
308
309 // If new internal has highest max_order, swap it to the end
310 // to maintain sorting invariant
311 if new_internal_max > parent.max_order {
312 children.indices.swap(best_child_pos, children_len - 1);
313 }
314 }
315 break;
316 }
317 } else {
318 // Best child is internal, continue descent
319 current_idx = best_child_idx;
320 }
321 }
322
323 // Propagate bounds and max_order updates up the tree
324 let mut updated_child_idx = None;
325 for &node_idx in self.insert_path.iter().rev() {
326 let node = &mut self.nodes[node_idx];
327 node.bounds = node.bounds.union(&bounds);
328
329 if node.max_order < order {
330 node.max_order = order;
331
332 // Swap updated child to end (skip first iteration since the invariant is already handled by previous cases)
333 if let Some(child_idx) = updated_child_idx {
334 if let NodeKind::Internal { children } = &mut node.kind {
335 if let Some(pos) = children.as_slice().iter().position(|&c| c == child_idx)
336 {
337 let last = children.len() - 1;
338 if pos != last {
339 children.indices.swap(pos, last);
340 }
341 }
342 }
343 }
344 }
345
346 updated_child_idx = Some(node_idx);
347 }
348
349 new_leaf_idx
350 }
351}
352
353impl<U> Default for BoundsTree<U>
354where
355 U: Clone + Debug + Default + PartialEq,
356{
357 fn default() -> Self {
358 BoundsTree {
359 nodes: Vec::new(),
360 root: None,
361 max_leaf: None,
362 insert_path: Vec::new(),
363 search_stack: Vec::new(),
364 }
365 }
366}
367
368#[cfg(test)]
369mod tests {
370 use super::*;
371 use crate::{Bounds, Point, Size};
372 use rand::{Rng, SeedableRng};
373
374 #[test]
375 fn test_insert() {
376 let mut tree = BoundsTree::<f32>::default();
377 let bounds1 = Bounds {
378 origin: Point { x: 0.0, y: 0.0 },
379 size: Size {
380 width: 10.0,
381 height: 10.0,
382 },
383 };
384 let bounds2 = Bounds {
385 origin: Point { x: 5.0, y: 5.0 },
386 size: Size {
387 width: 10.0,
388 height: 10.0,
389 },
390 };
391 let bounds3 = Bounds {
392 origin: Point { x: 10.0, y: 10.0 },
393 size: Size {
394 width: 10.0,
395 height: 10.0,
396 },
397 };
398
399 // Insert the bounds into the tree and verify the order is correct
400 assert_eq!(tree.insert(bounds1), 1);
401 assert_eq!(tree.insert(bounds2), 2);
402 assert_eq!(tree.insert(bounds3), 3);
403
404 // Insert non-overlapping bounds and verify they can reuse orders
405 let bounds4 = Bounds {
406 origin: Point { x: 20.0, y: 20.0 },
407 size: Size {
408 width: 10.0,
409 height: 10.0,
410 },
411 };
412 let bounds5 = Bounds {
413 origin: Point { x: 40.0, y: 40.0 },
414 size: Size {
415 width: 10.0,
416 height: 10.0,
417 },
418 };
419 let bounds6 = Bounds {
420 origin: Point { x: 25.0, y: 25.0 },
421 size: Size {
422 width: 10.0,
423 height: 10.0,
424 },
425 };
426 assert_eq!(tree.insert(bounds4), 1); // bounds4 does not overlap with bounds1, bounds2, or bounds3
427 assert_eq!(tree.insert(bounds5), 1); // bounds5 does not overlap with any other bounds
428 assert_eq!(tree.insert(bounds6), 2); // bounds6 overlaps with bounds4, so it should have a different order
429 }
430
431 #[test]
432 fn test_random_iterations() {
433 let max_bounds = 100;
434 for seed in 1..=1000 {
435 // let seed = 44;
436 let mut tree = BoundsTree::default();
437 let mut rng = rand::rngs::StdRng::seed_from_u64(seed as u64);
438 let mut expected_quads: Vec<(Bounds<f32>, u32)> = Vec::new();
439
440 // Insert a random number of random AABBs into the tree.
441 let num_bounds = rng.random_range(1..=max_bounds);
442 for _ in 0..num_bounds {
443 let min_x: f32 = rng.random_range(-100.0..100.0);
444 let min_y: f32 = rng.random_range(-100.0..100.0);
445 let width: f32 = rng.random_range(0.0..50.0);
446 let height: f32 = rng.random_range(0.0..50.0);
447 let bounds = Bounds {
448 origin: Point { x: min_x, y: min_y },
449 size: Size { width, height },
450 };
451
452 let expected_ordering = expected_quads
453 .iter()
454 .filter_map(|quad| quad.0.intersects(&bounds).then_some(quad.1))
455 .max()
456 .unwrap_or(0)
457 + 1;
458 expected_quads.push((bounds, expected_ordering));
459
460 // Insert the AABB into the tree and collect intersections.
461 let actual_ordering = tree.insert(bounds);
462 assert_eq!(actual_ordering, expected_ordering);
463 }
464 }
465 }
466}