z_order.rs

use num_traits::PrimInt;
use std::{
  mem::size_of,
  ops::{Shl, Shr},
};

#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd)]
pub struct ZOrderIndex<T: PrimInt>(T);

impl<T: PrimInt> ZOrderIndex<T> {
  pub const MAX_BITS_PER_ELEMENT: usize = size_of::<T>() * 8 / 3;
  pub const MAX_USABLE_BITS: usize = Self::MAX_BITS_PER_ELEMENT * 3;

  pub const ELEMENT_MIN: i32 = !0 & !Self::ELEMENT_MAX;
  pub const ELEMENT_MAX: i32 = !(!0 << (Self::MAX_BITS_PER_ELEMENT - 1));

  pub fn new(x: i32, y: i32, z: i32) -> Option<Self> {
    if (x < Self::ELEMENT_MIN || y < Self::ELEMENT_MIN || z < Self::ELEMENT_MIN)
      || (x > Self::ELEMENT_MAX || y > Self::ELEMENT_MAX || z > Self::ELEMENT_MAX)
    {
      return None;
    }

    let mut order = T::zero();
    let mut set_bit = T::one();

    for i in 0..(Self::MAX_BITS_PER_ELEMENT - 1) {
      if (x >> i) & 1 == 1 {
        order = order | set_bit;
      }
      set_bit = set_bit << 1;
      if (y >> i) & 1 == 1 {
        order = order | set_bit;
      }
      set_bit = set_bit << 1;
      if (z >> i) & 1 == 1 {
        order = order | set_bit;
      }
      set_bit = set_bit << 1;
    }

    // The final bit for each element is inverted, so that
    // negative indices are ordered before postitive ones.
    if (x >> Self::MAX_BITS_PER_ELEMENT - 1) & 1 == 0 {
      order = order | set_bit;
    }
    set_bit = set_bit << 1;
    if (y >> Self::MAX_BITS_PER_ELEMENT - 1) & 1 == 0 {
      order = order | set_bit;
    }
    set_bit = set_bit << 1;
    if (z >> Self::MAX_BITS_PER_ELEMENT - 1) & 1 == 0 {
      order = order | set_bit;
    }

    Some(ZOrderIndex(order))
  }

  pub fn x(&self) -> i32 {
    self.decode(0)
  }

  pub fn y(&self) -> i32 {
    self.decode(1)
  }

  pub fn z(&self) -> i32 {
    self.decode(2)
  }

  fn decode(&self, offset: usize) -> i32 {
    let mut n = 0;
    for i in 0..(Self::MAX_BITS_PER_ELEMENT - 1) {
      if ((self.0 >> (offset + i * 3)) & T::one()).is_one() {
        n |= 1 << i;
      }
    }
    // If the sign bit is "set", fill out the rest of the significant bits.
    // NOTE: The sign bit is considered "set" if it's zero.
    if ((self.0 >> (offset + (Self::MAX_BITS_PER_ELEMENT - 1) * 3)) & T::one()).is_zero() {
      for i in (Self::MAX_BITS_PER_ELEMENT - 1)..(size_of::<i32>() * 8) {
        n |= 1 << i;
      }
    }
    n
  }

  pub fn raw_order(&self) -> T {
    self.0
  }
}

impl<T: PrimInt> Shl<usize> for ZOrderIndex<T> {
  type Output = Self;
  fn shl(self, rhs: usize) -> Self {
    ZOrderIndex(self.0 << (rhs * 3))
  }
}

#[cfg(test)]
mod tests {
  use super::*;

  #[test]
  fn test() {
    assert_eq!(ZOrderIndex::<u32>::MAX_BITS_PER_ELEMENT, 10);
    assert_eq!(ZOrderIndex::<u32>::MAX_USABLE_BITS, 30);
    assert_eq!(ZOrderIndex::<u32>::ELEMENT_MIN, -512);
    assert_eq!(ZOrderIndex::<u32>::ELEMENT_MAX, 511);

    assert_eq!(ZOrderIndex::<u16>::MAX_BITS_PER_ELEMENT, 5);
    assert_eq!(ZOrderIndex::<u16>::MAX_USABLE_BITS, 15);
    assert_eq!(ZOrderIndex::<u16>::ELEMENT_MIN, -16);
    assert_eq!(ZOrderIndex::<u16>::ELEMENT_MAX, 15);

    let z = ZOrderIndex::<u16>::new(0, -10, 15).unwrap();
    assert_eq!(z.x(), 0);
    assert_eq!(z.y(), -10);
    assert_eq!(z.z(), 15);

    //                    -16 +  16 + 8 + 4 + 2 + 1
    //                          -------------------
    //                      x =   1   0   0   0   0
    //                      y =  0   0   1   1   0
    //                      z = 1   1   1   1   1
    assert_eq!(z.raw_order(), 0b101_100_110_110_100);
  }
}