grid_tree/
tree.rs

1use crate::allocator::{AllocPtr, NodeAllocator, EMPTY_ALLOC_PTR};
2use crate::{BranchShape, ChildIndex, Level, SmallKeyHashMap, VectorKey};
3
4use smallvec::SmallVec;
5use std::collections::{hash_map, VecDeque};
6use std::hash::Hash;
7use std::marker::PhantomData;
8use std::mem;
9
10/// Uniquely identifies a node slot in the [`Tree`].
11#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
12pub struct NodeKey<V> {
13    pub level: Level,
14    pub coordinates: V,
15}
16
17impl<V> NodeKey<V> {
18    #[inline]
19    pub fn new(level: Level, coordinates: V) -> Self {
20        Self { level, coordinates }
21    }
22}
23
24/// Uniquely and stably identifies an occupied node in the [`Tree`] (until the node is removed).
25#[derive(Clone, Copy, Debug, Eq, PartialEq)]
26pub struct NodePtr {
27    pub(crate) level: Level,
28    pub(crate) alloc_ptr: AllocPtr,
29}
30
31impl NodePtr {
32    pub fn new(level: Level, alloc_ptr: AllocPtr) -> Self {
33        Self { level, alloc_ptr }
34    }
35
36    #[inline]
37    pub fn alloc_ptr(&self) -> AllocPtr {
38        self.alloc_ptr
39    }
40
41    #[inline]
42    pub fn level(&self) -> Level {
43        self.level
44    }
45
46    /// Null pointers are only used to indicate a child does not exist.
47    #[inline]
48    pub fn is_null(&self) -> bool {
49        self.alloc_ptr == EMPTY_ALLOC_PTR
50    }
51}
52
53/// Info about a node's parent.
54#[derive(Clone, Debug, Eq, PartialEq)]
55pub struct Parent<V> {
56    pub ptr: NodePtr,
57    pub coordinates: V,
58    pub child_index: ChildIndex,
59}
60
61/// A child [`NodePtr`] and its [`Parent`].
62#[derive(Clone, Debug, Eq, PartialEq)]
63pub struct ChildRelation<V> {
64    pub child: NodePtr,
65    pub parent: Option<Parent<V>>,
66}
67
68/// All children pointers for some branch node. Some may be empty.
69#[derive(Clone, Debug, Eq, PartialEq)]
70pub struct ChildPointers<'a, const CHILDREN: usize> {
71    level: Level,
72    pointers: &'a [AllocPtr; CHILDREN],
73}
74
75impl<'a, const CHILDREN: usize> ChildPointers<'a, CHILDREN> {
76    /// Returns the child pointer at the given `child` linear index.
77    #[inline]
78    pub fn get_child(&self, child: ChildIndex) -> Option<NodePtr> {
79        let alloc_ptr = self.pointers[child as usize];
80        (alloc_ptr != EMPTY_ALLOC_PTR).then(|| NodePtr::new(self.level, alloc_ptr))
81    }
82
83    #[inline]
84    pub fn level(&self) -> Level {
85        self.level
86    }
87}
88
89pub enum NodeEntry<'a, T, const CHILDREN: usize> {
90    Occupied(OccupiedNodeEntry<'a, T, CHILDREN>),
91    Vacant(VacantNodeEntry<'a, T, CHILDREN>),
92}
93
94impl<'a, T, const CHILDREN: usize> NodeEntry<'a, T, CHILDREN> {
95    pub fn or_insert_with(&mut self, mut filler: impl FnMut() -> T) -> (AllocPtr, &mut T) {
96        match self {
97            Self::Occupied(o) => (o.ptr, o.get_mut()),
98            Self::Vacant(v) => {
99                let ptr = v.insert(filler());
100                (ptr, unsafe { v.alloc.get_value_unchecked_mut(ptr) })
101            }
102        }
103    }
104
105    pub fn pointer(&self) -> AllocPtr {
106        // PERF: this extra branch seems unnecessary when both variants have a pointer
107        match self {
108            Self::Occupied(o) => o.ptr,
109            Self::Vacant(v) => *v.ptr,
110        }
111    }
112}
113
114pub struct OccupiedNodeEntry<'a, T, const CHILDREN: usize> {
115    alloc: &'a mut NodeAllocator<T, CHILDREN>,
116    ptr: AllocPtr,
117}
118
119impl<'a, T, const CHILDREN: usize> OccupiedNodeEntry<'a, T, CHILDREN> {
120    #[inline]
121    pub fn get_mut(&mut self) -> &mut T {
122        unsafe { self.alloc.get_value_unchecked_mut(self.ptr) }
123    }
124}
125
126pub struct VacantNodeEntry<'a, T, const CHILDREN: usize> {
127    alloc: &'a mut NodeAllocator<T, CHILDREN>,
128    ptr: &'a mut AllocPtr,
129    level: Level,
130}
131
132impl<'a, T, const CHILDREN: usize> VacantNodeEntry<'a, T, CHILDREN> {
133    #[inline]
134    pub fn insert(&mut self, value: T) -> AllocPtr {
135        let new_ptr = if self.level == 0 {
136            self.alloc.insert_leaf(value)
137        } else {
138            self.alloc.insert_branch(value)
139        };
140        *self.ptr = new_ptr;
141        new_ptr
142    }
143}
144
145#[derive(Clone, Copy, Debug, Eq, PartialEq)]
146pub struct RootNode {
147    pub self_ptr: AllocPtr,
148    /// Roots may optionally point back to a node in some other [`Tree`]'s allocator.
149    pub parent_ptr: Option<AllocPtr>,
150}
151
152impl RootNode {
153    pub fn new(self_ptr: AllocPtr, parent_ptr: Option<AllocPtr>) -> Self {
154        Self {
155            self_ptr,
156            parent_ptr,
157        }
158    }
159
160    pub fn new_without_parent(self_ptr: AllocPtr) -> Self {
161        Self::new(self_ptr, None)
162    }
163}
164
165/// A generic "grid tree" which can be either a quadtree or an octree depending on the type parameters.
166#[derive(Clone, Debug)]
167pub struct Tree<V, S, T, const CHILDREN: usize> {
168    /// 2x2 square in 2D or 2x2x2 cube in 3D.
169    branch_shape: PhantomData<S>,
170    /// Every node at the highest LOD is a root.
171    root_nodes: SmallKeyHashMap<NodeKey<V>, RootNode>,
172    /// An allocator for each level.
173    allocators: Vec<NodeAllocator<T, CHILDREN>>,
174}
175
176impl<V, S, T, const CHILDREN: usize> Tree<V, S, T, CHILDREN>
177where
178    V: VectorKey,
179    NodeKey<V>: Hash,
180    S: BranchShape<V>,
181{
182    /// The maximum number of children a branch node can have. 4 for a quadtree and 8 for an octree.
183    pub const CHILDREN: ChildIndex = CHILDREN as ChildIndex;
184
185    /// This constructor is only necessary if you need to use a custom shape `S` or vector `V`. Otherwise use the constructor
186    /// for [`OctreeI32`](crate::OctreeI32) or [`QuadtreeI32`](crate::QuadtreeI32).
187    ///
188    /// # Safety
189    ///
190    /// You must guarantee that `S` correctly linearizes `V` vectors so that they don't go out of array bounds in a block.
191    /// Refer to [`QuadtreeShapeI32`](crate::QuadtreeShapeI32) and [`OctreeShapeI32`](crate::OctreeShapeI32) for examples.
192    pub unsafe fn new_generic(height: Level) -> Self {
193        assert!(height > 1);
194        Self {
195            branch_shape: PhantomData,
196            root_nodes: Default::default(),
197            allocators: (0..height).map(|_| NodeAllocator::default()).collect(),
198        }
199    }
200
201    /// The number of [`Level`]s in this tree.
202    #[inline]
203    pub fn height(&self) -> Level {
204        self.allocators.len() as Level
205    }
206
207    /// The [`Level`] where root nodes live.
208    #[inline]
209    pub fn root_level(&self) -> Level {
210        self.height() - 1
211    }
212
213    /// Returns the unique root at `self.root_level()` that is an ancestor of `descendant_key`.
214    #[inline]
215    pub fn ancestor_root_key(&self, descendant_key: NodeKey<V>) -> NodeKey<V> {
216        assert!(descendant_key.level <= self.root_level());
217        let levels_up = self.root_level() - descendant_key.level;
218        let coordinates = S::ancestor_key(descendant_key.coordinates, levels_up as u32);
219        NodeKey {
220            level: self.root_level(),
221            coordinates,
222        }
223    }
224
225    /// Iterate over all root keys.
226    pub fn iter_root_keys(&self) -> impl Iterator<Item = &NodeKey<V>> {
227        self.root_nodes.keys()
228    }
229
230    /// Iterate over all root nodes.
231    pub fn iter_roots(&self) -> impl Iterator<Item = (&NodeKey<V>, &RootNode)> {
232        self.root_nodes.iter()
233    }
234
235    /// Returns true iff this tree contains a node for `ptr`.
236    #[inline]
237    pub fn contains_node(&self, ptr: NodePtr) -> bool {
238        self.allocator(ptr.level).contains_node(ptr.alloc_ptr)
239    }
240
241    /// Returns true iff this tree contains a root node at `coords`.
242    #[inline]
243    pub fn contains_root(&self, key: NodeKey<V>) -> bool {
244        self.root_nodes.contains_key(&key)
245    }
246
247    #[inline]
248    pub fn get_value(&self, ptr: NodePtr) -> Option<&T> {
249        self.allocator(ptr.level).get_value(ptr.alloc_ptr)
250    }
251
252    #[inline]
253    pub fn get_value_mut(&mut self, ptr: NodePtr) -> Option<&mut T> {
254        self.allocator_mut(ptr.level).get_value_mut(ptr.alloc_ptr)
255    }
256
257    /// # Safety
258    ///
259    /// The node for `ptr` must exist.
260    #[inline]
261    pub unsafe fn get_value_unchecked(&self, ptr: NodePtr) -> &T {
262        self.allocator(ptr.level).get_value_unchecked(ptr.alloc_ptr)
263    }
264
265    /// # Safety
266    ///
267    /// The node for `ptr` must exist.
268    #[inline]
269    pub unsafe fn get_value_unchecked_mut(&mut self, ptr: NodePtr) -> &mut T {
270        self.allocator_mut(ptr.level)
271            .get_value_unchecked_mut(ptr.alloc_ptr)
272    }
273
274    /// Inserts `value` at the root node at `key`. Returns the old value.
275    #[inline]
276    pub fn insert_root(
277        &mut self,
278        key: NodeKey<V>,
279        parent_ptr: Option<AllocPtr>,
280        new_value: T,
281    ) -> (RootNode, Option<T>) {
282        let Self {
283            root_nodes,
284            allocators,
285            ..
286        } = self;
287        let mut old_value = None;
288        let alloc = &mut allocators[key.level as usize];
289        let root_node = match root_nodes.entry(key) {
290            hash_map::Entry::Occupied(occupied) => {
291                let root_node = *occupied.get();
292                let current_value = unsafe { alloc.get_value_unchecked_mut(root_node.self_ptr) };
293                old_value = Some(mem::replace(current_value, new_value));
294                root_node
295            }
296            hash_map::Entry::Vacant(vacant) => {
297                let root_ptr = alloc.insert_branch(new_value);
298                let node = RootNode::new(root_ptr, parent_ptr);
299                vacant.insert(node);
300                node
301            }
302        };
303        (root_node, old_value)
304    }
305
306    /// Gets the root pointer or calls `filler` to insert a value first.
307    #[inline]
308    pub fn get_or_create_root(
309        &mut self,
310        key: NodeKey<V>,
311        mut filler: impl FnMut() -> T,
312    ) -> RootNode {
313        let Self {
314            root_nodes,
315            allocators,
316            ..
317        } = self;
318        let alloc = &mut allocators[key.level as usize];
319        *root_nodes.entry(key).or_insert_with(|| {
320            let root_ptr = alloc.insert_branch(filler());
321            RootNode::new_without_parent(root_ptr)
322        })
323    }
324
325    /// Inserts a child node of `parent_ptr` storing `child_value`. Returns the old child value if one exists.
326    #[inline]
327    pub fn insert_child(
328        &mut self,
329        parent_ptr: NodePtr,
330        child_index: ChildIndex,
331        child_value: T,
332    ) -> (NodePtr, Option<T>) {
333        assert!(parent_ptr.level > 0);
334        let child_level = parent_ptr.level - 1;
335        match self.child_entry(parent_ptr, child_index) {
336            NodeEntry::Occupied(mut o) => {
337                let value = o.get_mut();
338                let old_value = mem::replace(value, child_value);
339                (NodePtr::new(child_level, o.ptr), Some(old_value))
340            }
341            NodeEntry::Vacant(mut v) => {
342                v.insert(child_value);
343                (NodePtr::new(child_level, *v.ptr), None)
344            }
345        }
346    }
347
348    /// Same as `insert_child` but `child_offset` is linearized into a [`ChildIndex`] based on the [`BranchShape`].
349    ///
350    /// This means any given coordinate in `child_offset` can only be 0 or 1!
351    #[inline]
352    pub fn insert_child_at_offset(
353        &mut self,
354        parent_ptr: NodePtr,
355        child_offset: V,
356        child_value: T,
357    ) -> (NodePtr, Option<T>) {
358        self.insert_child(parent_ptr, S::linearize_child(child_offset), child_value)
359    }
360
361    /// Equivalent to calling `fill_root` and then `fill_descendants` of that root.
362    #[inline]
363    pub fn fill_tree_from_root(
364        &mut self,
365        root_key: NodeKey<V>,
366        min_level: Level,
367        mut filler: impl FnMut(NodeKey<V>, &mut NodeEntry<'_, T, CHILDREN>) -> VisitCommand,
368    ) {
369        if let (Some(root_node), VisitCommand::Continue) =
370            self.fill_root(root_key, |entry| filler(root_key, entry))
371        {
372            self.fill_descendants(
373                NodePtr::new(root_key.level, root_node.self_ptr),
374                root_key.coordinates,
375                min_level,
376                filler,
377            )
378        }
379    }
380
381    /// May return the [`RootNode`] for convenience.
382    #[inline]
383    pub fn fill_root(
384        &mut self,
385        key: NodeKey<V>,
386        filler: impl FnOnce(&mut NodeEntry<'_, T, CHILDREN>) -> VisitCommand,
387    ) -> (Option<RootNode>, VisitCommand) {
388        let Self {
389            allocators,
390            root_nodes,
391            ..
392        } = self;
393        let root_alloc = &mut allocators[key.level as usize];
394
395        match root_nodes.entry(key) {
396            hash_map::Entry::Occupied(occupied_node) => {
397                let root_node = *occupied_node.get();
398                let mut entry = NodeEntry::Occupied(OccupiedNodeEntry {
399                    alloc: root_alloc,
400                    ptr: root_node.self_ptr,
401                });
402                (Some(root_node), filler(&mut entry))
403            }
404            hash_map::Entry::Vacant(vacant_node) => {
405                let mut new_ptr = EMPTY_ALLOC_PTR;
406                let mut entry = NodeEntry::Vacant(VacantNodeEntry {
407                    alloc: root_alloc,
408                    ptr: &mut new_ptr,
409                    level: key.level,
410                });
411                let command = filler(&mut entry);
412                if new_ptr == EMPTY_ALLOC_PTR {
413                    (None, command)
414                } else {
415                    let root_node = RootNode::new_without_parent(new_ptr);
416                    vacant_node.insert(root_node);
417                    (Some(root_node), command)
418                }
419            }
420        }
421    }
422
423    /// # Panics
424    ///
425    /// - If `parent_ptr` is at level 0 and hence cannot have children.
426    /// - If no node exists for `parent_ptr`.
427    #[inline]
428    pub fn child_entry(
429        &mut self,
430        parent_ptr: NodePtr,
431        child_index: ChildIndex,
432    ) -> NodeEntry<'_, T, CHILDREN> {
433        let [parent_alloc, child_alloc] =
434            Self::parent_and_child_allocators_mut(&mut self.allocators, parent_ptr.level)
435                .unwrap_or_else(|| {
436                    panic!("Tried getting child of invalid parent: {:?}", parent_ptr)
437                });
438
439        let children = parent_alloc.get_children_mut_or_panic(parent_ptr.alloc_ptr);
440        let child_ptr = &mut children[child_index as usize];
441        if *child_ptr == EMPTY_ALLOC_PTR {
442            NodeEntry::Vacant(VacantNodeEntry {
443                alloc: child_alloc,
444                ptr: child_ptr,
445                level: parent_ptr.level - 1,
446            })
447        } else {
448            NodeEntry::Occupied(OccupiedNodeEntry {
449                alloc: child_alloc,
450                ptr: *child_ptr,
451            })
452        }
453    }
454
455    /// Inserts data in descendant "slots" of `ancestor_ptr` using `filler()`.
456    ///
457    /// Any node N is skipped if `filler` returns [`VisitCommand::SkipDescendants`] for any ancestor of N.
458    #[inline]
459    pub fn fill_descendants(
460        &mut self,
461        ancestor_ptr: NodePtr,
462        ancestor_coordinates: V,
463        min_level: Level,
464        mut filler: impl FnMut(NodeKey<V>, &mut NodeEntry<'_, T, CHILDREN>) -> VisitCommand,
465    ) {
466        assert!(min_level < ancestor_ptr.level());
467        let mut stack = SmallVec::<[(NodePtr, V); 32]>::new();
468        stack.push((ancestor_ptr, ancestor_coordinates));
469        while let Some((parent_ptr, parent_coords)) = stack.pop() {
470            if parent_ptr.level > 0 {
471                let has_grandchildren = parent_ptr.level > min_level + 1;
472                let child_level = parent_ptr.level - 1;
473                for child_index in 0..Self::CHILDREN {
474                    let mut child_entry = self.child_entry(parent_ptr, child_index);
475                    let child_coords = parent_coords + S::delinearize_child(child_index);
476                    let child_key = NodeKey::new(child_level, child_coords);
477                    let command = filler(child_key, &mut child_entry);
478                    if let VisitCommand::Continue = command {
479                        if has_grandchildren {
480                            let child_ptr = child_entry.pointer();
481                            if child_ptr != EMPTY_ALLOC_PTR {
482                                stack.push((NodePtr::new(child_level, child_ptr), child_coords));
483                            }
484                        }
485                    }
486                }
487            }
488        }
489    }
490
491    /// Call `filler` on all nodes from the root ancestor to `target_key`.
492    pub fn fill_path_to_node_from_root(
493        &mut self,
494        target_key: NodeKey<V>,
495        mut filler: impl FnMut(NodeKey<V>, &mut NodeEntry<'_, T, CHILDREN>) -> VisitCommand,
496    ) {
497        // We need to start from the top to avoid allocating nodes that already exist.
498        let root_key = self.ancestor_root_key(target_key);
499        let (root_node, command) =
500            self.fill_root(root_key, |root_entry| filler(root_key, root_entry));
501
502        if let (Some(root_node), VisitCommand::Continue) = (root_node, command) {
503            self.fill_path_to_node(
504                root_key.coordinates,
505                NodePtr::new(root_key.level, root_node.self_ptr),
506                target_key,
507                filler,
508            );
509        }
510    }
511
512    /// Call `filler` on all nodes from the ancestor to `target_key`.
513    pub fn fill_path_to_node(
514        &mut self,
515        ancestor_coordinates: V,
516        ancestor_ptr: NodePtr,
517        target_key: NodeKey<V>,
518        mut filler: impl FnMut(NodeKey<V>, &mut NodeEntry<'_, T, CHILDREN>) -> VisitCommand,
519    ) {
520        assert!(ancestor_ptr.level() >= target_key.level);
521
522        let mut parent_ptr = ancestor_ptr;
523        let mut parent_coords = ancestor_coordinates;
524        for child_level in (target_key.level..ancestor_ptr.level()).rev() {
525            // Get the child index of the ancestor at this level.
526            let level_diff = child_level - target_key.level;
527            let ancestor_coords = S::ancestor_key(target_key.coordinates, level_diff as u32);
528            let child_index = S::linearize_child(ancestor_coords - S::min_child_key(parent_coords));
529            let child_coords = parent_coords + S::delinearize_child(child_index);
530            let node_key = NodeKey::new(child_level, child_coords);
531            let mut child_entry = self.child_entry(parent_ptr, child_index);
532            let command = filler(node_key, &mut child_entry);
533            if command == VisitCommand::SkipDescendants {
534                break;
535            }
536            let child_ptr = child_entry.pointer();
537            if child_ptr == EMPTY_ALLOC_PTR {
538                break;
539            }
540            parent_ptr = NodePtr::new(child_level, child_ptr);
541            parent_coords = ancestor_coords;
542        }
543    }
544
545    /// Looks up the root pointer for `key` in the top-level hash map.
546    #[inline]
547    pub fn find_root(&self, key: NodeKey<V>) -> Option<RootNode> {
548        self.root_nodes.get(&key).cloned()
549    }
550
551    /// Starting from the ancestor root, searches for the corresponding descendant node at `key`.
552    ///
553    /// A [`ChildRelation`] is returned because it contains some extra useful info that is conveniently accessible during the
554    /// search.
555    #[inline]
556    pub fn find_node(&self, key: NodeKey<V>) -> Option<ChildRelation<V>> {
557        if key.level == self.root_level() {
558            self.find_root(key).map(|root_node| ChildRelation {
559                child: NodePtr::new(key.level, root_node.self_ptr),
560                parent: None,
561            })
562        } else {
563            let root_key = self.ancestor_root_key(key);
564            self.find_root(root_key).and_then(|root_node| {
565                self.find_descendant(
566                    NodePtr::new(root_key.level, root_node.self_ptr),
567                    root_key.coordinates,
568                    key,
569                )
570            })
571        }
572    }
573
574    /// Starting from the node at `ancestor_ptr`, searches for the corresponding descendant node at `descendant_key`.
575    pub fn find_descendant(
576        &self,
577        ancestor_ptr: NodePtr,
578        ancestor_coordinates: V,
579        descendant_key: NodeKey<V>,
580    ) -> Option<ChildRelation<V>> {
581        assert!(
582            ancestor_ptr.level > descendant_key.level,
583            "{} > {}",
584            ancestor_ptr.level,
585            descendant_key.level
586        );
587        let level_diff = ancestor_ptr.level - descendant_key.level;
588
589        self.child_pointers(ancestor_ptr).and_then(|children| {
590            let child_coords = if level_diff == 1 {
591                descendant_key.coordinates
592            } else {
593                S::ancestor_key(descendant_key.coordinates, level_diff as u32 - 1)
594            };
595            let min_sibling_coords = S::min_child_key(ancestor_coordinates);
596            let offset = child_coords - min_sibling_coords;
597            let child_index = S::linearize_child(offset);
598            let child_ptr = children.pointers[child_index as usize];
599
600            (child_ptr != EMPTY_ALLOC_PTR)
601                .then(|| child_ptr)
602                .and_then(|child_ptr| {
603                    let child_ptr = NodePtr {
604                        level: children.level,
605                        alloc_ptr: child_ptr,
606                    };
607                    if level_diff == 1 {
608                        // Ancestor is the parent.
609                        Some(ChildRelation {
610                            child: child_ptr,
611                            parent: Some(Parent {
612                                ptr: ancestor_ptr,
613                                coordinates: ancestor_coordinates,
614                                child_index,
615                            }),
616                        })
617                    } else {
618                        self.find_descendant(child_ptr, child_coords, descendant_key)
619                    }
620                })
621        })
622    }
623
624    /// Visits pointers for all non-empty children of the node at `parent_ptr`.
625    ///
626    /// If `parent_ptr` does not exist or does not have any children, nothing happens.
627    #[inline]
628    pub fn visit_children(
629        &self,
630        parent_ptr: NodePtr,
631        mut visitor: impl FnMut(NodePtr, ChildIndex),
632    ) {
633        if let Some(children) = self.child_pointers(parent_ptr) {
634            for (child_index, &child_ptr) in children.pointers.iter().enumerate() {
635                if child_ptr != EMPTY_ALLOC_PTR {
636                    visitor(
637                        NodePtr {
638                            level: children.level,
639                            alloc_ptr: child_ptr,
640                        },
641                        child_index as ChildIndex,
642                    );
643                }
644            }
645        }
646    }
647
648    /// Visits pointers and coordinates for all non-empty children of the node at `parent_ptr` with coordinates `parent_coords`.
649    ///
650    /// If `parent_ptr` does not exist or does not have any children, nothing happens.
651    #[inline]
652    pub fn visit_children_with_coordinates(
653        &self,
654        parent_ptr: NodePtr,
655        parent_coordinates: V,
656        mut visitor: impl FnMut(NodePtr, V),
657    ) {
658        self.visit_children(parent_ptr, |child_ptr, child_index| {
659            let child_offset = S::delinearize_child(child_index as ChildIndex);
660            let child_coords = S::min_child_key(parent_coordinates) + child_offset;
661            visitor(child_ptr, child_coords)
662        })
663    }
664
665    /// Visit `ancestor_ptr` and all descendants in depth-first order.
666    ///
667    /// If `visitor` returns `false`, descendants of that node will not be visited.
668    #[inline]
669    pub fn visit_tree_depth_first(
670        &self,
671        ancestor_ptr: NodePtr,
672        ancestor_coordinates: V,
673        min_level: Level,
674        mut visitor: impl FnMut(NodePtr, V) -> VisitCommand,
675    ) {
676        let mut stack = SmallVec::<[(NodePtr, V); 32]>::new();
677        stack.push((ancestor_ptr, ancestor_coordinates));
678        while let Some((parent_ptr, parent_coords)) = stack.pop() {
679            let command = visitor(parent_ptr, parent_coords);
680            if let VisitCommand::Continue = command {
681                if parent_ptr.level > min_level {
682                    self.visit_children_with_coordinates(
683                        parent_ptr,
684                        parent_coords,
685                        |child_ptr, child_coords| {
686                            stack.push((child_ptr, child_coords));
687                        },
688                    );
689                }
690            }
691        }
692    }
693
694    /// Visit `ancestor_ptr` and all descendants in breadth-first order.
695    ///
696    /// If `visitor` returns `false`, descendants of that node will not be visited.
697    #[inline]
698    pub fn visit_tree_breadth_first(
699        &self,
700        ancestor_ptr: NodePtr,
701        ancestor_coordinates: V,
702        min_level: Level,
703        mut visitor: impl FnMut(NodePtr, V) -> VisitCommand,
704    ) {
705        let mut queue = VecDeque::new();
706        queue.push_back((ancestor_ptr, ancestor_coordinates));
707        while let Some((parent_ptr, parent_coords)) = queue.pop_front() {
708            let command = visitor(parent_ptr, parent_coords);
709            if command == VisitCommand::Continue && parent_ptr.level > min_level {
710                self.visit_children_with_coordinates(
711                    parent_ptr,
712                    parent_coords,
713                    |child_ptr, child_coords| {
714                        queue.push_back((child_ptr, child_coords));
715                    },
716                );
717            }
718        }
719    }
720
721    /// Returns an array of pointers to the children of `parent_ptr`.
722    ///
723    /// Returns `None` if `parent_ptr` is at level 0.
724    #[inline]
725    pub fn child_pointers(&self, parent_ptr: NodePtr) -> Option<ChildPointers<'_, CHILDREN>> {
726        self.allocator(parent_ptr.level)
727            .get_children(parent_ptr.alloc_ptr)
728            .map(|children| ChildPointers {
729                level: parent_ptr.level - 1,
730                pointers: children,
731            })
732    }
733
734    /// Drops the child node of `relation` and all descendants.
735    #[inline]
736    pub fn drop_tree(&mut self, relation: &ChildRelation<V>) {
737        self.unlink_child(relation);
738
739        // Drop node and all descendants.
740        let mut to_drop = SmallVec::<[NodePtr; 32]>::new();
741        to_drop.push(relation.child);
742        while let Some(ptr) = to_drop.pop() {
743            let (_value, children) = self.allocator_mut(ptr.level).remove(ptr.alloc_ptr);
744            if let Some(children) = children {
745                let child_level = ptr.level - 1;
746                for child_ptr in children.into_iter() {
747                    if child_ptr != EMPTY_ALLOC_PTR {
748                        to_drop.push(NodePtr {
749                            level: child_level,
750                            alloc_ptr: child_ptr,
751                        });
752                    }
753                }
754            }
755        }
756    }
757
758    /// Moves the child node of `relation` (with `coordinates`) and all descendants into `consumer`.
759    #[inline]
760    pub fn remove_tree(
761        &mut self,
762        relation: &ChildRelation<V>,
763        coordinates: V,
764        mut consumer: impl FnMut(NodeKey<V>, T),
765    ) {
766        self.unlink_child(relation);
767
768        let mut to_remove = SmallVec::<[(NodePtr, V); 32]>::new();
769        to_remove.push((relation.child, coordinates));
770        while let Some((ptr, coordinates)) = to_remove.pop() {
771            let (value, children) = self.allocator_mut(ptr.level).remove(ptr.alloc_ptr);
772            if let Some(value) = value {
773                let node_key = NodeKey {
774                    level: ptr.level,
775                    coordinates,
776                };
777                consumer(node_key, value);
778                if let Some(children) = children {
779                    let child_level = ptr.level - 1;
780                    let min_child = S::min_child_key(coordinates);
781                    for (child_index, child_ptr) in children.into_iter().enumerate() {
782                        if child_ptr != EMPTY_ALLOC_PTR {
783                            let child_coords =
784                                min_child + S::delinearize_child(child_index as ChildIndex);
785                            to_remove.push((
786                                NodePtr {
787                                    level: child_level,
788                                    alloc_ptr: child_ptr,
789                                },
790                                child_coords,
791                            ));
792                        }
793                    }
794                }
795            }
796        }
797    }
798
799    fn unlink_child(&mut self, relation: &ChildRelation<V>) {
800        if let Some(parent) = relation.parent.as_ref() {
801            self.root_nodes
802                .remove(&NodeKey::new(relation.child.level + 1, parent.coordinates));
803            self.allocator_mut(parent.ptr.level)
804                .unlink_child(parent.ptr.alloc_ptr, parent.child_index);
805        }
806    }
807
808    fn parent_and_child_allocators_mut(
809        allocators: &mut [NodeAllocator<T, CHILDREN>],
810        parent_level: Level,
811    ) -> Option<[&mut NodeAllocator<T, CHILDREN>; 2]> {
812        if parent_level == 0 {
813            return None;
814        }
815
816        let (left, right) = allocators.split_at_mut(parent_level as usize);
817        let child_alloc = left.last_mut().unwrap();
818        let parent_alloc = right.first_mut().unwrap();
819
820        Some([parent_alloc, child_alloc])
821    }
822
823    fn allocator(&self, level: Level) -> &NodeAllocator<T, CHILDREN> {
824        &self.allocators[level as usize]
825    }
826
827    fn allocator_mut(&mut self, level: Level) -> &mut NodeAllocator<T, CHILDREN> {
828        &mut self.allocators[level as usize]
829    }
830}
831
832#[derive(Clone, Copy, Debug, Eq, PartialEq)]
833pub enum VisitCommand {
834    Continue,
835    SkipDescendants,
836}
837
838// ████████╗███████╗███████╗████████╗
839// ╚══██╔══╝██╔════╝██╔════╝╚══██╔══╝
840//    ██║   █████╗  ███████╗   ██║
841//    ██║   ██╔══╝  ╚════██║   ██║
842//    ██║   ███████╗███████║   ██║
843//    ╚═╝   ╚══════╝╚══════╝   ╚═╝
844
845#[cfg(test)]
846mod test {
847    use super::*;
848
849    use crate::glam::IVec3;
850    use crate::OctreeI32;
851
852    #[test]
853    fn insert_root() {
854        let mut tree = OctreeI32::new(3);
855
856        let key = NodeKey::new(2, IVec3::ZERO);
857
858        assert_eq!(tree.find_root(key), None);
859
860        let (node, old_value) = tree.insert_root(key, None, "val1");
861        assert_eq!(old_value, None);
862        assert_eq!(node.parent_ptr, None);
863        let root_ptr = NodePtr::new(2, node.self_ptr);
864        assert!(tree.contains_node(root_ptr));
865        assert_eq!(tree.get_value(root_ptr), Some(&"val1"));
866        assert_eq!(tree.find_root(key), Some(node));
867
868        let (_node, old_value) = tree.insert_root(key, None, "val2");
869        assert_eq!(old_value, Some("val1"));
870        assert!(tree.contains_node(root_ptr));
871        assert_eq!(tree.get_value(root_ptr), Some(&"val2"));
872    }
873
874    #[test]
875    fn get_or_create_root() {
876        let mut tree = OctreeI32::new(3);
877
878        let key = NodeKey::new(2, IVec3::ZERO);
879
880        let ptr = tree.get_or_create_root(key, || ());
881        assert_eq!(ptr, tree.get_or_create_root(key, || ()));
882
883        let (ptr, _old_value) = tree.insert_root(key, None, ());
884        assert_eq!(ptr, tree.get_or_create_root(key, || ()));
885    }
886
887    #[test]
888    fn insert_child_of_root() {
889        let mut tree = OctreeI32::new(3);
890
891        let root_key = NodeKey::new(2, IVec3::ZERO);
892        let (root_node, _) = tree.insert_root(root_key, None, "val1");
893        let root_ptr = NodePtr::new(2, root_node.self_ptr);
894        let (child_ptr, old_val) = tree.insert_child(root_ptr, 0, "val2");
895        assert_eq!(old_val, None);
896        assert!(tree.contains_node(child_ptr));
897        assert_eq!(tree.get_value(child_ptr), Some(&"val2"));
898    }
899
900    #[test]
901    fn find_descendant_none() {
902        let mut tree = OctreeI32::new(3);
903
904        let root_key = NodeKey::new(2, IVec3::ZERO);
905        let (root_node, _) = tree.insert_root(root_key, None, ());
906        let root_ptr = NodePtr::new(2, root_node.self_ptr);
907
908        let descendant_key = NodeKey {
909            level: 1,
910            coordinates: IVec3::ZERO,
911        };
912        let found = tree.find_descendant(root_ptr, root_key.coordinates, descendant_key);
913        assert_eq!(found, None);
914        let found = tree.find_node(descendant_key);
915        assert_eq!(found, None);
916    }
917
918    #[test]
919    fn find_descendant_child() {
920        let mut tree = OctreeI32::new(3);
921
922        let root_key = NodeKey::new(2, IVec3::ZERO);
923        let (root_node, _) = tree.insert_root(root_key, None, ());
924        let root_ptr = NodePtr::new(2, root_node.self_ptr);
925
926        let child_key = NodeKey {
927            level: 1,
928            coordinates: IVec3::new(1, 1, 1),
929        };
930        let (child_ptr, _) = tree.insert_child_at_offset(root_ptr, child_key.coordinates, ());
931
932        let expected_find = Some(ChildRelation {
933            child: child_ptr,
934            parent: Some(Parent {
935                ptr: root_ptr,
936                coordinates: IVec3::ZERO,
937                child_index: 7,
938            }),
939        });
940        let found = tree.find_descendant(root_ptr, root_key.coordinates, child_key);
941        assert_eq!(found, expected_find);
942        let found = tree.find_node(child_key);
943        assert_eq!(found, expected_find);
944    }
945
946    #[test]
947    fn find_descendant_grandchild() {
948        let mut tree = OctreeI32::new(3);
949
950        let root_key = NodeKey::new(2, IVec3::ZERO);
951        let (root_node, _) = tree.insert_root(root_key, None, ());
952        let root_ptr = NodePtr::new(2, root_node.self_ptr);
953
954        let grandchild_key = NodeKey {
955            level: 0,
956            coordinates: IVec3::new(3, 3, 3),
957        };
958        let (child_ptr, _) = tree.insert_child_at_offset(root_ptr, IVec3::new(1, 1, 1), ());
959        let (grandchild_ptr, _) = tree.insert_child_at_offset(child_ptr, IVec3::new(1, 1, 1), ());
960
961        let expected_find = Some(ChildRelation {
962            child: grandchild_ptr,
963            parent: Some(Parent {
964                ptr: child_ptr,
965                coordinates: IVec3::new(1, 1, 1),
966                child_index: 7,
967            }),
968        });
969        let found = tree.find_descendant(root_ptr, root_key.coordinates, grandchild_key);
970        assert_eq!(found, expected_find);
971        let found = tree.find_node(grandchild_key);
972        assert_eq!(found, expected_find);
973    }
974
975    #[test]
976    fn visit_children_of_root() {
977        let mut tree = OctreeI32::new(3);
978
979        let root_key = NodeKey::new(2, IVec3::new(1, 1, 1));
980        let (root_node, _) = tree.insert_root(root_key, None, ());
981        let root_ptr = NodePtr::new(2, root_node.self_ptr);
982        let (child1_ptr, _) = tree.insert_child_at_offset(root_ptr, IVec3::new(0, 0, 0), ());
983        let (child2_ptr, _) = tree.insert_child_at_offset(root_ptr, IVec3::new(1, 1, 1), ());
984
985        let mut visited = Vec::new();
986        tree.visit_children_with_coordinates(
987            root_ptr,
988            root_key.coordinates,
989            |child_ptr, child_coords| {
990                visited.push((child_coords, child_ptr));
991            },
992        );
993
994        assert_eq!(
995            visited.as_slice(),
996            &[
997                (IVec3::new(2, 2, 2), child1_ptr),
998                (IVec3::new(3, 3, 3), child2_ptr)
999            ]
1000        );
1001    }
1002
1003    #[test]
1004    fn visit_tree_of_root() {
1005        let mut tree = OctreeI32::new(3);
1006
1007        let root_key = NodeKey::new(2, IVec3::new(1, 1, 1));
1008        let (root_node, _) = tree.insert_root(root_key, None, ());
1009        let root_ptr = NodePtr::new(2, root_node.self_ptr);
1010        let (child1_ptr, _) = tree.insert_child_at_offset(root_ptr, IVec3::new(0, 0, 0), ());
1011        let (child2_ptr, _) = tree.insert_child_at_offset(root_ptr, IVec3::new(1, 1, 1), ());
1012        let (grandchild1_ptr, _) = tree.insert_child_at_offset(child1_ptr, IVec3::new(1, 1, 1), ());
1013        let (grandchild2_ptr, _) = tree.insert_child_at_offset(child2_ptr, IVec3::new(0, 0, 0), ());
1014
1015        let mut visited = Vec::new();
1016        tree.visit_tree_depth_first(
1017            root_ptr,
1018            root_key.coordinates,
1019            0,
1020            |child_ptr, child_coords| {
1021                visited.push((child_coords, child_ptr));
1022                VisitCommand::Continue
1023            },
1024        );
1025        assert_eq!(
1026            visited.as_slice(),
1027            &[
1028                (IVec3::new(1, 1, 1), root_ptr),
1029                (IVec3::new(3, 3, 3), child2_ptr),
1030                (IVec3::new(6, 6, 6), grandchild2_ptr),
1031                (IVec3::new(2, 2, 2), child1_ptr),
1032                (IVec3::new(5, 5, 5), grandchild1_ptr),
1033            ]
1034        );
1035
1036        let mut visited = Vec::new();
1037        tree.visit_tree_breadth_first(
1038            root_ptr,
1039            root_key.coordinates,
1040            0,
1041            |child_ptr, child_coords| {
1042                visited.push((child_coords, child_ptr));
1043                VisitCommand::Continue
1044            },
1045        );
1046        assert_eq!(
1047            visited.as_slice(),
1048            &[
1049                (IVec3::new(1, 1, 1), root_ptr),
1050                (IVec3::new(2, 2, 2), child1_ptr),
1051                (IVec3::new(3, 3, 3), child2_ptr),
1052                (IVec3::new(5, 5, 5), grandchild1_ptr),
1053                (IVec3::new(6, 6, 6), grandchild2_ptr),
1054            ]
1055        );
1056    }
1057
1058    #[test]
1059    fn drop_tree() {
1060        let mut tree = OctreeI32::new(3);
1061
1062        let root_key = NodeKey::new(2, IVec3::new(1, 1, 1));
1063        let (root_node, _) = tree.insert_root(root_key, None, ());
1064        let root_ptr = NodePtr::new(2, root_node.self_ptr);
1065        let (child1_ptr, _) = tree.insert_child_at_offset(root_ptr, IVec3::new(0, 0, 0), ());
1066        let (child2_ptr, _) = tree.insert_child_at_offset(root_ptr, IVec3::new(1, 1, 1), ());
1067        let (grandchild1_ptr, _) = tree.insert_child_at_offset(child1_ptr, IVec3::new(1, 1, 1), ());
1068        let (grandchild2_ptr, _) = tree.insert_child_at_offset(child2_ptr, IVec3::new(0, 0, 0), ());
1069
1070        let child1_relation = tree
1071            .find_node(NodeKey {
1072                level: 1,
1073                coordinates: IVec3::new(2, 2, 2),
1074            })
1075            .unwrap();
1076        tree.drop_tree(&child1_relation);
1077
1078        assert!(!tree.contains_node(child1_ptr));
1079        assert!(!tree.contains_node(grandchild1_ptr));
1080
1081        let mut visited = Vec::new();
1082        tree.visit_tree_breadth_first(root_ptr, root_key.coordinates, 0, |ptr, coords| {
1083            visited.push((ptr, coords));
1084            VisitCommand::Continue
1085        });
1086
1087        assert_eq!(
1088            visited.as_slice(),
1089            &[
1090                (root_ptr, root_key.coordinates),
1091                (child2_ptr, IVec3::new(3, 3, 3)),
1092                (grandchild2_ptr, IVec3::new(6, 6, 6))
1093            ]
1094        );
1095    }
1096
1097    #[test]
1098    fn remove_tree() {
1099        let mut tree = OctreeI32::new(3);
1100
1101        let root_key = NodeKey::new(2, IVec3::new(1, 1, 1));
1102        let (root_node, _) = tree.insert_root(root_key, None, ());
1103        let root_ptr = NodePtr::new(2, root_node.self_ptr);
1104        let (child1_ptr, _) = tree.insert_child_at_offset(root_ptr, IVec3::new(0, 0, 0), ());
1105        let (child2_ptr, _) = tree.insert_child_at_offset(root_ptr, IVec3::new(1, 1, 1), ());
1106        let (grandchild1_ptr, _) = tree.insert_child_at_offset(child1_ptr, IVec3::new(1, 1, 1), ());
1107        let (grandchild2_ptr, _) = tree.insert_child_at_offset(child2_ptr, IVec3::new(0, 0, 0), ());
1108
1109        let child1_relation = tree
1110            .find_node(NodeKey {
1111                level: 1,
1112                coordinates: IVec3::new(2, 2, 2),
1113            })
1114            .unwrap();
1115        let mut removed = Vec::new();
1116        let child1_coords = IVec3::new(2, 2, 2);
1117        tree.remove_tree(&child1_relation, child1_coords, |key, value| {
1118            removed.push((key, value))
1119        });
1120
1121        assert_eq!(
1122            removed.as_slice(),
1123            &[
1124                (
1125                    NodeKey {
1126                        level: 1,
1127                        coordinates: IVec3::new(2, 2, 2)
1128                    },
1129                    ()
1130                ),
1131                (
1132                    NodeKey {
1133                        level: 0,
1134                        coordinates: IVec3::new(5, 5, 5)
1135                    },
1136                    ()
1137                ),
1138            ]
1139        );
1140
1141        assert!(!tree.contains_node(child1_ptr));
1142        assert!(!tree.contains_node(grandchild1_ptr));
1143
1144        let mut visited = Vec::new();
1145        tree.visit_tree_breadth_first(root_ptr, root_key.coordinates, 0, |ptr, coords| {
1146            visited.push((ptr, coords));
1147            VisitCommand::Continue
1148        });
1149
1150        assert_eq!(
1151            visited.as_slice(),
1152            &[
1153                (root_ptr, root_key.coordinates),
1154                (child2_ptr, IVec3::new(3, 3, 3)),
1155                (grandchild2_ptr, IVec3::new(6, 6, 6))
1156            ]
1157        );
1158    }
1159
1160    #[test]
1161    fn fill_some_descendants() {
1162        let mut tree = OctreeI32::new(3);
1163
1164        let root_key = NodeKey::new(2, IVec3::new(1, 1, 1));
1165        let (root_node, _) = tree.insert_root(root_key, None, ());
1166        let root_ptr = NodePtr::new(2, root_node.self_ptr);
1167        tree.fill_descendants(root_ptr, root_key.coordinates, 0, |_child_coords, entry| {
1168            let (ptr, &mut ()) = entry.or_insert_with(|| ());
1169            assert_ne!(ptr, EMPTY_ALLOC_PTR);
1170            VisitCommand::Continue
1171        });
1172
1173        let mut visited_lvl0 = SmallKeyHashMap::new();
1174        tree.visit_tree_depth_first(
1175            root_ptr,
1176            root_key.coordinates,
1177            0,
1178            |child_ptr, child_coords| {
1179                if child_ptr.level() == 0 && child_coords % 2 == IVec3::ZERO {
1180                    visited_lvl0.insert(child_coords, child_ptr);
1181                }
1182                VisitCommand::Continue
1183            },
1184        );
1185
1186        let mut expected_lvl0 = SmallKeyHashMap::new();
1187        for z in 4..8 {
1188            for y in 4..8 {
1189                for x in 4..8 {
1190                    let p = IVec3::new(x, y, z);
1191                    if p % 2 == IVec3::ZERO {
1192                        let relation = tree
1193                            .find_node(NodeKey {
1194                                level: 0,
1195                                coordinates: p,
1196                            })
1197                            .unwrap();
1198                        expected_lvl0.insert(p, relation.child);
1199                    }
1200                }
1201            }
1202        }
1203        assert_eq!(visited_lvl0, expected_lvl0);
1204    }
1205
1206    #[test]
1207    fn fill_root() {
1208        let mut tree = OctreeI32::new(5);
1209
1210        let root_key = NodeKey::new(2, IVec3::new(1, 1, 1));
1211        let (root_node, _) = tree.fill_root(root_key, |entry| {
1212            match entry {
1213                NodeEntry::Occupied(_) => {
1214                    panic!("Unexpected occupied entry");
1215                }
1216                NodeEntry::Vacant(v) => {
1217                    v.insert(());
1218                }
1219            }
1220            VisitCommand::Continue
1221        });
1222        assert!(tree.contains_root(root_key));
1223        assert_eq!(tree.find_root(root_key).unwrap(), root_node.unwrap());
1224    }
1225
1226    #[test]
1227    fn fill_path_to_node() {
1228        let mut tree = OctreeI32::new(5);
1229
1230        let target_key = NodeKey::new(1, IVec3::new(1, 1, 1));
1231        let mut path = Vec::new();
1232        tree.fill_path_to_node_from_root(target_key, |key, entry| {
1233            match entry {
1234                NodeEntry::Occupied(_) => {
1235                    panic!("Unexpected occupied entry");
1236                }
1237                NodeEntry::Vacant(v) => {
1238                    let new_ptr = v.insert(());
1239                    path.push((NodePtr::new(key.level, new_ptr), key.coordinates));
1240                }
1241            }
1242            VisitCommand::Continue
1243        });
1244
1245        assert_eq!(path.len(), 4);
1246
1247        for (ptr, coords) in path.into_iter() {
1248            let key = NodeKey::new(ptr.level(), coords);
1249            assert!(tree.contains_node(ptr));
1250            let relation = tree.find_node(key).unwrap();
1251            assert_eq!(relation.child, ptr);
1252        }
1253    }
1254}
grid_tree/tree.rs

grid_tree/
tree.rs
