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