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::*;

/// Split an unsigned integer such that its bits are spaced by 2 zeros.
/// For instance, morton_split(0b00110010) = 0b000000001001000000001000.
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);
    }
}