grid_tree/
allocator.rs

1use crate::ChildIndex;
2
3use slab::Slab;
4
5/// Allocate branch or leaf nodes.
6///
7/// One [`AllocPtr`] corresponds to a single node at the given [`Level`](crate::Level). Branches always have children pointers,
8/// although they may be [`EMPTY_ALLOC_PTR`]. Leaves do not allocate child pointers. For a given [`Level`](crate::Level), only one
9/// type of node should be allocated. Level 0 is only leaves and all other levels are only branches.
10#[derive(Clone, Debug)]
11pub struct NodeAllocator<T, const CHILDREN: usize> {
12    /// A slab of pointer blocks. Empty at level 0.
13    pointers: Slab<[AllocPtr; CHILDREN]>,
14    /// A slab of values.
15    values: Slab<T>,
16}
17
18/// Points to a node owned by an internal allocator.
19pub type AllocPtr = u32;
20
21/// An [`AllocPtr`] that doesn't point to anything.
22pub const EMPTY_ALLOC_PTR: AllocPtr = AllocPtr::MAX;
23
24impl<T, const CHILDREN: usize> Default for NodeAllocator<T, CHILDREN> {
25    fn default() -> Self {
26        Self {
27            pointers: Default::default(),
28            values: Default::default(),
29        }
30    }
31}
32
33impl<T, const CHILDREN: usize> NodeAllocator<T, CHILDREN> {
34    #[inline]
35    pub fn insert_leaf(&mut self, value: T) -> AllocPtr {
36        self.values.insert(value) as AllocPtr
37    }
38
39    #[inline]
40    pub fn insert_branch(&mut self, value: T) -> AllocPtr {
41        let ptr = self.values.insert(value);
42        let pointer_entry = self.pointers.vacant_entry();
43        debug_assert_eq!(ptr, pointer_entry.key());
44        pointer_entry.insert([EMPTY_ALLOC_PTR; CHILDREN]);
45        ptr as AllocPtr
46    }
47
48    #[inline]
49    pub fn remove(&mut self, ptr: AllocPtr) -> (Option<T>, Option<[AllocPtr; CHILDREN]>) {
50        (
51            self.values.try_remove(ptr as usize),
52            self.pointers.try_remove(ptr as usize),
53        )
54    }
55
56    #[inline]
57    pub fn contains_node(&self, ptr: AllocPtr) -> bool {
58        // This assumes that every node has a value.
59        self.values.contains(ptr as usize)
60    }
61
62    #[inline]
63    pub unsafe fn get_value_unchecked(&self, ptr: AllocPtr) -> &T {
64        self.values.get_unchecked(ptr as usize)
65    }
66
67    #[inline]
68    pub unsafe fn get_value_unchecked_mut(&mut self, ptr: AllocPtr) -> &mut T {
69        self.values.get_unchecked_mut(ptr as usize)
70    }
71
72    #[inline]
73    pub fn get_children_mut_or_panic(&mut self, ptr: AllocPtr) -> &mut [AllocPtr; CHILDREN] {
74        self.get_children_mut(ptr).unwrap_or_else(|| {
75            panic!(
76                "Tried inserting children of {:?} which has no child pointers",
77                ptr,
78            )
79        })
80    }
81
82    #[inline]
83    pub fn get_value(&self, ptr: AllocPtr) -> Option<&T> {
84        self.values.get(ptr as usize)
85    }
86
87    #[inline]
88    pub fn get_value_mut(&mut self, ptr: AllocPtr) -> Option<&mut T> {
89        self.values.get_mut(ptr as usize)
90    }
91
92    #[inline]
93    pub fn get_children(&self, ptr: AllocPtr) -> Option<&[AllocPtr; CHILDREN]> {
94        self.pointers.get(ptr as usize)
95    }
96
97    #[inline]
98    pub fn get_children_mut(&mut self, ptr: AllocPtr) -> Option<&mut [AllocPtr; CHILDREN]> {
99        self.pointers.get_mut(ptr as usize)
100    }
101
102    #[inline]
103    pub fn set_child_pointer(
104        &mut self,
105        parent_ptr: AllocPtr,
106        child_index: ChildIndex,
107        child_ptr: AllocPtr,
108    ) {
109        if let Some(children) = self.get_children_mut(parent_ptr) {
110            children[child_index as usize] = child_ptr;
111        }
112    }
113
114    #[inline]
115    pub fn unlink_child(&mut self, parent_ptr: AllocPtr, child_index: ChildIndex) {
116        self.set_child_pointer(parent_ptr, child_index, EMPTY_ALLOC_PTR)
117    }
118}
grid_tree/allocator.rs

grid_tree/
allocator.rs
