Skip to main content

ndslive_math/
morton.rs

1// SPDX-License-Identifier: BSD-3-Clause
2//! Morton (Z-order) encoding for 2D NDS coordinates.
3//!
4//! Faithful port of `python/src/ndslive/math/morton.py` and cross-checked
5//! against `cpp/include/ndsmath/mortoncode.h`.
6
7/// A 64-bit Morton (Z-order) code.
8///
9/// Wraps a raw `u64`. The encoder masks off bit 63 to preserve the NDS
10/// semantics of the reference implementation.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
12pub struct MortonCode(pub u64);
13
14impl MortonCode {
15    /// Construct a `MortonCode` from a raw 64-bit value.
16    ///
17    /// Most callers use [`MortonCode::from_nds_coordinates`] instead.
18    pub fn new(morton_code: u64) -> Self {
19        // Python masks to 64 bits; in Rust a u64 is already 64-bit.
20        MortonCode(morton_code)
21    }
22
23    /// Encode NDS integer coordinates into a Morton (Z-order) code.
24    ///
25    /// `x` is NDS longitude (signed 32-bit), `y` is NDS latitude
26    /// (signed 31-bit). Out-of-range inputs are wrapped, exactly as in the
27    /// Python reference. Bit 63 is masked off.
28    pub fn from_nds_coordinates(x: i32, y: i32) -> MortonCode {
29        let x_base: i64 = 1 << 31;
30        let y_base: i64 = 1 << 30;
31
32        let mut x: i64 = x as i64;
33        let mut y: i64 = y as i64;
34
35        // Wrap x into [-2^31, 2^31) and y into [-2^30, 2^30).
36        while x >= x_base {
37            x -= 1 << 32;
38        }
39        while x < -x_base {
40            x += 1 << 32;
41        }
42        while y >= y_base {
43            y -= 1 << 31;
44        }
45        while y < -y_base {
46            y += 1 << 31;
47        }
48
49        // Reinterpret the (now reduced) signed values as their two's-complement
50        // bit patterns so that the AND/shift logic matches Python's behavior on
51        // negative ints. Everything is then done with wrapping shifts in u64.
52        let mut x: u64 = x as u64;
53        let mut y: u64 = y as u64;
54
55        let mut bit: u64 = 1;
56        let mut morton_code: u64 = 0;
57
58        y = y.wrapping_shl(1);
59
60        for _ in 0..31 {
61            morton_code |= x & bit;
62            x = x.wrapping_shl(1);
63            bit = bit.wrapping_shl(1);
64
65            morton_code |= y & bit;
66            y = y.wrapping_shl(1);
67            bit = bit.wrapping_shl(1);
68        }
69
70        morton_code |= x & bit;
71        // The final `x <<= 1` / `bit <<= 1` in the reference are dead stores
72        // (their results are never read), so we omit them.
73
74        morton_code &= !(1u64 << 63);
75
76        MortonCode(morton_code)
77    }
78
79    /// Decode this Morton code back into NDS integer coordinates.
80    ///
81    /// Inverse of [`MortonCode::from_nds_coordinates`]. Returns `(x, y)` as
82    /// signed integers: NDS longitude (32-bit) and NDS latitude (31-bit).
83    pub fn to_nds_coordinates(&self) -> (i32, i32) {
84        const YBASE: i64 = 1 << 30;
85        const XBASE: i64 = 1 << 31;
86
87        let mut bit: u64 = 1;
88        let mut morton_code: u64 = self.0;
89        let mut x: u64 = 0;
90        let mut y: u64 = 0;
91
92        for _ in 0..31 {
93            x |= morton_code & bit;
94            morton_code >>= 1;
95            y |= morton_code & bit;
96            bit <<= 1;
97        }
98
99        x |= morton_code & bit;
100        // Reference does a final `morton_code >>= 1` here whose result is unused.
101
102        let mut x = x as i64;
103        let mut y = y as i64;
104
105        if y >= YBASE {
106            y -= 1 << 31;
107        }
108        if x >= XBASE {
109            x -= 1 << 32;
110        }
111
112        (x as i32, y as i32)
113    }
114
115    /// Get the raw Morton code value (matches the C++/Python `value()` API).
116    pub fn value(&self) -> u64 {
117        self.0
118    }
119}
120
121impl std::fmt::Display for MortonCode {
122    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
123        write!(f, "MortonCode(value={})", self.0)
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130
131    fn round_trip(x: i32, y: i32, expected: u64) {
132        let m = MortonCode::from_nds_coordinates(x, y);
133        assert_eq!(m.value(), expected, "encode ({x},{y})");
134        let (dx, dy) = m.to_nds_coordinates();
135        assert_eq!((dx, dy), (x, y), "decode ({x},{y})");
136    }
137
138    #[test]
139    fn small_values() {
140        round_trip(0, 0, 0);
141        round_trip(1, 0, 1);
142        round_trip(0, 1, 2);
143        round_trip(1, 1, 3);
144        round_trip(12345, 6789, 126387555);
145    }
146
147    #[test]
148    fn negative_values() {
149        round_trip(-12345, -6789, 9223372036728388255);
150        round_trip(i32::MIN, -(1i32 << 30), 6917529027641081856);
151    }
152
153    #[test]
154    fn extremes() {
155        round_trip(i32::MAX, (1i32 << 30) - 1, 2305843009213693951);
156        round_trip(1000000000, -1000000000, 2694675840375717888);
157    }
158
159    #[test]
160    fn bit63_masked() {
161        // No encoded value may ever have bit 63 set.
162        for &(x, y) in &[
163            (0, 0),
164            (-1, -1),
165            (i32::MAX, (1i32 << 30) - 1),
166            (i32::MIN, -(1i32 << 30)),
167        ] {
168            let m = MortonCode::from_nds_coordinates(x, y);
169            assert_eq!(m.value() & (1u64 << 63), 0);
170        }
171    }
172}