baby_shark/voxel/internal_node/
mod.rs

1mod csg;
2mod flood_fill;
3mod tree_node;
4
5use super::*;
6use crate::{
7    data_structures::bitset::*,
8    helpers::{
9        aliases::{Vec3i, Vec3u},
10        one_of::OneOf,
11    },
12};
13use std::{alloc::Layout, fmt::Debug, mem::ManuallyDrop};
14
15#[derive(Debug)]
16pub(super) struct InternalNode<
17    TValue: Value,
18    TChild: TreeNode,
19    const BRANCHING: usize,
20    const BRANCHING_TOTAL: usize,
21    const SIZE: usize,
22    const BIT_SIZE: usize,
23    const PARALLEL: bool,
24> {
25    origin: Vec3i,
26    child_mask: BitArray<SIZE, BIT_SIZE>,
27    value_mask: BitArray<SIZE, BIT_SIZE>,
28    childs: [ChildUnion<TValue, TChild>; SIZE], // `child_mask` and `value_mask` are discriminating between branch and tile
29}
30
31impl<
32        TChild,
33        const BRANCHING: usize,
34        const BRANCHING_TOTAL: usize,
35        const SIZE: usize,
36        const BIT_SIZE: usize,
37        const PARALLEL: bool,
38    > InternalNode<TChild::Value, TChild, BRANCHING, BRANCHING_TOTAL, SIZE, BIT_SIZE, PARALLEL>
39where
40    TChild: TreeNode,
41{
42    fn offset_to_global_index(&self, offset: usize) -> Vec3i {
43        let mut local = Self::offset_to_local_index(offset);
44
45        for i in 0..3 {
46            local[i] <<= TChild::BRANCHING_TOTAL;
47        }
48
49        local.cast() + self.origin
50    }
51
52    fn add_branch(&mut self, offset: usize) -> &mut TChild {
53        if self.child_mask.is_on(offset) {
54            return unsafe { &mut self.childs[offset].branch };
55        }
56
57        self.child_mask.on(offset);
58        self.value_mask.off(offset);
59
60        let child_origin = self.offset_to_global_index(offset);
61        let child_node = TChild::empty(child_origin);
62        self.childs[offset] = ChildUnion {
63            branch: ManuallyDrop::new(child_node),
64        };
65
66        unsafe { &mut self.childs[offset].branch }
67    }
68
69    #[inline]
70    fn remove_branch(&mut self, offset: usize) -> Option<Box<TChild>> {
71        if self.child_mask.is_off(offset) {
72            return None;
73        }
74
75        self.child_mask.off(offset);
76        let child = unsafe { ManuallyDrop::take(&mut self.childs[offset].branch) };
77        Some(child)
78    }
79
80    #[inline]
81    fn remove_child(&mut self, offset: usize) -> Option<ChildOwned<TChild>> {
82        if self.child_mask.is_on(offset) {
83            self.child_mask.off(offset);
84            let branch = unsafe { ManuallyDrop::take(&mut self.childs[offset].branch) };
85            Some(ChildOwned::T1(branch))
86        } else if self.value_mask.is_on(offset) {
87            self.value_mask.off(offset);
88            let value = unsafe { self.childs[offset].tile };
89            Some(ChildOwned::T2(value))
90        } else {
91            None
92        }
93    }
94
95    #[inline]
96    fn child(&self, offset: usize) -> Option<Child<'_, TChild>> {
97        let child = &self.childs[offset];
98
99        unsafe {
100            if self.child_mask.is_on(offset) {
101                return Some(OneOf::T1(child.branch.as_ref()));
102            } else if self.value_mask.is_on(offset) {
103                return Some(OneOf::T2(&child.tile));
104            }
105        }
106
107        None
108    }
109
110    #[inline]
111    fn child_mut(&mut self, offset: usize) -> Option<ChildMut<'_, TChild>> {
112        let child = &mut self.childs[offset];
113
114        unsafe {
115            if self.child_mask.is_on(offset) {
116                return Some(OneOf::T1(child.branch.as_mut()));
117            } else if self.value_mask.is_on(offset) {
118                return Some(OneOf::T2(&mut child.tile));
119            }
120        }
121
122        None
123    }
124
125    fn childs(&self) -> impl Iterator<Item = (usize, Child<'_, TChild>)> + '_ {
126        (0..SIZE).filter_map(|offset| self.child(offset).map(|child| (offset, child)))
127    }
128
129    #[inline]
130    fn offset(index: &Vec3i) -> usize {
131        let offset = (((index.x & ((1 << Self::BRANCHING_TOTAL) - 1)) >> <Self as TreeNode>::Child::BRANCHING_TOTAL) << (Self::BRANCHING + Self::BRANCHING))
132            + (((index.y & ((1 << Self::BRANCHING_TOTAL) - 1)) >> <Self as TreeNode>::Child::BRANCHING_TOTAL) << Self::BRANCHING)
133            + ((index.z & ((1 << Self::BRANCHING_TOTAL) - 1)) >> <Self as TreeNode>::Child::BRANCHING_TOTAL);
134
135        offset as usize
136    }
137
138    #[inline]
139    fn offset_to_local_index(mut offset: usize) -> Vec3u {
140        debug_assert!(offset < (1 << (3 * BRANCHING)));
141
142        let x = offset >> (2 * BRANCHING);
143        offset &= (1 << (2 * BRANCHING)) - 1;
144        let y = offset >> BRANCHING;
145        let z = offset & ((1 << BRANCHING) - 1);
146
147        Vec3u::new(x, y, z)
148    }
149
150    const fn childs_per_dim() -> usize {
151        1 << BRANCHING
152    }
153
154    ///
155    /// Allocates memory for the node on the heap and initializes it with default values.
156    ///
157    unsafe fn alloc_on_heap(origin: Vec3i) -> Box<Self> {
158        let layout = Layout::new::<Self>();
159        let ptr = std::alloc::alloc(layout) as *mut Self;
160
161        if ptr.is_null() {
162            std::alloc::handle_alloc_error(layout);
163        }
164
165        (*ptr).origin = origin;
166        (*ptr).child_mask.off_all();
167        (*ptr).value_mask.off_all();
168
169        Box::from_raw(ptr)
170    }
171}
172
173impl<
174        TValue: Value,
175        TChild: TreeNode,
176        const BRANCHING: usize,
177        const BRANCHING_TOTAL: usize,
178        const SIZE: usize,
179        const BIT_SIZE: usize,
180        const PARALLEL: bool,
181    > Drop
182    for InternalNode<TValue, TChild, BRANCHING, BRANCHING_TOTAL, SIZE, BIT_SIZE, PARALLEL>
183where
184    TChild: TreeNode,
185{
186    fn drop(&mut self) {
187        for offset in 0..SIZE {
188            if self.child_mask.is_on(offset) {
189                unsafe {
190                    ManuallyDrop::drop(&mut self.childs[offset].branch);
191                }
192            }
193        }
194    }
195}
196
197pub type Child<'node, T> = OneOf<&'node T, &'node <T as TreeNode>::Value>;
198pub type ChildMut<'node, T> = OneOf<&'node mut T, &'node mut <T as TreeNode>::Value>;
199pub type ChildOwned<T> = OneOf<Box<T>, <T as TreeNode>::Value>;
200
201pub const fn internal_node_size(branching: usize) -> usize {
202    1 << (branching * 3)
203}
204
205pub const fn internal_node_bit_size(branching: usize) -> usize {
206    let size = internal_node_size(branching) / usize::BITS as usize;
207
208    if size == 0 {
209        1
210    } else {
211        size
212    }
213}
214
215pub const fn internal_node_branching<TChild: TreeNode>(branching: usize) -> usize {
216    branching + TChild::BRANCHING_TOTAL
217}
218
219union ChildUnion<TValue: Value, TChild: TreeNode> {
220    branch: ManuallyDrop<Box<TChild>>,
221    tile: TValue,
222}
223
224impl<TValue: Value, TChild: TreeNode> Debug for ChildUnion<TValue, TChild> {
225    #[inline]
226    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
227        f.write_str("unknown")
228    }
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234
235    #[test]
236    fn test_alloc_internal_node() {
237        const LEAF_LOG2: usize = 2;
238        const INTERNAL_LOG2: usize = 3;
239
240        type Leaf = LeafNode<
241            f32,
242            LEAF_LOG2,
243            LEAF_LOG2,
244            { leaf_node_size(LEAF_LOG2) },
245            { leaf_node_bit_size(LEAF_LOG2) },
246        >;
247        type Internal = InternalNode<
248            f32,
249            Leaf,
250            INTERNAL_LOG2,
251            { LEAF_LOG2 + INTERNAL_LOG2 },
252            { internal_node_size(INTERNAL_LOG2) },
253            { internal_node_bit_size(INTERNAL_LOG2) },
254            false,
255        >;
256
257        let origin = Vec3i::new(1, 2, 3);
258        let node = Internal::empty(origin);
259        assert_eq!(node.child_mask, BitArray::zeroes());
260        assert_eq!(node.value_mask, BitArray::zeroes());
261        assert_eq!(node.origin, origin);
262    }
263}