Skip to main content

nodedb_array/coord/
zorder.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! ND Z-order (Morton) encode/decode.
4//!
5//! Bit-interleaving fallback for non-uniform dim sizes — cheap to
6//! compute, less locality than Hilbert but tolerates dims of differing
7//! granularity. Same `n × bits <= 64` constraint as
8//! [`super::hilbert`].
9
10use super::normalize::MAX_DIMS;
11use crate::error::{ArrayError, ArrayResult};
12
13/// Interleave per-dim bits into a single Z-order index, MSB first.
14pub fn encode(coords: &[u64], bits: u32) -> ArrayResult<u64> {
15    check_shape(coords.len(), bits)?;
16    let n = coords.len();
17    let mut idx: u64 = 0;
18    for b in (0..bits).rev() {
19        for c in coords.iter().take(n) {
20            let bit = (*c >> b) & 1;
21            idx = (idx << 1) | bit;
22        }
23    }
24    Ok(idx)
25}
26
27/// Inverse of [`encode`].
28pub fn decode(idx: u64, n: usize, bits: u32) -> ArrayResult<Vec<u64>> {
29    check_shape(n, bits)?;
30    let mut out = vec![0u64; n];
31    for b in (0..bits).rev() {
32        for (i, slot) in out.iter_mut().enumerate().take(n) {
33            let shift = b * n as u32 + (n as u32 - 1 - i as u32);
34            let bit = (idx >> shift) & 1;
35            *slot = (*slot << 1) | bit;
36        }
37    }
38    Ok(out)
39}
40
41fn check_shape(n: usize, bits: u32) -> ArrayResult<()> {
42    if n > MAX_DIMS {
43        return Err(ArrayError::InvalidSchema {
44            array: String::new(),
45            detail: format!("zorder: arity {n} exceeds MAX_DIMS={MAX_DIMS}"),
46        });
47    }
48    if bits == 0 || (n as u32) * bits > 64 {
49        return Err(ArrayError::InvalidSchema {
50            array: String::new(),
51            detail: format!("zorder: {n} dims × {bits} bits exceeds 64-bit prefix"),
52        });
53    }
54    Ok(())
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60
61    #[test]
62    fn zorder_round_trip_2d() {
63        for x in 0..16u64 {
64            for y in 0..16u64 {
65                let idx = encode(&[x, y], 4).unwrap();
66                assert_eq!(decode(idx, 2, 4).unwrap(), vec![x, y]);
67            }
68        }
69    }
70
71    #[test]
72    fn zorder_round_trip_3d() {
73        for x in 0..8u64 {
74            for y in 0..8u64 {
75                for z in 0..8u64 {
76                    let idx = encode(&[x, y, z], 3).unwrap();
77                    assert_eq!(decode(idx, 3, 3).unwrap(), vec![x, y, z]);
78                }
79            }
80        }
81    }
82
83    #[test]
84    fn zorder_unique_2d() {
85        use std::collections::HashSet;
86        let mut seen = HashSet::new();
87        for x in 0..8u64 {
88            for y in 0..8u64 {
89                assert!(seen.insert(encode(&[x, y], 3).unwrap()));
90            }
91        }
92        assert_eq!(seen.len(), 64);
93    }
94
95    #[test]
96    fn zorder_rejects_overflow() {
97        let coords = vec![0u64; 9];
98        assert!(encode(&coords, 8).is_err());
99    }
100}