khora_core/physics/
dynamic_tree.rs

1// Copyright 2025 eraflo
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! # Dynamic AABB Tree
16//!
17//! A high-performance spatial partitioning structure for broad-phase collision detection.
18//! It uses an incrementally balanced binary tree of AABBs.
19
20use crate::math::Aabb;
21use crate::math::Vec3;
22
23const NULL_NODE: i32 = -1;
24
25/// A node in the dynamic tree.
26#[derive(Debug, Clone)]
27pub struct DynamicTreeNode<T: Clone> {
28    /// Enlarged AABB for this node.
29    pub aabb: Aabb,
30    /// User data (e.g., EntityId or index).
31    pub user_data: Option<T>,
32    /// Index of the parent node.
33    pub parent: i32,
34    /// Indices of child nodes (if internal).
35    pub children: [i32; 2],
36    /// Height of the node in the tree (0 for leaves).
37    pub height: i32,
38}
39
40impl<T: Clone> DynamicTreeNode<T> {
41    /// Returns true if this node is a leaf (has no children).
42    #[inline]
43    pub fn is_leaf(&self) -> bool {
44        self.children[0] == NULL_NODE
45    }
46}
47
48/// A dynamic AABB tree for efficient spatial queries.
49#[derive(Debug, Clone)]
50pub struct DynamicTree<T: Clone> {
51    root: i32,
52    nodes: Vec<DynamicTreeNode<T>>,
53    free_list: i32,
54    node_count: usize,
55}
56
57impl<T: Clone> Default for DynamicTree<T> {
58    fn default() -> Self {
59        Self::new()
60    }
61}
62
63impl<T: Clone> DynamicTree<T> {
64    /// Creates a new, empty dynamic tree.
65    pub fn new() -> Self {
66        Self {
67            root: NULL_NODE,
68            nodes: Vec::new(),
69            free_list: NULL_NODE,
70            node_count: 0,
71        }
72    }
73
74    /// Inserts a new leaf into the tree and returns its index.
75    pub fn insert(&mut self, aabb: Aabb, user_data: T) -> i32 {
76        let leaf = self.allocate_node();
77        self.nodes[leaf as usize].aabb = aabb;
78        self.nodes[leaf as usize].user_data = Some(user_data);
79        self.nodes[leaf as usize].height = 0;
80
81        self.insert_leaf(leaf);
82        leaf
83    }
84
85    /// Removes a leaf from the tree.
86    pub fn remove(&mut self, leaf: i32) {
87        debug_assert!(leaf != NULL_NODE);
88        debug_assert!(self.nodes[leaf as usize].is_leaf());
89
90        self.remove_leaf(leaf);
91        self.deallocate_node(leaf);
92    }
93
94    /// Updates a leaf with a new AABB.
95    /// If the new AABB is still within the "fat" AABB, no tree update is needed.
96    pub fn update(
97        &mut self,
98        leaf: i32,
99        aabb: Aabb,
100        displacement: Vec3,
101        force_update: bool,
102    ) -> bool {
103        debug_assert!(leaf != NULL_NODE);
104        debug_assert!(self.nodes[leaf as usize].is_leaf());
105
106        if !force_update && self.nodes[leaf as usize].aabb.contains_aabb(&aabb) {
107            return false;
108        }
109
110        self.remove_leaf(leaf);
111
112        // Fatten the AABB
113        let mut fat_aabb = aabb;
114        let extension = Vec3::ONE * 0.1; // Magic constant for fattening
115        fat_aabb.min = fat_aabb.min - extension;
116        fat_aabb.max = fat_aabb.max + extension;
117
118        // Predictive encoding: add displacement to the AABB in the direction of travel
119        if displacement.x < 0.0 {
120            fat_aabb.min.x += displacement.x * 2.0;
121        } else {
122            fat_aabb.max.x += displacement.x * 2.0;
123        }
124        if displacement.y < 0.0 {
125            fat_aabb.min.y += displacement.y * 2.0;
126        } else {
127            fat_aabb.max.y += displacement.y * 2.0;
128        }
129        if displacement.z < 0.0 {
130            fat_aabb.min.z += displacement.z * 2.0;
131        } else {
132            fat_aabb.max.z += displacement.z * 2.0;
133        }
134
135        self.nodes[leaf as usize].aabb = fat_aabb;
136        self.insert_leaf(leaf);
137        true
138    }
139
140    /// Returns the user data for a given leaf.
141    pub fn get_user_data(&self, leaf: i32) -> &T {
142        self.nodes[leaf as usize].user_data.as_ref().unwrap()
143    }
144
145    /// Queries the tree for all pairs of overlapping leaves.
146    pub fn query_pairs<F>(&self, mut callback: F)
147    where
148        F: FnMut(&T, &T),
149    {
150        if self.root == NULL_NODE {
151            return;
152        }
153
154        // Simultaneous tree traversal for all pairs.
155        let mut stack: Vec<i32> = Vec::with_capacity(256);
156        stack.push(self.root);
157
158        while let Some(node_id) = stack.pop() {
159            let node = &self.nodes[node_id as usize];
160            if node.height < 1 {
161                continue;
162            }
163
164            let left = node.children[0];
165            let right = node.children[1];
166
167            self.query_overlap(left, right, &mut callback);
168
169            stack.push(left);
170            stack.push(right);
171        }
172    }
173
174    fn query_overlap<F>(&self, i_a: i32, i_b: i32, callback: &mut F)
175    where
176        F: FnMut(&T, &T),
177    {
178        let node_a = &self.nodes[i_a as usize];
179        let node_b = &self.nodes[i_b as usize];
180
181        if !node_a.aabb.intersects_aabb(&node_b.aabb) {
182            return;
183        }
184
185        if node_a.is_leaf() && node_b.is_leaf() {
186            if i_a != i_b {
187                callback(
188                    node_a.user_data.as_ref().unwrap(),
189                    node_b.user_data.as_ref().unwrap(),
190                );
191            }
192        } else if node_a.is_leaf() {
193            self.query_overlap(i_a, node_b.children[0], callback);
194            self.query_overlap(i_a, node_b.children[1], callback);
195        } else if node_b.is_leaf() {
196            self.query_overlap(node_a.children[0], i_b, callback);
197            self.query_overlap(node_a.children[1], i_b, callback);
198        } else {
199            self.query_overlap(node_a.children[0], node_b.children[0], callback);
200            self.query_overlap(node_a.children[0], node_b.children[1], callback);
201            self.query_overlap(node_a.children[1], node_b.children[0], callback);
202            self.query_overlap(node_a.children[1], node_b.children[1], callback);
203        }
204    }
205
206    /// Queries the tree for leaves overlapping the given AABB.
207    pub fn query<F>(&self, aabb: &Aabb, mut callback: F)
208    where
209        F: FnMut(&T) -> bool,
210    {
211        if self.root == NULL_NODE {
212            return;
213        }
214
215        let mut stack = Vec::with_capacity(64);
216        stack.push(self.root);
217
218        while let Some(node_id) = stack.pop() {
219            let node = &self.nodes[node_id as usize];
220            if node.aabb.intersects_aabb(aabb) {
221                if node.is_leaf() {
222                    let proceed = callback(node.user_data.as_ref().unwrap());
223                    if !proceed {
224                        return;
225                    }
226                } else {
227                    stack.push(node.children[0]);
228                    stack.push(node.children[1]);
229                }
230            }
231        }
232    }
233
234    // --- Internal Leaf Management ---
235
236    fn insert_leaf(&mut self, leaf: i32) {
237        if self.root == NULL_NODE {
238            self.root = leaf;
239            self.nodes[self.root as usize].parent = NULL_NODE;
240            return;
241        }
242
243        // Find the best sibling for the new leaf
244        let leaf_aabb = self.nodes[leaf as usize].aabb;
245        let mut index = self.root;
246        while !self.nodes[index as usize].is_leaf() {
247            let node = &self.nodes[index as usize];
248            let child1 = node.children[0];
249            let child2 = node.children[1];
250
251            let area = node.aabb.surface_area();
252            let combined_aabb = node.aabb.merge(&leaf_aabb);
253            let combined_area = combined_aabb.surface_area();
254
255            // Cost of creating a new parent for this node and the new leaf
256            let cost = 2.0 * combined_area;
257
258            // Minimum cost of pushing the leaf further down the tree
259            let inheritance_cost = 2.0 * (combined_area - area);
260
261            // Cost of descending into child1
262            let cost1 = {
263                let child1_node = &self.nodes[child1 as usize];
264                let new_area = child1_node.aabb.merge(&leaf_aabb).surface_area();
265                if child1_node.is_leaf() {
266                    new_area + inheritance_cost
267                } else {
268                    let old_area = child1_node.aabb.surface_area();
269                    (new_area - old_area) + inheritance_cost
270                }
271            };
272
273            // Cost of descending into child2
274            let cost2 = {
275                let child2_node = &self.nodes[child2 as usize];
276                let new_area = child2_node.aabb.merge(&leaf_aabb).surface_area();
277                if child2_node.is_leaf() {
278                    new_area + inheritance_cost
279                } else {
280                    let old_area = child2_node.aabb.surface_area();
281                    (new_area - old_area) + inheritance_cost
282                }
283            };
284
285            // Descend according to the minimum cost
286            if cost < cost1 && cost < cost2 {
287                break;
288            }
289
290            index = if cost1 < cost2 { child1 } else { child2 };
291        }
292
293        let sibling = index;
294
295        // Create a new parent
296        let old_parent = self.nodes[sibling as usize].parent;
297        let new_parent = self.allocate_node();
298        self.nodes[new_parent as usize].parent = old_parent;
299        self.nodes[new_parent as usize].user_data = None;
300        self.nodes[new_parent as usize].aabb = leaf_aabb.merge(&self.nodes[sibling as usize].aabb);
301        self.nodes[new_parent as usize].height = self.nodes[sibling as usize].height + 1;
302
303        if old_parent != NULL_NODE {
304            // Sibling was not root
305            if self.nodes[old_parent as usize].children[0] == sibling {
306                self.nodes[old_parent as usize].children[0] = new_parent;
307            } else {
308                self.nodes[old_parent as usize].children[1] = new_parent;
309            }
310            self.nodes[new_parent as usize].children[0] = sibling;
311            self.nodes[new_parent as usize].children[1] = leaf;
312            self.nodes[sibling as usize].parent = new_parent;
313            self.nodes[leaf as usize].parent = new_parent;
314        } else {
315            // Sibling was root
316            self.nodes[new_parent as usize].children[0] = sibling;
317            self.nodes[new_parent as usize].children[1] = leaf;
318            self.nodes[sibling as usize].parent = new_parent;
319            self.nodes[leaf as usize].parent = new_parent;
320            self.root = new_parent;
321        }
322
323        // Walk back up the tree fixing heights and AABBs
324        index = self.nodes[leaf as usize].parent;
325        while index != NULL_NODE {
326            index = self.balance(index);
327
328            let child1 = self.nodes[index as usize].children[0];
329            let child2 = self.nodes[index as usize].children[1];
330
331            debug_assert!(child1 != NULL_NODE);
332            debug_assert!(child2 != NULL_NODE);
333
334            self.nodes[index as usize].height = 1 + self.nodes[child1 as usize]
335                .height
336                .max(self.nodes[child2 as usize].height);
337            self.nodes[index as usize].aabb = self.nodes[child1 as usize]
338                .aabb
339                .merge(&self.nodes[child2 as usize].aabb);
340
341            index = self.nodes[index as usize].parent;
342        }
343    }
344
345    fn remove_leaf(&mut self, leaf: i32) {
346        if leaf == self.root {
347            self.root = NULL_NODE;
348            return;
349        }
350
351        let parent = self.nodes[leaf as usize].parent;
352        let grand_parent = self.nodes[parent as usize].parent;
353        let sibling = if self.nodes[parent as usize].children[0] == leaf {
354            self.nodes[parent as usize].children[1]
355        } else {
356            self.nodes[parent as usize].children[0]
357        };
358
359        if grand_parent != NULL_NODE {
360            // Destroy parent and connect sibling to grandParent
361            if self.nodes[grand_parent as usize].children[0] == parent {
362                self.nodes[grand_parent as usize].children[0] = sibling;
363            } else {
364                self.nodes[grand_parent as usize].children[1] = sibling;
365            }
366            self.nodes[sibling as usize].parent = grand_parent;
367            self.deallocate_node(parent);
368
369            // Adjust ancestor bounds
370            let mut index = grand_parent;
371            while index != NULL_NODE {
372                index = self.balance(index);
373
374                let child1 = self.nodes[index as usize].children[0];
375                let child2 = self.nodes[index as usize].children[1];
376
377                self.nodes[index as usize].aabb = self.nodes[child1 as usize]
378                    .aabb
379                    .merge(&self.nodes[child2 as usize].aabb);
380                self.nodes[index as usize].height = 1 + self.nodes[child1 as usize]
381                    .height
382                    .max(self.nodes[child2 as usize].height);
383
384                index = self.nodes[index as usize].parent;
385            }
386        } else {
387            self.root = sibling;
388            self.nodes[sibling as usize].parent = NULL_NODE;
389            self.deallocate_node(parent);
390        }
391    }
392
393    // --- Node Allocation ---
394
395    fn allocate_node(&mut self) -> i32 {
396        if self.free_list != NULL_NODE {
397            let index = self.free_list;
398            self.free_list = self.nodes[index as usize].parent;
399            self.nodes[index as usize].parent = NULL_NODE;
400            self.nodes[index as usize].children = [NULL_NODE, NULL_NODE];
401            self.nodes[index as usize].height = 0;
402            self.node_count += 1;
403            index
404        } else {
405            let index = self.nodes.len() as i32;
406            self.nodes.push(DynamicTreeNode {
407                aabb: Aabb::INVALID,
408                user_data: None,
409                parent: NULL_NODE,
410                children: [NULL_NODE, NULL_NODE],
411                height: 0,
412            });
413            self.node_count += 1;
414            index
415        }
416    }
417
418    fn deallocate_node(&mut self, index: i32) {
419        debug_assert!(index != NULL_NODE);
420        self.nodes[index as usize].parent = self.free_list;
421        self.nodes[index as usize].user_data = None;
422        self.free_list = index;
423        self.node_count -= 1;
424    }
425
426    // --- Balancing (Tree Rotations) ---
427
428    fn balance(&mut self, i_a: i32) -> i32 {
429        debug_assert!(i_a != NULL_NODE);
430
431        let node_a = &self.nodes[i_a as usize];
432        if node_a.is_leaf() || node_a.height < 2 {
433            return i_a;
434        }
435
436        let i_b = node_a.children[0];
437        let i_c = node_a.children[1];
438
439        let balance = self.nodes[i_c as usize].height - self.nodes[i_b as usize].height;
440
441        // Rotate C up
442        if balance > 1 {
443            let i_f = self.nodes[i_c as usize].children[0];
444            let i_g = self.nodes[i_c as usize].children[1];
445
446            // Swap A and C
447            self.nodes[i_c as usize].children[0] = i_a;
448            self.nodes[i_c as usize].parent = self.nodes[i_a as usize].parent;
449            self.nodes[i_a as usize].parent = i_c;
450
451            // A's old parent should point to C
452            if self.nodes[i_c as usize].parent != NULL_NODE {
453                let p = self.nodes[i_c as usize].parent;
454                if self.nodes[p as usize].children[0] == i_a {
455                    self.nodes[p as usize].children[0] = i_c;
456                } else {
457                    self.nodes[p as usize].children[1] = i_c;
458                }
459            } else {
460                self.root = i_c;
461            }
462
463            // Rotate
464            if self.nodes[i_f as usize].height > self.nodes[i_g as usize].height {
465                self.nodes[i_c as usize].children[1] = i_f;
466                self.nodes[i_a as usize].children[1] = i_g;
467                self.nodes[i_g as usize].parent = i_a;
468
469                self.update_node_meta(i_a);
470                self.update_node_meta(i_c);
471            } else {
472                self.nodes[i_c as usize].children[1] = i_g;
473                self.nodes[i_a as usize].children[1] = i_f;
474                self.nodes[i_f as usize].parent = i_a;
475
476                self.update_node_meta(i_a);
477                self.update_node_meta(i_c);
478            }
479
480            return i_c;
481        }
482
483        // Rotate B up
484        if balance < -1 {
485            let i_d = self.nodes[i_b as usize].children[0];
486            let i_e = self.nodes[i_b as usize].children[1];
487
488            // Swap A and B
489            self.nodes[i_b as usize].children[0] = i_a;
490            self.nodes[i_b as usize].parent = self.nodes[i_a as usize].parent;
491            self.nodes[i_a as usize].parent = i_b;
492
493            // A's old parent should point to B
494            if self.nodes[i_b as usize].parent != NULL_NODE {
495                let p = self.nodes[i_b as usize].parent;
496                if self.nodes[p as usize].children[0] == i_a {
497                    self.nodes[p as usize].children[0] = i_b;
498                } else {
499                    self.nodes[p as usize].children[1] = i_b;
500                }
501            } else {
502                self.root = i_b;
503            }
504
505            // Rotate
506            if self.nodes[i_d as usize].height > self.nodes[i_e as usize].height {
507                self.nodes[i_b as usize].children[1] = i_d;
508                self.nodes[i_a as usize].children[0] = i_e;
509                self.nodes[i_e as usize].parent = i_a;
510
511                self.update_node_meta(i_a);
512                self.update_node_meta(i_b);
513            } else {
514                self.nodes[i_b as usize].children[1] = i_e;
515                self.nodes[i_a as usize].children[0] = i_d;
516                self.nodes[i_d as usize].parent = i_a;
517
518                self.update_node_meta(i_a);
519                self.update_node_meta(i_b);
520            }
521
522            return i_b;
523        }
524
525        i_a
526    }
527
528    fn update_node_meta(&mut self, index: i32) {
529        let child1 = self.nodes[index as usize].children[0];
530        let child2 = self.nodes[index as usize].children[1];
531        self.nodes[index as usize].aabb = self.nodes[child1 as usize]
532            .aabb
533            .merge(&self.nodes[child2 as usize].aabb);
534        self.nodes[index as usize].height = 1 + self.nodes[child1 as usize]
535            .height
536            .max(self.nodes[child2 as usize].height);
537    }
538}
539
540/// Iterator for traversing all user data in the tree's leaves.
541pub struct DynamicTreeIterator<'a, T: Clone> {
542    tree: &'a DynamicTree<T>,
543    stack: Vec<i32>,
544}
545
546impl<'a, T: Clone> Iterator for DynamicTreeIterator<'a, T> {
547    type Item = &'a T;
548
549    fn next(&mut self) -> Option<Self::Item> {
550        while let Some(node_id) = self.stack.pop() {
551            let node = &self.tree.nodes[node_id as usize];
552            if node.is_leaf() {
553                return Some(node.user_data.as_ref().unwrap());
554            } else {
555                self.stack.push(node.children[0]);
556                self.stack.push(node.children[1]);
557            }
558        }
559        None
560    }
561}
khora_core/physics/dynamic_tree.rs

khora_core/physics/
dynamic_tree.rs
