1use ahash::RandomState;
5use bitflags::bitflags;
6use cgmath::*;
7use dashmap::{mapref::one::*, DashMap};
8
9bitflags! {
10    pub struct OctreeIndex : u8 {
12        const RIGHT = 1 << 0;
14        const TOP = 1 << 1;
16        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#[derive(Hash, PartialEq, Eq, Clone, Copy, Debug)]
52pub struct MortonKey(pub u64);
53
54impl MortonKey {
55    pub const fn root() -> Self {
57        Self(1)
58    }
59
60    pub const fn none() -> Self {
62        Self(0)
63    }
64
65    pub fn parent(self) -> Self {
68        Self(self.0 >> 3)
69    }
70
71    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    pub fn level(self) -> u32 {
83        assert_ne!(self.0, 0);
84        (63 - self.0.leading_zeros()) / 3
85    }
86
87    pub const fn max_level() -> u32 {
89        63 / 3
90    }
91
92    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    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    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#[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    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    pub fn value(&self, key: MortonKey) -> Option<Ref<'_, MortonKey, T, RandomState>> {
200        self.values.get(&key)
201    }
202
203    pub fn value_mut(&mut self, key: MortonKey) -> Option<RefMut<'_, MortonKey, T, RandomState>> {
208        self.values.get_mut(&key)
209    }
210
211    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    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    pub fn node_exists(&self, key: MortonKey) -> bool {
252        self.values.contains_key(&key)
253    }
254
255    pub fn is_subdivided(&self, key: MortonKey) -> bool {
260        self.values.contains_key(&(key.child(OctreeIndex::empty())))
261    }
262
263    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    pub fn node_count(&self) -> usize {
285        self.values.len()
286    }
287}