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}