hurray_core/layout/addressing/
morton.rs

1//! Morton (Z-order curve) element offset computation.
2//!
3//! Spec: docs/spec/layouts/morton.md § Morton Index Computation
4
5use crate::layout::MortonLayout;
6use crate::{Error, Result, Shape};
7
8use super::{validate_index, ElementAddress};
9
10impl ElementAddress for MortonLayout {
11    fn element_offset(&self, index: &[u64], shape: &Shape) -> Result<u64> {
12        validate_index(index, shape)?;
13        if self.morton_bits.len() != shape.rank() {
14            return Err(Error::IndexRankMismatch {
15                index_rank: self.morton_bits.len(),
16                shape_rank: shape.rank(),
17            });
18        }
19
20        let rank = index.len();
21        let total_bits: u32 = self.morton_bits.iter().sum();
22        if total_bits > 64 {
23            return Err(Error::IndexArithmeticOverflow);
24        }
25
26        let max_bits = self.morton_bits.iter().copied().max().unwrap_or(0);
27        let mut morton_code: u64 = 0;
28
29        // Bit interleaving: for each bit position b and dimension d, place bit b
30        // of index[d] at output position b*rank+d. Round-robin LSB-first per spec.
31        for b in 0..max_bits {
32            for (d, (&idx_val, &bits_for_dim)) in
33                index.iter().zip(self.morton_bits.iter()).enumerate()
34            {
35                if b < bits_for_dim {
36                    let bit = (idx_val >> b) & 1;
37                    let shift = b * rank as u32 + d as u32;
38                    if shift >= 64 {
39                        return Err(Error::IndexArithmeticOverflow);
40                    }
41                    morton_code |= bit << shift;
42                }
43            }
44        }
45        Ok(morton_code)
46    }
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52    use crate::{Error, Shape};
53
54    // Spec docs/spec/layouts/morton.md § Morton Index Computation:
55    // Bit interleaving: for index [i_0, i_1, …], place bit b of i_d at output
56    // bit position b*rank + d (round-robin, LSB-first).
57    //
58    // Full 4×4 conformance table (morton_bits=[2,2], shape=[4,4]):
59    //   [i_0, i_1] → morton code
60    //   i_0 occupies bits 0,2; i_1 occupies bits 1,3.
61    fn layout_2d_2bits() -> (MortonLayout, Shape) {
62        let layout = MortonLayout::new(vec![2, 2]).unwrap();
63        let shape = Shape::new(vec![4, 4]).unwrap();
64        (layout, shape)
65    }
66
67    #[test]
68    fn morton_4x4_full_table() {
69        // Expected offsets from the spec conformance table.
70        // Entry format: ([i_0, i_1], expected_offset)
71        let table: &[([u64; 2], u64)] = &[
72            ([0, 0], 0),
73            ([1, 0], 1),
74            ([0, 1], 2),
75            ([1, 1], 3),
76            ([2, 0], 4),
77            ([3, 0], 5),
78            ([2, 1], 6),
79            ([3, 1], 7),
80            ([0, 2], 8),
81            ([1, 2], 9),
82            ([0, 3], 10),
83            ([1, 3], 11),
84            ([2, 2], 12),
85            ([3, 2], 13),
86            ([2, 3], 14),
87            ([3, 3], 15),
88        ];
89        let (layout, shape) = layout_2d_2bits();
90        for &(index, expected) in table {
91            let got = layout.element_offset(&index, &shape).unwrap();
92            assert_eq!(
93                got, expected,
94                "morton[{},{}]: expected {expected}, got {got}",
95                index[0], index[1]
96            );
97        }
98    }
99
100    // Spec example from task description:
101    // shape [4,4], morton_bits [2,2], element [2,3]:
102    // i_0=2=0b10, i_1=3=0b11
103    // bit0(i_0)=0, bit0(i_1)=1, bit1(i_0)=1, bit1(i_1)=1 → 0b1110 = 14
104    #[test]
105    fn morton_4x4_spec_example_2_3() {
106        let (layout, shape) = layout_2d_2bits();
107        assert_eq!(layout.element_offset(&[2, 3], &shape).unwrap(), 14);
108    }
109
110    // Rank-3 Morton (morton_bits=[1,1,1], shape=[2,2,2]):
111    // i_0 → bit 0,3; i_1 → bit 1,4; i_2 → bit 2,5
112    // Each dimension contributes 1 bit.
113    #[test]
114    fn morton_2x2x2_full_table() {
115        let layout = MortonLayout::new(vec![1, 1, 1]).unwrap();
116        let shape = Shape::new(vec![2, 2, 2]).unwrap();
117        let table: &[([u64; 3], u64)] = &[
118            ([0, 0, 0], 0),
119            ([1, 0, 0], 1),
120            ([0, 1, 0], 2),
121            ([1, 1, 0], 3),
122            ([0, 0, 1], 4),
123            ([1, 0, 1], 5),
124            ([0, 1, 1], 6),
125            ([1, 1, 1], 7),
126        ];
127        for &(index, expected) in table {
128            let got = layout.element_offset(&index, &shape).unwrap();
129            assert_eq!(
130                got, expected,
131                "morton[{},{},{}]: expected {expected}, got {got}",
132                index[0], index[1], index[2]
133            );
134        }
135    }
136
137    // Overflow guard: total_bits > 64 must return IndexArithmeticOverflow.
138    // morton_bits=[33, 33] → total = 66 > 64.
139    #[test]
140    fn morton_overflow_total_bits_exceeds_64() {
141        let layout = MortonLayout::new(vec![33, 33]).unwrap();
142        // Shape must accommodate 2^33 per dimension; use DYNAMIC to avoid
143        // a shape-construction overflow.
144        let _shape = Shape::new(vec![crate::DYNAMIC, crate::DYNAMIC]).unwrap();
145        // validate_index will reject DYNAMIC before the overflow check; use a
146        // shape whose dims are within index range but trigger the overflow path.
147        // Since validate_index fires first on DYNAMIC, construct with large static dims.
148        // 2^33 = 8_589_934_592
149        let shape2 = Shape::new(vec![8_589_934_592u64, 8_589_934_592u64]).unwrap();
150        let err = layout.element_offset(&[0, 0], &shape2).unwrap_err();
151        assert!(
152            matches!(err, Error::IndexArithmeticOverflow),
153            "expected IndexArithmeticOverflow, got {err:?}"
154        );
155    }
156
157    // Error: index rank mismatch.
158    #[test]
159    fn morton_rank_mismatch() {
160        let (layout, shape) = layout_2d_2bits();
161        let err = layout.element_offset(&[1], &shape).unwrap_err();
162        assert!(
163            matches!(err, Error::IndexRankMismatch { .. }),
164            "expected IndexRankMismatch, got {err:?}"
165        );
166    }
167
168    // Error: index out of range.
169    #[test]
170    fn morton_index_out_of_range() {
171        let (layout, shape) = layout_2d_2bits();
172        // Index [4,0] where shape[0]=4 → out of range.
173        let err = layout.element_offset(&[4, 0], &shape).unwrap_err();
174        assert!(
175            matches!(err, Error::IndexOutOfRange { dim: 0, .. }),
176            "expected IndexOutOfRange, got {err:?}"
177        );
178    }
179}
hurray_core/layout/addressing/morton.rs

hurray_core/layout/addressing/
morton.rs
