dmc/
octree.rs

1//! Contains the container for the mesh generation [`HashedOctree`], which relies on
2//! [`MortonKey`] for indexing.
3
4use ahash::RandomState;
5use bitflags::bitflags;
6use cgmath::*;
7use dashmap::{mapref::one::*, DashMap};
8
9bitflags! {
10    /// Indicates a node for when indexing an octree for any of its children.
11    pub struct OctreeIndex : u8 {
12        /// Index +X (1) or -X (0).
13        const RIGHT = 1 << 0;
14        /// Index +Y (1) or -Y (0).
15        const TOP = 1 << 1;
16        /// Index +Z (1) or -Z (0).
17        const FOREMOST = 1 << 2;
18    }
19}
20
21impl<T> From<OctreeIndex> for Vector3<T>
22where
23    T: From<i8>,
24{
25    fn from(x: OctreeIndex) -> Self {
26        let result = Vector3::new(
27            if x.contains(OctreeIndex::RIGHT) {
28                T::from(1)
29            } else {
30                T::from(-1)
31            },
32            if x.contains(OctreeIndex::TOP) {
33                T::from(1)
34            } else {
35                T::from(-1)
36            },
37            if x.contains(OctreeIndex::FOREMOST) {
38                T::from(-1)
39            } else {
40                T::from(1)
41            },
42        );
43
44        result
45    }
46}
47
48/// A 64-bit integer that indicates a single node in a hashed octree.
49/// Its internal representation starts with 1 as the root node and appends 3 bits for each following
50/// child node, indicating its position in the X, Y and Z axis.
51#[derive(Hash, PartialEq, Eq, Clone, Copy, Debug)]
52pub struct MortonKey(pub u64);
53
54impl MortonKey {
55    /// Returns a Morton key pointing to the octree root node.
56    pub const fn root() -> Self {
57        Self(1)
58    }
59
60    /// Returns a Morton key pointing to no node.
61    pub const fn none() -> Self {
62        Self(0)
63    }
64
65    /// Returns the parent of this Morton key.
66    /// If this key is pointing to the root node, the node given will be equal to [`MortonKey::none()`].
67    pub fn parent(self) -> Self {
68        Self(self.0 >> 3)
69    }
70
71    /// Returns a child of the node this key is pointing to.
72    /// # Panics
73    /// Panics if the level of this key is equal to the maximum level possible.
74    pub fn child(self, index: OctreeIndex) -> Self {
75        assert!(self.level() < Self::max_level());
76        Self(self.0 << 3 | index.bits as u64)
77    }
78
79    /// Returns the level or depth of this node, where 0 is the root node.
80    /// # Panics
81    /// Panics if this is a Morton key pointing to no node.
82    pub fn level(self) -> u32 {
83        assert_ne!(self.0, 0);
84        (63 - self.0.leading_zeros()) / 3
85    }
86
87    /// Returns the maximum depth of the nodes that this type can point to.
88    pub const fn max_level() -> u32 {
89        63 / 3
90    }
91
92    /// Returns the position of the center of the node this key is pointing to.
93    /// All axis are in the (0, 1) interval.
94    ///
95    pub fn position(self) -> Point3<f32> {
96        let (mut x, mut y, mut z) = (0, 0, 0);
97        let level = self.level();
98
99        let mut bits = self.0;
100
101        for i in 1..=level {
102            x |= (bits & 1) << i;
103            bits >>= 1;
104            y |= (bits & 1) << i;
105            bits >>= 1;
106            z |= (bits & 1) << i;
107            bits >>= 1;
108        }
109
110        let max_level_k = (1 << level) as f32;
111
112        Point3::new(
113            (x + 1) as f32 / max_level_k - 1.,
114            (y + 1) as f32 / max_level_k - 1.,
115            (z + 1) as f32 / max_level_k - 1.,
116        )
117    }
118
119    /// Returns the Morton key with level `level` of an octree of bounds \[-1]³ to \[1]³ which is
120    /// closest to the position `position`.
121    /// # Panics
122    /// Panics if `level < Self::max_level()`.
123    pub fn closest_to_position(position: Point3<f32>, level: u32) -> Self {
124        assert!(level < Self::max_level());
125
126        let position = point3(
127            position.x.clamp(-1., 1.),
128            position.y.clamp(-1., 1.),
129            position.z.clamp(-1., 1.),
130        );
131
132        let max_level_k = (1 << level) as f32;
133
134        let (x, y, z) = (
135            ((position.x + 1.) * max_level_k - 1.) as u64,
136            ((position.y + 1.) * max_level_k - 1.) as u64,
137            ((position.z + 1.) * max_level_k - 1.) as u64,
138        );
139
140        let mut bits = 1u64;
141
142        for i in (1..=level).rev() {
143            bits <<= 1;
144            bits |= (z >> i) & 1;
145            bits <<= 1;
146            bits |= (y >> i) & 1;
147            bits <<= 1;
148            bits |= (x >> i) & 1;
149        }
150
151        Self(bits)
152    }
153
154    /// Returns a key that goes along the current key until a given level.
155    /// # Panics
156    /// If `self.level() < level`.
157    pub fn until_level(self, level: u32) -> Self {
158        assert!(self.level() >= level);
159        Self(self.0 >> (self.level() - level) * 3)
160    }
161}
162
163/// A hashed (linear) octree implementation which uses [`MortonKey`] for indexing.
164/// This type supports concurrent access and uses a fast hashing algorithm for best performance.
165/// # Example
166/// ```rust
167/// use dmc::octree::*;
168///
169/// let mut octree = HashedOctree::new(0usize);
170///
171/// let mut i = 0;
172/// octree.subdivide(MortonKey::root()).for_each(|child| {
173///     *octree.value_mut(child).unwrap() = i; i += 1;
174/// });
175/// dbg!(octree);
176/// ```
177#[derive(Clone, Debug)]
178#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
179pub struct HashedOctree<T: Copy> {
180    values: DashMap<MortonKey, T, RandomState>,
181}
182
183impl<T: Copy> HashedOctree<T> {
184    /// Creates a new hashed octree, with a value for the root node.
185    pub fn new(root_value: T) -> Self {
186        Self {
187            values: {
188                let x = DashMap::with_capacity_and_hasher(256, RandomState::new());
189                x.insert(MortonKey::root(), root_value);
190                x
191            },
192        }
193    }
194
195    /// Returns a reference to a node if it exists.
196    ///
197    /// # Concurrent Behaviour
198    /// May deadlock if called when holding a mutable reference into the octree.
199    pub fn value(&self, key: MortonKey) -> Option<Ref<'_, MortonKey, T, RandomState>> {
200        self.values.get(&key)
201    }
202
203    /// Returns a mutable reference to a node if it exists.
204    ///
205    /// # Concurrent Behaviour
206    /// May deadlock if called when holding any sort of reference into the octree.
207    pub fn value_mut(&mut self, key: MortonKey) -> Option<RefMut<'_, MortonKey, T, RandomState>> {
208        self.values.get_mut(&key)
209    }
210
211    /// Subdivides a node and returns an iterator over its newly created children.
212    /// The value of the children will be copied from the parent.
213    ///
214    /// # Panics
215    /// Panics if the node passed is already subdivided.
216    pub fn subdivide(&mut self, key: MortonKey) -> impl std::iter::Iterator<Item = MortonKey> {
217        assert!(
218            !self.is_subdivided(key),
219            "Tried to subdivide already subdivided node"
220        );
221        let child_value = *self
222            .value(key)
223            .expect("Tried to subdivide null octree node");
224
225        let vec = (0..8)
226            .into_iter()
227            .map(move |child| {
228                let child_key = key.child(unsafe { OctreeIndex::from_bits_unchecked(child) });
229                self.values.insert(child_key, child_value);
230                child_key
231            })
232            .collect::<Vec<_>>();
233
234        vec.into_iter()
235    }
236
237    /// Returns an iterator over the children of a node, or None if the children (or parent) don't
238    /// exist.
239    pub fn children(&self, key: MortonKey) -> Option<impl std::iter::Iterator<Item = MortonKey>> {
240        if self.is_subdivided(key) {
241            Some((0..8).into_iter().map(move |child| {
242                let child_key = key.child(unsafe { OctreeIndex::from_bits_unchecked(child) });
243                child_key
244            }))
245        } else {
246            None
247        }
248    }
249
250    /// Returns true if a node exists.
251    pub fn node_exists(&self, key: MortonKey) -> bool {
252        self.values.contains_key(&key)
253    }
254
255    /// Returns true if the children of a node exists.
256    ///
257    /// # Panics
258    /// Panics if the level of the key given is equal to the maximum level possible.
259    pub fn is_subdivided(&self, key: MortonKey) -> bool {
260        self.values.contains_key(&(key.child(OctreeIndex::empty())))
261    }
262
263    /// Finds all the leaf nodes belonging to `parent` and returns a vector with a key for each of
264    /// them.
265    pub fn leaves(&self, parent: MortonKey) -> Vec<MortonKey> {
266        let mut leaves = Vec::with_capacity(128);
267        let mut to_process = Vec::with_capacity(64);
268        to_process.push(parent);
269
270        while let Some(node_to_process) = to_process.pop() {
271            if let Some(children) = self.children(node_to_process) {
272                for child in children {
273                    to_process.push(child);
274                }
275            } else {
276                leaves.push(node_to_process);
277            }
278        }
279
280        leaves
281    }
282
283    /// Returns the total node count, including leaf and branch nodes.
284    pub fn node_count(&self) -> usize {
285        self.values.len()
286    }
287}