1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
use crate::{utils::*, *};
use glam::*;
use std::fmt::Debug;
#[cfg(not(feature = "wasm_support"))]
use rayon::prelude::*;
pub fn morton_split(int: u32) -> u32 {
let log_bits = round_up_log2(std::mem::size_of::<u32>() as u32 * 8, 0);
let mut x = int as usize;
let mut mask = !0_usize;
let mut i = log_bits as usize;
let mut n = 1 << log_bits as usize;
while i > 0 {
mask = (mask | (mask << n)) & !(mask << (n / 2));
x = (x | (x << n)) & mask;
n >>= 1;
i -= 1;
}
x as u32
}
pub struct MortonEncoder {
world_to_grid: Vec3,
grid_offset: Vec3,
grid_dim: usize,
}
impl MortonEncoder {
pub const MAX_GRID_DIM: usize = 1 << (std::mem::size_of::<u32>() * 8 / 3);
pub fn new(aabb: &Aabb, grid_dim: usize) -> MortonEncoder {
debug_assert!(grid_dim <= Self::MAX_GRID_DIM);
let world_to_grid = grid_dim as f32 * (1.0 / aabb.diagonal());
let grid_offset = -Vec3::from(aabb.min) * world_to_grid;
Self {
world_to_grid,
grid_offset,
grid_dim,
}
}
pub fn morton_encode(x: u32, y: u32, z: u32) -> u32 {
morton_split(x) | (morton_split(y) << 1) | (morton_split(z) << 2)
}
pub fn encode<T: Into<[f32; 3]> + Debug + Copy>(&self, point: T) -> u32 {
let grid_pos = Vec3::from(point.into()) * self.world_to_grid + self.grid_offset;
let min = (self.grid_dim - 1) as i32;
let x: u32 = min.min((grid_pos[0] as i32).max(0)) as u32;
let y: u32 = min.min((grid_pos[1] as i32).max(0)) as u32;
let z: u32 = min.min((grid_pos[2] as i32).max(0)) as u32;
Self::morton_encode(x, y, z)
}
pub fn get_sorted_indices<E: Debug + Copy + Send + Sync, T: Primitive<E>>(
&self,
aabbs: &[Aabb<E>],
primitives: &[T],
) -> (Vec<u32>, Vec<u32>) {
debug_assert_eq!(aabbs.len(), primitives.len());
let prim_count = aabbs.len();
let mut indices: Vec<u32> = (0..(prim_count as u32)).collect();
#[cfg(not(feature = "wasm_support"))]
let morton_codes: Vec<u32> = (0..prim_count)
.into_par_iter()
.map(|i| self.encode(Vec3::from(primitives[i].center())))
.collect();
#[cfg(feature = "wasm_support")]
let morton_codes: Vec<u32> = (0..prim_count)
.into_iter()
.map(|i| self.encode(Vec3::from(primitives[i].center())))
.collect();
#[cfg(not(feature = "wasm_support"))]
indices.par_sort_by(|a, b| {
let a = (*a) as usize;
let b = (*b) as usize;
morton_codes[a].cmp(&morton_codes[b])
});
#[cfg(feature = "wasm_support")]
indices.sort_by(|a, b| {
let a = (*a) as usize;
let b = (*b) as usize;
morton_codes[a].cmp(&morton_codes[b])
});
(indices, morton_codes)
}
}
#[cfg(test)]
mod tests {
use crate::morton::*;
#[test]
fn morton_split_works() {
assert_eq!(morton_split(2), 8);
assert_eq!(morton_split(4), 64);
assert_eq!(morton_split(8), 512);
assert_eq!(morton_split(32), 32768);
}
}