1use geo_types::Coord;
2use wide::u32x8;
3
4use crate::decoder::Morton;
5use crate::encoder::model::CurveParams;
6use crate::{Decoder, MltError, MltResult};
7
8const LANES: usize = 8;
9
10#[must_use]
18#[inline]
19pub fn interleave_bits(coord: Coord<u32>) -> u32 {
20 let mut sx = coord.x & 0xFFFF;
24 sx = (sx | (sx << 8)) & 0x00FF_00FF;
25 sx = (sx | (sx << 4)) & 0x0F0F_0F0F;
26 sx = (sx | (sx << 2)) & 0x3333_3333;
27 sx = (sx | (sx << 1)) & 0x5555_5555;
28
29 let mut sy = coord.y & 0xFFFF;
30 sy = (sy | (sy << 8)) & 0x00FF_00FF;
31 sy = (sy | (sy << 4)) & 0x0F0F_0F0F;
32 sy = (sy | (sy << 2)) & 0x3333_3333;
33 sy = (sy | (sy << 1)) & 0x5555_5555;
34
35 sx | (sy << 1)
36}
37
38#[cfg(feature = "unstable-v2")]
40#[must_use]
41#[inline]
42pub fn interleave_u32(even: u32, odd: u32) -> u64 {
43 spread_u32(even) | (spread_u32(odd) << 1)
44}
45
46#[cfg(feature = "unstable-v2")]
48#[must_use]
49#[inline]
50pub fn deinterleave_u64(code: u64) -> (u32, u32) {
51 (compact_u64(code), compact_u64(code >> 1))
52}
53
54#[cfg(feature = "unstable-v2")]
55fn spread_u32(value: u32) -> u64 {
56 let mut s = u64::from(value);
57 s = (s | (s << 16)) & 0x0000_FFFF_0000_FFFF;
58 s = (s | (s << 8)) & 0x00FF_00FF_00FF_00FF;
59 s = (s | (s << 4)) & 0x0F0F_0F0F_0F0F_0F0F;
60 s = (s | (s << 2)) & 0x3333_3333_3333_3333;
61 s = (s | (s << 1)) & 0x5555_5555_5555_5555;
62 s
63}
64
65#[cfg(feature = "unstable-v2")]
66fn compact_u64(code: u64) -> u32 {
67 let mut c = code & 0x5555_5555_5555_5555;
68 c = (c | (c >> 1)) & 0x3333_3333_3333_3333;
69 c = (c | (c >> 2)) & 0x0F0F_0F0F_0F0F_0F0F;
70 c = (c | (c >> 4)) & 0x00FF_00FF_00FF_00FF;
71 c = (c | (c >> 8)) & 0x0000_FFFF_0000_FFFF;
72 c = (c | (c >> 16)) & 0x0000_0000_FFFF_FFFF;
73 u32::try_from(c).expect("infallible: the last mask keeps only the low 32 bits")
74}
75
76#[must_use]
87#[inline]
88pub fn morton_sort_key(c: Coord<i32>, params: CurveParams) -> u32 {
89 debug_assert!(params.bits >= 1);
90 #[expect(
91 clippy::cast_possible_truncation,
92 clippy::cast_sign_loss,
93 reason = "shift brings value into [0, extent]; masked to 16 bits immediately after"
94 )]
95 let sx = ((i64::from(c.x) + i64::from(params.shift)) as u32) & 0xFFFF;
96 #[expect(
97 clippy::cast_possible_truncation,
98 clippy::cast_sign_loss,
99 reason = "shift brings value into [0, extent]; masked to 16 bits immediately after"
100 )]
101 let sy = ((i64::from(c.y) + i64::from(params.shift)) as u32) & 0xFFFF;
102 interleave_bits((sx, sy).into())
103}
104
105impl Morton {
107 pub fn from_vertices(vertices: &[i32]) -> MltResult<Self> {
112 let min_v = vertices.iter().copied().min().unwrap_or(0);
113 let max_v = vertices.iter().copied().max().unwrap_or(0);
114 let shift: u32 = if min_v < 0 { min_v.unsigned_abs() } else { 0 };
115 let tile_extent = i64::from(max_v) + i64::from(shift);
116 let extent =
117 u32::try_from(tile_extent).expect("infallible: max_v + shift lies in 0..=u32::MAX");
118 let required_bits = extent.bit_width();
122 if required_bits > 16 {
123 return Err(MltError::VertexMortonNotCompatibleWithExtent {
124 extent,
125 required_bits,
126 });
127 }
128 Self::new(required_bits, shift)
129 }
130
131 #[inline]
136 pub fn encode_morton(self, x: i32, y: i32) -> MltResult<u32> {
137 let sx = u32::try_from(i64::from(x) + i64::from(self.shift))?;
138 let sy = u32::try_from(i64::from(y) + i64::from(self.shift))?;
139 let mut code = 0u32;
140 for i in 0..self.bits {
141 code |= ((sx >> i) & 1) << (2 * i);
143 code |= ((sy >> i) & 1) << (2 * i + 1);
144 }
145 Ok(code)
146 }
147}
148
149impl Morton {
150 #[inline]
152 fn decode_one(self, morton_code: u32) -> Coord<i32> {
153 let mut x = 0u32;
154 let mut y = 0u32;
155 for i in 0..self.bits {
156 let bit_mask = 1u32 << (2 * i);
157 x |= (morton_code & bit_mask) >> i;
158 y |= ((morton_code >> 1) & bit_mask) >> i;
159 }
160 Coord::<i32> {
161 x: x.wrapping_sub(self.shift).cast_signed(),
162 y: y.wrapping_sub(self.shift).cast_signed(),
163 }
164 }
165
166 pub fn decode_codes(self, data: &[u32], dec: &mut Decoder) -> MltResult<Vec<i32>> {
172 let alloc_size = data.len() * 2;
173 let mut out = dec.alloc(alloc_size)?;
174 let shift_vec = u32x8::splat(self.shift);
175
176 let (chunks, remainder) = data.as_chunks::<LANES>();
177
178 for &chunk in chunks {
179 self.decode_chunk(chunk, shift_vec, &mut out);
180 }
181
182 for &code in remainder {
184 let coord = self.decode_one(code);
185 out.push(coord.x);
186 out.push(coord.y);
187 }
188
189 dec.adjust_alloc(&out, alloc_size)
190 .expect("infallible: two coordinates pushed per code fill alloc_size exactly");
191 Ok(out)
192 }
193
194 pub fn decode_delta(self, data: &[u32], dec: &mut Decoder) -> MltResult<Vec<i32>> {
201 let alloc_size = data.len() * 2;
202 let mut out = dec.alloc(alloc_size)?;
203 let shift_vec = u32x8::splat(self.shift);
204
205 let mut prev = 0i32;
206 let (chunks, remainder) = data.as_chunks::<LANES>();
207
208 for chunk in chunks {
209 let mut buf = [0u32; LANES];
211 for (b, &d) in buf.iter_mut().zip(chunk.iter()) {
212 prev = prev.wrapping_add(d.cast_signed());
213 *b = prev.cast_unsigned();
214 }
215 self.decode_chunk(buf, shift_vec, &mut out);
216 }
217
218 for &d in remainder {
220 prev = prev.wrapping_add(d.cast_signed());
221 let coord = self.decode_one(prev.cast_unsigned());
222 out.push(coord.x);
223 out.push(coord.y);
224 }
225
226 dec.adjust_alloc(&out, alloc_size)
227 .expect("infallible: two coordinates pushed per code fill alloc_size exactly");
228 Ok(out)
229 }
230
231 #[inline]
236 fn decode_chunk(self, buf: [u32; LANES], shift_vec: u32x8, out: &mut Vec<i32>) {
237 let codes = u32x8::from(buf);
238 let codes_y = codes >> 1;
240
241 let mut x_vec = u32x8::ZERO;
242 let mut y_vec = u32x8::ZERO;
243
244 for i in 0..self.bits {
245 let bit_mask = u32x8::splat(1u32 << (2 * i));
247 x_vec |= (codes & bit_mask) >> i;
249 y_vec |= (codes_y & bit_mask) >> i;
250 }
251
252 let xs: [u32; LANES] = (x_vec - shift_vec).into();
253 let ys: [u32; LANES] = (y_vec - shift_vec).into();
254
255 for lane in 0..LANES {
256 out.push(xs[lane].cast_signed());
257 out.push(ys[lane].cast_signed());
258 }
259 }
260}
261
262#[cfg(test)]
263mod tests {
264
265 use rstest::rstest;
266
267 use super::*;
268 use crate::test_helpers::{dec, starved_dec};
269
270 const fn c(x: i32, y: i32) -> Coord<i32> {
271 Coord::<i32> { x, y }
272 }
273
274 const fn p(shift: u32, bits: u32) -> CurveParams {
275 CurveParams { shift, bits }
276 }
277
278 fn spread_bits(mut tx: u32) -> u32 {
281 tx = (tx | (tx << 8)) & 0x00FF_00FF;
282 tx = (tx | (tx << 4)) & 0x0F0F_0F0F;
283 tx = (tx | (tx << 2)) & 0x3333_3333;
284 tx = (tx | (tx << 1)) & 0x5555_5555;
285 tx
286 }
287
288 fn compact_bits(mut tx: u32) -> u32 {
289 tx &= 0x5555_5555;
290 tx = (tx | (tx >> 1)) & 0x3333_3333;
291 tx = (tx | (tx >> 2)) & 0x0F0F_0F0F;
292 tx = (tx | (tx >> 4)) & 0x00FF_00FF;
293 tx = (tx | (tx >> 8)) & 0x0000_FFFF;
294 tx
295 }
296
297 #[test]
298 fn spread_then_compact_is_identity() {
299 for x in 0u32..=0xFFFF {
300 assert_eq!(compact_bits(spread_bits(x)), x, "round-trip failed for {x}");
301 }
302 }
303
304 #[test]
305 fn spread_bits_places_bit0_at_position0() {
306 assert_eq!(spread_bits(1), 1);
307 }
308
309 #[test]
310 fn spread_bits_places_bit1_at_position2() {
311 assert_eq!(spread_bits(2), 4);
312 }
313
314 #[test]
315 fn spread_bits_places_bit2_at_position4() {
316 assert_eq!(spread_bits(4), 16);
317 }
318
319 #[cfg(feature = "unstable-v2")]
320 #[rstest]
321 #[case::zero(0, 0, 0)]
322 #[case::even_only(1, 0, 1)]
323 #[case::odd_only(0, 1, 2)]
324 #[case::one_byte_varint(15, 7, 0x7F)]
325 #[case::u32_max_both(u32::MAX, u32::MAX, u64::MAX)]
326 #[case::u32_max_even(u32::MAX, 0, 0x5555_5555_5555_5555)]
327 fn interleave_u32_places_even_and_odd_bits(
328 #[case] even: u32,
329 #[case] odd: u32,
330 #[case] code: u64,
331 ) {
332 assert_eq!(interleave_u32(even, odd), code);
333 assert_eq!(deinterleave_u64(code), (even, odd));
334 }
335
336 #[test]
337 fn origin_maps_to_zero() {
338 assert_eq!(morton_sort_key(c(0, 0), p(0, 16)), 0);
339 }
340
341 #[test]
342 fn x_axis_produces_even_bits() {
343 assert_eq!(
344 morton_sort_key(c(1, 0), p(0, 16)),
345 1,
346 "x bit 0 lands at Morton bit 0"
347 );
348 assert_eq!(
349 morton_sort_key(c(2, 0), p(0, 16)),
350 4,
351 "x bit 1 lands at Morton bit 2"
352 );
353 }
354
355 #[test]
356 fn y_axis_produces_odd_bits() {
357 assert_eq!(
358 morton_sort_key(c(0, 1), p(0, 16)),
359 2,
360 "y bit 0 lands at Morton bit 1"
361 );
362 assert_eq!(
363 morton_sort_key(c(0, 2), p(0, 16)),
364 8,
365 "y bit 1 lands at Morton bit 3"
366 );
367 }
368
369 #[test]
370 fn negative_coords_shift_correctly() {
371 assert_eq!(
372 morton_sort_key(c(-1, -1), p(1, 16)),
373 0,
374 "a shift of 1 maps (-1, -1) onto (0, 0)"
375 );
376 assert_eq!(
377 morton_sort_key(c(-1, 0), p(1, 16)),
378 2,
379 "a shift of 1 maps (-1, 0) onto (0, 1)"
380 );
381 }
382
383 #[test]
384 fn spatial_locality_z_order() {
385 let k00 = morton_sort_key(c(0, 0), p(0, 16));
386 let k10 = morton_sort_key(c(1, 0), p(0, 16));
387 let k01 = morton_sort_key(c(0, 1), p(0, 16));
388 let k11 = morton_sort_key(c(1, 1), p(0, 16));
389 assert!(k00 < k10, "(0,0) precedes (1,0) in Z-order");
390 assert!(k10 < k01, "(1,0) precedes (0,1) in Z-order");
391 assert!(k01 < k11, "(0,1) precedes (1,1) in Z-order");
392 }
393
394 #[test]
395 fn interleave_round_trips_via_deinterleave() {
396 for x in 0u32..16 {
397 for y in 0u32..16 {
398 let code = interleave_bits((x, y).into());
399 let mut rx = 0u32;
400 let mut ry = 0u32;
401 for bit in 0..16 {
402 rx |= ((code >> (2 * bit)) & 1) << bit;
403 ry |= ((code >> (2 * bit + 1)) & 1) << bit;
404 }
405 assert_eq!(rx, x, "x mismatch for ({x}, {y})");
406 assert_eq!(ry, y, "y mismatch for ({x}, {y})");
407 }
408 }
409 }
410
411 const NUM_BITS: u32 = 15;
414 const COORD_SHIFT: u32 = 1 << (NUM_BITS - 1);
415 const MORTON: Morton = Morton {
416 bits: NUM_BITS,
417 shift: COORD_SHIFT,
418 };
419
420 #[must_use]
421 #[inline]
422 pub fn encode_morton_15(coord: Coord<u32>) -> u32 {
423 let mut code = 0u32;
424 for bit in 0..15 {
425 code |= ((coord.x >> bit) & 1) << (2 * bit);
426 code |= ((coord.y >> bit) & 1) << (2 * bit + 1);
427 }
428 code
429 }
430
431 #[test]
432 fn test_decode_morton_codes_empty() {
433 assert_eq!(
434 MORTON.decode_codes(&[], &mut dec()).unwrap(),
435 [] as [i32; 0]
436 );
437 }
438
439 #[test]
440 fn test_decode_morton_codes_origin() {
441 let code = encode_morton_15((COORD_SHIFT, COORD_SHIFT).into());
442 let decoded = MORTON.decode_codes(&[code], &mut dec()).unwrap();
443 assert_eq!(
444 decoded,
445 [0, 0],
446 "the code for (COORD_SHIFT, COORD_SHIFT) sits at the origin"
447 );
448 }
449
450 #[test]
451 fn test_decode_morton_codes_known_values() {
452 let x: u32 = 1;
453 let y: u32 = 2;
454 let code = encode_morton_15((x, y).into());
455 let expected_x = x.cast_signed() - COORD_SHIFT.cast_signed();
456 let expected_y = y.cast_signed() - COORD_SHIFT.cast_signed();
457 let decoded = MORTON.decode_codes(&[code], &mut dec()).unwrap();
458 assert_eq!(
459 decoded,
460 [expected_x, expected_y],
461 "decoding subtracts COORD_SHIFT from each axis"
462 );
463 }
464
465 #[test]
466 fn test_decode_morton_codes_scalar_tail() {
467 let pairs: [Coord<u32>; _] = [(0, 1).into(), (2, 3).into(), (4, 5).into()];
468 let codes: Vec<u32> = pairs.iter().map(|&c| encode_morton_15(c)).collect();
469 let result = MORTON.decode_codes(&codes, &mut dec()).unwrap();
470 let expected = expected_coords(&pairs);
471 assert_eq!(result, expected, "3 codes take the scalar tail alone");
472 }
473
474 #[test]
475 fn test_decode_morton_codes_full_simd_chunk() {
476 let pairs: [Coord<u32>; _] = [
477 (0, 0).into(),
478 (1, 0).into(),
479 (0, 1).into(),
480 (1, 1).into(),
481 (2, 3).into(),
482 (7, 5).into(),
483 (10, 9).into(),
484 (15, 15).into(),
485 ];
486 let codes: Vec<u32> = pairs.iter().map(|&c| encode_morton_15(c)).collect();
487 let result = MORTON.decode_codes(&codes, &mut dec()).unwrap();
488 let expected = expected_coords(&pairs);
489 assert_eq!(result, expected, "8 codes fill one SIMD chunk with no tail");
490 }
491
492 #[test]
493 fn test_decode_morton_codes_simd_plus_tail() {
494 let pairs: Vec<Coord<u32>> = (0..11u32)
495 .map(|i| (i * 3 % 100, i * 7 % 100).into())
496 .collect();
497 let codes: Vec<u32> = pairs.iter().map(|&c| encode_morton_15(c)).collect();
498 let result = MORTON.decode_codes(&codes, &mut dec()).unwrap();
499 let expected = expected_coords(&pairs);
500 assert_eq!(
501 result, expected,
502 "11 codes take one SIMD chunk plus a tail of 3"
503 );
504 }
505
506 #[test]
507 fn test_decode_morton_delta_empty() {
508 assert_eq!(
509 MORTON.decode_delta(&[], &mut dec()).unwrap(),
510 [] as [i32; 0]
511 );
512 }
513
514 #[test]
515 fn test_decode_morton_delta_identity_with_zero_deltas() {
516 let deltas = vec![0u32; 3];
517 let result = MORTON.decode_delta(&deltas, &mut dec()).unwrap();
518 let shift = -COORD_SHIFT.cast_signed();
519 assert_eq!(
520 result,
521 vec![shift, shift, shift, shift, shift, shift],
522 "zero deltas resolve every code to 0, which decodes to (-COORD_SHIFT, -COORD_SHIFT)"
523 );
524 }
525
526 #[test]
527 fn test_decode_morton_delta_matches_codes_after_prefix_sum() {
528 let pairs: Vec<Coord<u32>> = (0..11u32)
529 .map(|i| (i * 5 % 200, i * 9 % 200).into())
530 .collect();
531 let codes: Vec<u32> = pairs.iter().map(|&c| encode_morton_15(c)).collect();
532 let deltas = signed_deltas(&codes);
533
534 let from_codes = MORTON.decode_codes(&codes, &mut dec()).unwrap();
535 let from_deltas = MORTON.decode_delta(&deltas, &mut dec()).unwrap();
536 assert_eq!(
537 from_codes, from_deltas,
538 "the prefix sum rebuilds the absolute codes"
539 );
540 }
541
542 #[test]
543 fn test_decode_morton_delta_scalar_tail() {
544 let codes: Vec<u32> = vec![
545 encode_morton_15((10, 20).into()),
546 encode_morton_15((30, 40).into()),
547 encode_morton_15((50, 60).into()),
548 ];
549 let deltas = signed_deltas(&codes);
550 let from_codes = MORTON.decode_codes(&codes, &mut dec()).unwrap();
551 let from_deltas = MORTON.decode_delta(&deltas, &mut dec()).unwrap();
552 assert_eq!(
553 from_codes, from_deltas,
554 "3 deltas take the scalar tail alone"
555 );
556 }
557
558 #[test]
559 fn test_decode_morton_delta_wrapping() {
560 let code_a = encode_morton_15((500, 300).into());
561 let code_b = encode_morton_15((10, 10).into());
562 let delta_b = code_b
563 .cast_signed()
564 .wrapping_sub(code_a.cast_signed())
565 .cast_unsigned();
566 assert_eq!(
567 MORTON.decode_delta(&[code_a, delta_b], &mut dec()).unwrap(),
568 MORTON.decode_codes(&[code_a, code_b], &mut dec()).unwrap(),
569 "a delta onto a smaller code wraps and still resolves"
570 );
571 }
572
573 #[test]
574 fn decode_codes_past_the_memory_budget_is_rejected() {
575 let err = MORTON
576 .decode_codes(&[0, 1], &mut starved_dec())
577 .unwrap_err();
578 assert!(
579 matches!(err, MltError::MemoryLimitExceeded { .. }),
580 "{err:?}"
581 );
582 }
583
584 #[test]
585 fn decode_delta_past_the_memory_budget_is_rejected() {
586 let err = MORTON
587 .decode_delta(&[0, 1], &mut starved_dec())
588 .unwrap_err();
589 assert!(
590 matches!(err, MltError::MemoryLimitExceeded { .. }),
591 "{err:?}"
592 );
593 }
594
595 #[rstest]
598 #[case::empty(&[], 0, 0)]
599 #[case::non_negative_needs_no_shift(&[0, 100], 0, 7)]
600 #[case::negative_min_shifts_onto_zero(&[-5, 10], 5, 4)]
601 #[case::widest_accepted_extent(&[0, 65_535], 0, 16)]
602 fn from_vertices_derives_shift_and_bits(
603 #[case] vertices: &[i32],
604 #[case] shift: u32,
605 #[case] bits: u32,
606 ) {
607 let morton = Morton::from_vertices(vertices).unwrap();
608 assert_eq!(
609 (morton.shift, morton.bits),
610 (shift, bits),
611 "shift cancels a negative min, bits covers max + shift"
612 );
613 }
614
615 #[rstest]
616 #[case::past_the_cap_outright(&[0, 65_536])]
617 #[case::pushed_past_the_cap_by_the_shift(&[-1, 65_535])]
618 fn from_vertices_rejects_an_extent_wider_than_sixteen_bits(#[case] vertices: &[i32]) {
619 let err = Morton::from_vertices(vertices).unwrap_err();
620 assert!(
621 matches!(
622 err,
623 MltError::VertexMortonNotCompatibleWithExtent {
624 extent: 65_536,
625 required_bits: 17,
626 }
627 ),
628 "{err:?}"
629 );
630 }
631
632 #[rstest]
633 #[case::origin(0, 0)]
634 #[case::x_only(1, 0)]
635 #[case::y_only(0, 1)]
636 #[case::both_axes(500, 300)]
637 #[case::negative_coords(-4000, -9000)]
638 fn encode_morton_round_trips_through_decode_codes(#[case] x: i32, #[case] y: i32) {
639 let code = MORTON.encode_morton(x, y).unwrap();
640 assert_eq!(MORTON.decode_codes(&[code], &mut dec()).unwrap(), [x, y]);
641 }
642
643 #[rstest]
644 #[case::x_below_the_shift(-20_000, 0)]
645 #[case::y_below_the_shift(0, -20_000)]
646 fn encode_morton_rejects_a_coordinate_the_shift_cannot_lift(#[case] x: i32, #[case] y: i32) {
647 let err = MORTON.encode_morton(x, y).unwrap_err();
648 assert!(matches!(err, MltError::TryFromIntError(_)), "{err:?}");
649 }
650
651 fn expected_coords(pairs: &[Coord<u32>]) -> Vec<i32> {
652 pairs
653 .iter()
654 .flat_map(|&Coord { x, y }| {
655 [
656 x.cast_signed() - COORD_SHIFT.cast_signed(),
657 y.cast_signed() - COORD_SHIFT.cast_signed(),
658 ]
659 })
660 .collect()
661 }
662
663 fn signed_deltas(codes: &[u32]) -> Vec<u32> {
664 let mut prev = 0i32;
665 codes
666 .iter()
667 .map(|&c| {
668 let delta = c.cast_signed().wrapping_sub(prev).cast_unsigned();
669 prev = c.cast_signed();
670 delta
671 })
672 .collect()
673 }
674}