Skip to main content

feagi_brain_development/spatial/
morton.rs

1// Copyright 2025 Neuraville Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4/*!
5Morton encoding utilities for 3D spatial hashing.
6
7Implements Z-order curve encoding to preserve spatial locality.
8*/
9
10/// Morton encode 3D coordinates into a single u64.
11///
12/// Interleaves bits of x, y, z coordinates to preserve spatial locality.
13/// Each dimension limited to 21 bits (0-2,097,151).
14#[inline]
15pub fn morton_encode_3d(x: u32, y: u32, z: u32) -> u64 {
16    // Limit to 21 bits per dimension (63 bits total)
17    debug_assert!(x < (1 << 21), "x coordinate exceeds 21-bit limit");
18    debug_assert!(y < (1 << 21), "y coordinate exceeds 21-bit limit");
19    debug_assert!(z < (1 << 21), "z coordinate exceeds 21-bit limit");
20
21    let mut result = 0u64;
22
23    // Interleave bits: ...z2y2x2z1y1x1z0y0x0
24    for i in 0..21 {
25        result |= ((x as u64 & (1 << i)) << (2 * i))
26            | ((y as u64 & (1 << i)) << (2 * i + 1))
27            | ((z as u64 & (1 << i)) << (2 * i + 2));
28    }
29
30    result
31}
32
33/// Morton decode a u64 back to 3D coordinates.
34#[inline]
35pub fn morton_decode_3d(morton_code: u64) -> (u32, u32, u32) {
36    let mut x = 0u32;
37    let mut y = 0u32;
38    let mut z = 0u32;
39
40    // Extract interleaved bits
41    for i in 0..21 {
42        x |= ((morton_code & (1 << (3 * i))) >> (2 * i)) as u32;
43        y |= ((morton_code & (1 << (3 * i + 1))) >> (2 * i + 1)) as u32;
44        z |= ((morton_code & (1 << (3 * i + 2))) >> (2 * i + 2)) as u32;
45    }
46
47    (x, y, z)
48}
49
50/// Encode a 3D region into a vector of Morton codes.
51///
52/// Returns all Morton codes for coordinates in the inclusive range
53/// [x1..=x2, y1..=y2, z1..=z2].
54pub fn morton_encode_region_3d(x1: u32, y1: u32, z1: u32, x2: u32, y2: u32, z2: u32) -> Vec<u64> {
55    let mut codes = Vec::new();
56
57    for z in z1..=z2 {
58        for y in y1..=y2 {
59            for x in x1..=x2 {
60                codes.push(morton_encode_3d(x, y, z));
61            }
62        }
63    }
64
65    codes
66}
67
68#[cfg(test)]
69mod tests {
70    use super::*;
71
72    #[test]
73    fn test_morton_encode_decode() {
74        let coords = [(0, 0, 0), (1, 2, 3), (100, 200, 300), (1000, 2000, 3000)];
75
76        for (x, y, z) in coords {
77            let encoded = morton_encode_3d(x, y, z);
78            let (dx, dy, dz) = morton_decode_3d(encoded);
79            assert_eq!(
80                (x, y, z),
81                (dx, dy, dz),
82                "Round-trip failed for ({}, {}, {})",
83                x,
84                y,
85                z
86            );
87        }
88    }
89
90    #[test]
91    fn test_spatial_locality() {
92        // Nearby points should have nearby Morton codes
93        let c1 = morton_encode_3d(10, 10, 10);
94        let c2 = morton_encode_3d(11, 10, 10);
95        let c3 = morton_encode_3d(100, 100, 100);
96
97        // c1 and c2 should be closer than c1 and c3
98        let diff_near = c1.abs_diff(c2);
99        let diff_far = c1.abs_diff(c3);
100        assert!(diff_near < diff_far, "Spatial locality not preserved");
101    }
102
103    #[test]
104    fn test_region_encoding() {
105        let codes = morton_encode_region_3d(0, 0, 0, 1, 1, 1);
106        assert_eq!(codes.len(), 8, "2x2x2 region should have 8 points");
107
108        // All codes should be unique
109        let unique: std::collections::HashSet<_> = codes.iter().collect();
110        assert_eq!(unique.len(), 8, "All Morton codes should be unique");
111    }
112}