1use crate::collider::AABB;
15use crate::math::{Fix128, Vec3Fix};
16use crate::metric::MetricWeights;
17
18#[cfg(not(feature = "std"))]
19use alloc::vec::Vec;
20
21pub(crate) const NULL_NODE: u32 = u32::MAX;
23
24const FAT_MARGIN: Fix128 = Fix128 {
26 hi: 0,
27 lo: 0x8000000000000000,
28}; #[derive(Clone, Debug)]
36pub(crate) struct DynamicNode {
37 pub aabb: AABB,
39 pub parent: u32,
41 pub left: u32,
43 pub right: u32,
45 pub height: i32,
47 pub user_data: u32,
49 pub is_leaf: bool,
51}
52
53impl DynamicNode {
54 const fn new_leaf(aabb: AABB, user_data: u32) -> Self {
55 Self {
56 aabb,
57 parent: NULL_NODE,
58 left: NULL_NODE,
59 right: NULL_NODE,
60 height: 0,
61 user_data,
62 is_leaf: true,
63 }
64 }
65
66 const fn new_internal() -> Self {
67 Self {
68 aabb: AABB::new(Vec3Fix::ZERO, Vec3Fix::ZERO),
69 parent: NULL_NODE,
70 left: NULL_NODE,
71 right: NULL_NODE,
72 height: 0,
73 user_data: NULL_NODE,
74 is_leaf: false,
75 }
76 }
77}
78
79pub struct DynamicAabbTree {
81 pub(crate) nodes: Vec<DynamicNode>,
83 pub(crate) free_list: Vec<u32>,
85 pub(crate) root: u32,
87 pub margin: Fix128,
89 pub metric: MetricWeights,
97}
98
99impl DynamicAabbTree {
100 #[must_use]
102 pub const fn new() -> Self {
103 Self {
104 nodes: Vec::new(),
105 free_list: Vec::new(),
106 root: NULL_NODE,
107 margin: FAT_MARGIN,
108 metric: MetricWeights::L2,
109 }
110 }
111
112 pub fn insert(&mut self, aabb: AABB, user_data: u32) -> u32 {
114 let fat_aabb = self.fatten(aabb);
115 let node_id = self.alloc_node();
116
117 self.nodes[node_id as usize] = DynamicNode::new_leaf(fat_aabb, user_data);
118
119 self.insert_leaf(node_id);
120 node_id
121 }
122
123 pub fn remove(&mut self, proxy_id: u32) {
125 if proxy_id as usize >= self.nodes.len() {
126 return;
127 }
128 self.remove_leaf(proxy_id);
129 self.free_node(proxy_id);
130 }
131
132 pub fn update(&mut self, proxy_id: u32, new_aabb: AABB) -> bool {
136 if proxy_id as usize >= self.nodes.len() {
137 return false;
138 }
139
140 let fat = &self.nodes[proxy_id as usize].aabb;
141 if fat.min.x <= new_aabb.min.x
143 && fat.min.y <= new_aabb.min.y
144 && fat.min.z <= new_aabb.min.z
145 && fat.max.x >= new_aabb.max.x
146 && fat.max.y >= new_aabb.max.y
147 && fat.max.z >= new_aabb.max.z
148 {
149 return false; }
151
152 self.remove_leaf(proxy_id);
153 self.nodes[proxy_id as usize].aabb = self.fatten(new_aabb);
154 self.insert_leaf(proxy_id);
155 true
156 }
157
158 #[inline]
160 #[must_use]
161 pub fn user_data(&self, proxy_id: u32) -> u32 {
162 self.nodes[proxy_id as usize].user_data
163 }
164
165 #[inline]
167 #[must_use]
168 pub fn get_aabb(&self, proxy_id: u32) -> AABB {
169 self.nodes[proxy_id as usize].aabb
170 }
171
172 #[must_use]
174 pub fn query(&self, aabb: &AABB) -> Vec<u32> {
175 let mut result = Vec::new();
176 if self.root == NULL_NODE {
177 return result;
178 }
179
180 let mut stack = Vec::with_capacity(64);
181 stack.push(self.root);
182
183 while let Some(node_id) = stack.pop() {
184 if node_id == NULL_NODE {
185 continue;
186 }
187
188 let node = &self.nodes[node_id as usize];
189 if !node.aabb.intersects(aabb) {
190 continue;
191 }
192
193 if node.is_leaf {
194 result.push(node.user_data);
195 } else {
196 stack.push(node.left);
197 stack.push(node.right);
198 }
199 }
200
201 result
202 }
203
204 pub fn query_callback<F: FnMut(u32)>(&self, aabb: &AABB, mut callback: F) {
206 if self.root == NULL_NODE {
207 return;
208 }
209
210 let mut stack = Vec::with_capacity(64);
211 stack.push(self.root);
212
213 while let Some(node_id) = stack.pop() {
214 if node_id == NULL_NODE {
215 continue;
216 }
217
218 let node = &self.nodes[node_id as usize];
219 if !node.aabb.intersects(aabb) {
220 continue;
221 }
222
223 if node.is_leaf {
224 callback(node.user_data);
225 } else {
226 stack.push(node.left);
227 stack.push(node.right);
228 }
229 }
230 }
231
232 #[must_use]
234 pub fn find_pairs(&self) -> Vec<(u32, u32)> {
235 let mut pairs = Vec::new();
236
237 if self.root == NULL_NODE {
238 return pairs;
239 }
240
241 let mut leaves = Vec::new();
243 self.collect_leaves(self.root, &mut leaves);
244
245 for &leaf_id in &leaves {
247 let aabb = &self.nodes[leaf_id as usize].aabb;
248 let ud_a = self.nodes[leaf_id as usize].user_data;
249
250 self.query_callback(aabb, |ud_b| {
251 if ud_a < ud_b {
252 pairs.push((ud_a, ud_b));
253 }
254 });
255 }
256
257 pairs.sort_unstable();
258 pairs.dedup();
259 pairs
260 }
261
262 #[must_use]
264 pub fn proxy_count(&self) -> usize {
265 self.nodes
266 .iter()
267 .filter(|n| n.is_leaf && n.user_data != NULL_NODE)
268 .count()
269 }
270
271 #[inline]
273 #[must_use]
274 pub fn node_count(&self) -> usize {
275 self.nodes.len() - self.free_list.len()
276 }
277
278 #[must_use]
280 pub fn height(&self) -> i32 {
281 if self.root == NULL_NODE {
282 0
283 } else {
284 self.nodes[self.root as usize].height
285 }
286 }
287
288 fn fatten(&self, aabb: AABB) -> AABB {
291 let pad = self.metric.euclidean_radius(self.margin);
293 let m = Vec3Fix::new(pad, pad, pad);
294 AABB::new(aabb.min - m, aabb.max + m)
295 }
296
297 fn alloc_node(&mut self) -> u32 {
298 if let Some(id) = self.free_list.pop() {
299 id
300 } else {
301 let id = self.nodes.len() as u32;
302 self.nodes.push(DynamicNode::new_internal());
303 id
304 }
305 }
306
307 fn free_node(&mut self, node_id: u32) {
308 self.nodes[node_id as usize].height = -1;
309 self.nodes[node_id as usize].user_data = NULL_NODE;
310 self.nodes[node_id as usize].is_leaf = false;
311 self.nodes[node_id as usize].left = NULL_NODE;
312 self.nodes[node_id as usize].right = NULL_NODE;
313 self.nodes[node_id as usize].parent = NULL_NODE;
314 self.free_list.push(node_id);
315 }
316
317 fn insert_leaf(&mut self, leaf: u32) {
318 if self.root == NULL_NODE {
319 self.root = leaf;
320 self.nodes[leaf as usize].parent = NULL_NODE;
321 return;
322 }
323
324 let leaf_aabb = self.nodes[leaf as usize].aabb;
326 let mut sibling = self.root;
327
328 while !self.nodes[sibling as usize].is_leaf {
329 let left = self.nodes[sibling as usize].left;
330 let right = self.nodes[sibling as usize].right;
331
332 let area = self.nodes[sibling as usize].aabb.surface_area();
333 let combined = leaf_aabb.union(&self.nodes[sibling as usize].aabb);
334 let combined_area = combined.surface_area();
335
336 let cost = combined_area * Fix128::from_int(2);
337 let inheritance_cost = (combined_area - area) * Fix128::from_int(2);
338
339 let cost_left = self.child_insertion_cost(left, &leaf_aabb, inheritance_cost);
340 let cost_right = self.child_insertion_cost(right, &leaf_aabb, inheritance_cost);
341
342 if cost < cost_left && cost < cost_right {
343 break;
344 }
345
346 sibling = if cost_left < cost_right { left } else { right };
347 }
348
349 let old_parent = self.nodes[sibling as usize].parent;
351 let new_parent = self.alloc_node();
352 self.nodes[new_parent as usize] = DynamicNode::new_internal();
353 self.nodes[new_parent as usize].parent = old_parent;
354 self.nodes[new_parent as usize].aabb = leaf_aabb.union(&self.nodes[sibling as usize].aabb);
355 self.nodes[new_parent as usize].height = self.nodes[sibling as usize].height + 1;
356
357 if old_parent == NULL_NODE {
358 self.root = new_parent;
359 } else if self.nodes[old_parent as usize].left == sibling {
360 self.nodes[old_parent as usize].left = new_parent;
361 } else {
362 self.nodes[old_parent as usize].right = new_parent;
363 }
364
365 self.nodes[new_parent as usize].left = sibling;
366 self.nodes[new_parent as usize].right = leaf;
367 self.nodes[sibling as usize].parent = new_parent;
368 self.nodes[leaf as usize].parent = new_parent;
369
370 self.fix_upwards(new_parent);
372 }
373
374 fn child_insertion_cost(&self, child: u32, leaf_aabb: &AABB, inheritance: Fix128) -> Fix128 {
375 let combined = leaf_aabb.union(&self.nodes[child as usize].aabb);
376 if self.nodes[child as usize].is_leaf {
377 combined.surface_area() + inheritance
378 } else {
379 let old_area = self.nodes[child as usize].aabb.surface_area();
380 let new_area = combined.surface_area();
381 (new_area - old_area) + inheritance
382 }
383 }
384
385 fn remove_leaf(&mut self, leaf: u32) {
386 if leaf == self.root {
387 self.root = NULL_NODE;
388 return;
389 }
390
391 let parent = self.nodes[leaf as usize].parent;
392 let grand_parent = self.nodes[parent as usize].parent;
393 let sibling = if self.nodes[parent as usize].left == leaf {
394 self.nodes[parent as usize].right
395 } else {
396 self.nodes[parent as usize].left
397 };
398
399 if grand_parent == NULL_NODE {
400 self.root = sibling;
401 self.nodes[sibling as usize].parent = NULL_NODE;
402 self.free_node(parent);
403 } else {
404 if self.nodes[grand_parent as usize].left == parent {
406 self.nodes[grand_parent as usize].left = sibling;
407 } else {
408 self.nodes[grand_parent as usize].right = sibling;
409 }
410 self.nodes[sibling as usize].parent = grand_parent;
411 self.free_node(parent);
412
413 self.fix_upwards(grand_parent);
414 }
415 }
416
417 fn fix_upwards(&mut self, start: u32) {
418 let mut node_id = start;
419 while node_id != NULL_NODE {
420 node_id = self.balance(node_id);
421
422 let left = self.nodes[node_id as usize].left;
423 let right = self.nodes[node_id as usize].right;
424
425 if left != NULL_NODE && right != NULL_NODE {
426 let lh = self.nodes[left as usize].height;
427 let rh = self.nodes[right as usize].height;
428 self.nodes[node_id as usize].height = 1 + if lh > rh { lh } else { rh };
429 self.nodes[node_id as usize].aabb = self.nodes[left as usize]
430 .aabb
431 .union(&self.nodes[right as usize].aabb);
432 }
433
434 node_id = self.nodes[node_id as usize].parent;
435 }
436 }
437
438 fn balance(&mut self, node_id: u32) -> u32 {
440 if self.nodes[node_id as usize].is_leaf || self.nodes[node_id as usize].height < 2 {
441 return node_id;
442 }
443
444 let left = self.nodes[node_id as usize].left;
445 let right = self.nodes[node_id as usize].right;
446
447 let balance_factor = self.nodes[right as usize].height - self.nodes[left as usize].height;
448
449 if balance_factor > 1 {
450 self.rotate_left(node_id)
451 } else if balance_factor < -1 {
452 self.rotate_right(node_id)
453 } else {
454 node_id
455 }
456 }
457
458 fn rotate_left(&mut self, node_id: u32) -> u32 {
459 let right = self.nodes[node_id as usize].right;
460 let right_left = self.nodes[right as usize].left;
461 let right_right = self.nodes[right as usize].right;
462 let parent = self.nodes[node_id as usize].parent;
463
464 self.nodes[right as usize].left = node_id;
466 self.nodes[right as usize].parent = parent;
467 self.nodes[node_id as usize].parent = right;
468
469 if parent == NULL_NODE {
470 self.root = right;
471 } else if self.nodes[parent as usize].left == node_id {
472 self.nodes[parent as usize].left = right;
473 } else {
474 self.nodes[parent as usize].right = right;
475 }
476
477 let rl_h = if right_left == NULL_NODE {
479 -1
480 } else {
481 self.nodes[right_left as usize].height
482 };
483 let rr_h = if right_right == NULL_NODE {
484 -1
485 } else {
486 self.nodes[right_right as usize].height
487 };
488
489 if rl_h > rr_h {
490 self.nodes[right as usize].right = right_left;
491 self.nodes[node_id as usize].right = right_right;
492 if right_right != NULL_NODE {
493 self.nodes[right_right as usize].parent = node_id;
494 }
495 if right_left != NULL_NODE {
496 self.nodes[right_left as usize].parent = right;
497 }
498 } else {
499 self.nodes[node_id as usize].right = right_left;
500 if right_left != NULL_NODE {
501 self.nodes[right_left as usize].parent = node_id;
502 }
503 }
504
505 let left = self.nodes[node_id as usize].left;
507 let nr = self.nodes[node_id as usize].right;
508 if left != NULL_NODE && nr != NULL_NODE {
509 self.nodes[node_id as usize].aabb = self.nodes[left as usize]
510 .aabb
511 .union(&self.nodes[nr as usize].aabb);
512 let lh = self.nodes[left as usize].height;
513 let rh = self.nodes[nr as usize].height;
514 self.nodes[node_id as usize].height = 1 + if lh > rh { lh } else { rh };
515 }
516
517 let rl = self.nodes[right as usize].left;
518 let rr = self.nodes[right as usize].right;
519 if rl != NULL_NODE && rr != NULL_NODE {
520 self.nodes[right as usize].aabb = self.nodes[rl as usize]
521 .aabb
522 .union(&self.nodes[rr as usize].aabb);
523 let lh = self.nodes[rl as usize].height;
524 let rh = self.nodes[rr as usize].height;
525 self.nodes[right as usize].height = 1 + if lh > rh { lh } else { rh };
526 }
527
528 right
529 }
530
531 fn rotate_right(&mut self, node_id: u32) -> u32 {
532 let left = self.nodes[node_id as usize].left;
533 let left_left = self.nodes[left as usize].left;
534 let left_right = self.nodes[left as usize].right;
535 let parent = self.nodes[node_id as usize].parent;
536
537 self.nodes[left as usize].right = node_id;
538 self.nodes[left as usize].parent = parent;
539 self.nodes[node_id as usize].parent = left;
540
541 if parent == NULL_NODE {
542 self.root = left;
543 } else if self.nodes[parent as usize].left == node_id {
544 self.nodes[parent as usize].left = left;
545 } else {
546 self.nodes[parent as usize].right = left;
547 }
548
549 let ll_h = if left_left == NULL_NODE {
550 -1
551 } else {
552 self.nodes[left_left as usize].height
553 };
554 let lr_h = if left_right == NULL_NODE {
555 -1
556 } else {
557 self.nodes[left_right as usize].height
558 };
559
560 if lr_h > ll_h {
561 self.nodes[left as usize].left = left_right;
562 self.nodes[node_id as usize].left = left_left;
563 if left_left != NULL_NODE {
564 self.nodes[left_left as usize].parent = node_id;
565 }
566 if left_right != NULL_NODE {
567 self.nodes[left_right as usize].parent = left;
568 }
569 } else {
570 self.nodes[node_id as usize].left = left_right;
571 if left_right != NULL_NODE {
572 self.nodes[left_right as usize].parent = node_id;
573 }
574 }
575
576 let nl = self.nodes[node_id as usize].left;
578 let right = self.nodes[node_id as usize].right;
579 if nl != NULL_NODE && right != NULL_NODE {
580 self.nodes[node_id as usize].aabb = self.nodes[nl as usize]
581 .aabb
582 .union(&self.nodes[right as usize].aabb);
583 let lh = self.nodes[nl as usize].height;
584 let rh = self.nodes[right as usize].height;
585 self.nodes[node_id as usize].height = 1 + if lh > rh { lh } else { rh };
586 }
587
588 let ll = self.nodes[left as usize].left;
589 let lr = self.nodes[left as usize].right;
590 if ll != NULL_NODE && lr != NULL_NODE {
591 self.nodes[left as usize].aabb = self.nodes[ll as usize]
592 .aabb
593 .union(&self.nodes[lr as usize].aabb);
594 let lh = self.nodes[ll as usize].height;
595 let rh = self.nodes[lr as usize].height;
596 self.nodes[left as usize].height = 1 + if lh > rh { lh } else { rh };
597 }
598
599 left
600 }
601
602 fn collect_leaves(&self, node_id: u32, leaves: &mut Vec<u32>) {
603 if node_id == NULL_NODE {
604 return;
605 }
606
607 if self.nodes[node_id as usize].is_leaf {
608 leaves.push(node_id);
609 } else {
610 self.collect_leaves(self.nodes[node_id as usize].left, leaves);
611 self.collect_leaves(self.nodes[node_id as usize].right, leaves);
612 }
613 }
614}
615
616impl Default for DynamicAabbTree {
617 fn default() -> Self {
618 Self::new()
619 }
620}
621
622#[cfg(test)]
623mod tests {
624 use super::*;
625
626 fn make_aabb(x: i64, y: i64, z: i64) -> AABB {
627 AABB::new(
628 Vec3Fix::from_int(x, y, z),
629 Vec3Fix::from_int(x + 1, y + 1, z + 1),
630 )
631 }
632
633 #[test]
634 fn test_insert_and_query() {
635 let mut tree = DynamicAabbTree::new();
636
637 let _p0 = tree.insert(make_aabb(0, 0, 0), 0);
638 let _p1 = tree.insert(make_aabb(10, 10, 10), 1);
639 let _p2 = tree.insert(make_aabb(20, 20, 20), 2);
640
641 assert_eq!(tree.proxy_count(), 3);
642
643 let results = tree.query(&make_aabb(-1, -1, -1));
645 assert!(results.contains(&0));
646 assert!(!results.contains(&2));
647
648 let all = tree.query(&AABB::new(
650 Vec3Fix::from_int(-100, -100, -100),
651 Vec3Fix::from_int(100, 100, 100),
652 ));
653 assert_eq!(all.len(), 3);
654 }
655
656 #[test]
657 fn test_remove() {
658 let mut tree = DynamicAabbTree::new();
659
660 let _p0 = tree.insert(make_aabb(0, 0, 0), 0);
661 let p1 = tree.insert(make_aabb(5, 5, 5), 1);
662 let _p2 = tree.insert(make_aabb(10, 10, 10), 2);
663
664 assert_eq!(tree.proxy_count(), 3);
665
666 tree.remove(p1);
667 assert_eq!(tree.proxy_count(), 2);
668
669 let all = tree.query(&AABB::new(
670 Vec3Fix::from_int(-100, -100, -100),
671 Vec3Fix::from_int(100, 100, 100),
672 ));
673 assert!(!all.contains(&1));
674 }
675
676 #[test]
677 fn test_update_no_reinsert() {
678 let mut tree = DynamicAabbTree::new();
679
680 let p0 = tree.insert(make_aabb(0, 0, 0), 0);
681
682 let tiny_move = AABB::new(
684 Vec3Fix::new(Fix128::from_ratio(1, 10), Fix128::ZERO, Fix128::ZERO),
685 Vec3Fix::new(
686 Fix128::ONE + Fix128::from_ratio(1, 10),
687 Fix128::ONE,
688 Fix128::ONE,
689 ),
690 );
691 let reinserted = tree.update(p0, tiny_move);
692 assert!(!reinserted, "Small move should not trigger reinsert");
693 }
694
695 #[test]
696 fn test_update_reinsert() {
697 let mut tree = DynamicAabbTree::new();
698
699 let p0 = tree.insert(make_aabb(0, 0, 0), 0);
700
701 let far_move = make_aabb(100, 100, 100);
703 let reinserted = tree.update(p0, far_move);
704 assert!(reinserted, "Large move should trigger reinsert");
705
706 let results = tree.query(&make_aabb(99, 99, 99));
708 assert!(results.contains(&0));
709 }
710
711 #[test]
712 fn test_find_pairs() {
713 let mut tree = DynamicAabbTree::new();
714
715 tree.insert(
717 AABB::new(Vec3Fix::from_int(0, 0, 0), Vec3Fix::from_int(2, 2, 2)),
718 0,
719 );
720 tree.insert(
721 AABB::new(Vec3Fix::from_int(1, 1, 1), Vec3Fix::from_int(3, 3, 3)),
722 1,
723 );
724 tree.insert(make_aabb(100, 100, 100), 2);
726
727 let pairs = tree.find_pairs();
728 assert!(
729 pairs.contains(&(0, 1)),
730 "Overlapping bodies should form a pair"
731 );
732 assert!(
733 !pairs.contains(&(0, 2)),
734 "Far body should not pair with body 0"
735 );
736 }
737
738 #[test]
739 fn test_tree_balance() {
740 let mut tree = DynamicAabbTree::new();
741
742 for i in 0..100 {
744 tree.insert(make_aabb(i * 3, 0, 0), i as u32);
745 }
746
747 assert_eq!(tree.proxy_count(), 100);
748 assert!(
750 tree.height() < 20,
751 "Tree should be balanced, height={}",
752 tree.height()
753 );
754 }
755
756 #[test]
757 fn test_empty_tree() {
758 let tree = DynamicAabbTree::new();
759 assert_eq!(tree.proxy_count(), 0);
760 assert_eq!(tree.height(), 0);
761 assert!(tree.query(&make_aabb(0, 0, 0)).is_empty());
762 assert!(tree.find_pairs().is_empty());
763 }
764
765 #[test]
766 fn get_aabb_returns_fattened_proxy_box_and_tracks_update() {
767 let mut tree = DynamicAabbTree::new();
768 tree.margin = Fix128::from_ratio(1, 2);
769 let m = Vec3Fix::new(tree.margin, tree.margin, tree.margin);
770
771 let tight0 = make_aabb(0, 0, 0);
773 let tight1 = AABB::new(Vec3Fix::from_int(-7, 3, -12), Vec3Fix::from_int(5, 9, -2));
774 let p0 = tree.insert(tight0, 10);
775 let p1 = tree.insert(tight1, 11);
776 assert_eq!(tree.get_aabb(p0), AABB::new(tight0.min - m, tight0.max + m));
777 assert_eq!(tree.get_aabb(p1), AABB::new(tight1.min - m, tight1.max + m));
778 assert_ne!(tree.get_aabb(p0), tree.get_aabb(p1));
780 assert_eq!(tree.user_data(p0), 10);
781 assert_eq!(tree.user_data(p1), 11);
782
783 let inside = AABB::new(
785 Vec3Fix::new(Fix128::from_ratio(1, 4), Fix128::ZERO, Fix128::ZERO),
786 Vec3Fix::new(Fix128::from_ratio(5, 4), Fix128::ONE, Fix128::ONE),
787 );
788 let before = tree.get_aabb(p0);
789 assert!(!tree.update(p0, inside));
790 assert_eq!(tree.get_aabb(p0), before);
791
792 let moved = make_aabb(30, 0, 0);
794 assert!(tree.update(p0, moved));
795 assert_eq!(tree.get_aabb(p0), AABB::new(moved.min - m, moved.max + m));
796 assert!(tree.query(&make_aabb(30, 0, 0)).contains(&10));
798 assert!(!tree.query(&make_aabb(0, 0, 0)).contains(&10));
799 let fat = tree.get_aabb(p0);
801 let touching = AABB::new(
802 Vec3Fix::new(fat.max.x, Fix128::ZERO, Fix128::ZERO),
803 Vec3Fix::new(fat.max.x + Fix128::ONE, Fix128::ONE, Fix128::ONE),
804 );
805 assert!(tree.query(&touching).contains(&10));
806
807 let mut exact = DynamicAabbTree::new();
809 exact.margin = Fix128::ZERO;
810 let q = exact.insert(tight1, 0);
811 assert_eq!(exact.get_aabb(q), tight1);
812 }
813}