use super::{LOOKUP_IJ, get_u_norm, get_v_norm};
use crate::geometry::{
K_INVERT_MASK, K_MAX_CELL_LEVEL, K_SWAP_MASK, LOOKUP_POS, LonLat, S2Point, ST_TO_UV, UV_TO_ST,
face_uv_to_xyz, ij_to_st, si_ti_to_st, st_to_ij, xyz_to_face_uv,
};
use alloc::{string::String, vec::Vec};
use core::{fmt, ops::Deref};
use libm::{fmax, fmin};
use s2json::{BBox, VectorPoint};
use serde::{Deserialize, Serialize};
pub type CellId = S2CellId;
pub const K_FACE_BITS: u8 = 3;
pub const K_NUM_FACES: u8 = 6;
pub const K_MAX_LEVEL: u64 = K_MAX_CELL_LEVEL as u64; pub const K_POS_BITS: u64 = 2 * K_MAX_LEVEL + 1;
pub const K_MAX_SIZE: u32 = 1 << K_MAX_LEVEL;
pub const K_LOOKUP_BITS: u8 = 4;
pub const K_WRAP_OFFSET: u64 = (K_NUM_FACES as u64) << K_POS_BITS;
#[derive(
Debug, Default, Copy, Clone, PartialEq, PartialOrd, Ord, Eq, Hash, Serialize, Deserialize,
)]
#[repr(C)]
pub struct S2CellId {
pub id: u64,
}
impl Deref for S2CellId {
type Target = u64;
fn deref(&self) -> &Self::Target {
&self.id
}
}
impl S2CellId {
pub fn new(id: u64) -> Self {
S2CellId { id }
}
pub fn none() -> Self {
0_u64.into()
}
pub fn sentinel() -> Self {
(!0_u64).into()
}
pub fn from_face(face: u8) -> S2CellId {
(((face as u64) << K_POS_BITS) + lsb_for_level(0)).into()
}
pub fn from_lon_lat<M: Clone + Default>(ll: &LonLat<M>) -> S2CellId {
S2CellId::from_s2_point(&ll.to_point())
}
pub fn from_s2_point(p: &S2Point) -> Self {
let (face, u, v) = xyz_to_face_uv(p);
let i: u32 = st_to_ij(UV_TO_ST(u));
let j: u32 = st_to_ij(UV_TO_ST(v));
Self::from_face_ij(face, i, j, None)
}
pub fn from_face_uv(face: u8, u: f64, v: f64) -> Self {
S2CellId::from_face_st(face, UV_TO_ST(u), UV_TO_ST(v))
}
pub fn from_face_st(face: u8, s: f64, t: f64) -> Self {
let i: u32 = st_to_ij(s);
let j: u32 = st_to_ij(t);
Self::from_face_ij(face, i, j, None)
}
pub fn from_face_ij(face: u8, i: u32, j: u32, level: Option<u8>) -> S2CellId {
let mut i = i as u64;
let mut j = j as u64;
if let Some(level) = level {
i <<= K_MAX_LEVEL - level as u64;
j <<= K_MAX_LEVEL - level as u64;
}
let mut n: u64 = (face as u64) << (K_POS_BITS - 1);
let mut bits: u64 = (face & K_SWAP_MASK) as u64;
let mut k: u8 = 7;
loop {
let mask: u64 = (1 << 4) - 1;
bits += ((i >> (k * 4)) & mask) << (4 + 2);
bits += ((j >> (k * 4)) & mask) << 2;
bits = LOOKUP_POS[bits as usize] as u64;
n |= (bits >> 2) << (k * 2 * 4);
bits &= (K_SWAP_MASK | K_INVERT_MASK) as u64;
if k == 0 {
break;
}
k -= 1;
}
let id: S2CellId = ((n << 1) + 1).into();
if let Some(level) = level { id.parent(Some(level)) } else { id }
}
pub fn from_distance(distance: u64, level: Option<u8>) -> S2CellId {
let mut level: u64 = level.unwrap_or(K_MAX_LEVEL as u8) as u64;
level = 2 * (K_MAX_LEVEL - level);
((distance << (level + 1)) + (1 << level)).into()
}
pub fn to_face_ij(&self) -> (u8, u8, u32, u32) {
let level = self.level();
let (face, i, j, _or) = self.to_face_ij_orientation(Some(level));
(face, level, i, j)
}
pub fn to_face_ij_orientation(&self, level: Option<u8>) -> (u8, u32, u32, u8) {
let mut i: u32 = 0;
let mut j: u32 = 0;
let face: u8 = self.face();
let mut bits: u64 = (face & K_SWAP_MASK) as u64;
let mut k = 7;
while k >= 0 {
let nbits: u64 = if k == 7 { 2 } else { 4 };
let kk: u64 = k as u64;
bits += ((self.id >> (kk * 8 + 1)) & ((1 << (2 * nbits)) - 1)) << 2;
bits = LOOKUP_IJ[bits as usize] as u64;
i += (bits as u32 >> K_NUM_FACES as u32) << (kk * 4);
j += (((bits >> 2) & 15) << (kk * 4)) as u32;
bits &= (K_SWAP_MASK | K_INVERT_MASK) as u64;
k -= 1;
}
let lsb = self.id & ((!self.id).wrapping_add(1));
if (lsb & 0x1111111111111110) != 0 {
bits ^= K_SWAP_MASK as u64;
}
if let Some(level) = level {
i >>= K_MAX_LEVEL as u32 - level as u32;
j >>= K_MAX_LEVEL as u32 - level as u32;
}
(face, i, j, bits as u8)
}
pub fn get_edges(&self) -> [S2Point; 4] {
let mut edges = self.get_edges_raw();
for e in edges.iter_mut() {
e.normalize();
}
edges
}
pub fn get_edges_raw(&self) -> [S2Point; 4] {
let f = self.face();
let BBox { left: u_low, right: u_high, bottom: v_low, top: v_high } = self.get_bound_uv();
[
get_v_norm::<S2Point>(f, v_low),
get_u_norm::<S2Point>(f, u_high),
get_v_norm::<S2Point>(f, v_high).invert(),
get_u_norm::<S2Point>(f, u_low).invert(),
]
}
pub fn get_vertices(&self) -> [S2Point; 4] {
self.get_vertices_raw().map(|mut v| {
v.normalize();
v
})
}
pub fn get_vertices_raw(&self) -> [S2Point; 4] {
let f = self.face();
let BBox { left: u_low, right: u_high, bottom: v_low, top: v_high } = self.get_bound_uv();
[
face_uv_to_xyz(f, u_low, v_low),
face_uv_to_xyz(f, u_high, v_low),
face_uv_to_xyz(f, u_high, v_high),
face_uv_to_xyz(f, u_low, v_high),
]
}
pub fn get_bound_uv(&self) -> BBox {
let (_, i, j, _) = self.to_face_ij_orientation(None);
ij_level_to_bound_uv(i, j, self.level())
}
pub fn get_size_ij(&self) -> u64 {
1_u64 << (K_MAX_LEVEL - self.level() as u64)
}
pub fn child_position(&self, level: u8) -> u8 {
if !self.is_valid() || level > self.level() {
unreachable!();
}
((self.id >> (2 * (K_MAX_LEVEL - level as u64) + 1)) & 3) as u8
}
pub fn from_string(val: &str) -> S2CellId {
let val_bytes = val.as_bytes();
let level = (val.len() - 2) as u64;
if !(0..=K_MAX_LEVEL).contains(&level) {
return S2CellId::none();
}
let face: u8 = val_bytes[0] - b'0';
if !(0..=5).contains(&face) || val_bytes[1] != b'/' {
return S2CellId::none();
}
let mut id = S2CellId::from_face(face);
let mut i = 2;
while i < val.len() {
let child_pos = val_bytes[i] - b'0';
if !(0..=3).contains(&child_pos) {
return S2CellId::none();
}
id = id.child(child_pos);
i += 1;
}
id
}
pub fn child(&self, position: u8) -> S2CellId {
if !self.is_valid() || self.is_leaf() || position > 3 {
unreachable!();
}
let new_lsb = self.lsb() >> 2;
(self.id - (3 * new_lsb) + (2 * (position as u64) * new_lsb)).into()
}
pub fn face(&self) -> u8 {
(self.id >> K_POS_BITS) as u8
}
pub fn pos(&self) -> u64 {
self.id & (!0u64 >> K_FACE_BITS)
}
pub fn level(&self) -> u8 {
if self.id == 0 {
unreachable!();
}
K_MAX_LEVEL as u8 - (self.id.trailing_zeros() as u8 >> 1u8)
}
pub fn is_valid(&self) -> bool {
self.face() < K_NUM_FACES && (self.lsb() & 0x1555555555555555) != 0_u64
}
pub fn is_leaf(&self) -> bool {
(self.id & 1u64) != 0
}
pub fn to_point_raw(&self) -> S2Point {
let (face, u, v) = self.to_uv();
face_uv_to_xyz(face, u, v)
}
pub fn to_point(&self) -> S2Point {
let mut p = self.to_point_raw();
p.normalize();
p
}
pub fn to_st(&self) -> (u8, f64, f64) {
let (face, i, j, _or) = self.to_face_ij_orientation(None);
let s = ij_to_st(i);
let t = ij_to_st(j);
(face, s, t)
}
pub fn to_uv(&self) -> (u8, f64, f64) {
let (face, s, t) = self.to_st();
let u = ST_TO_UV(s);
let v = ST_TO_UV(t);
(face, u, v)
}
pub fn is_face(&self) -> bool {
(self.id & ((1 << 60) - 1)) == 0
}
pub fn distance(&self, lev: Option<u8>) -> f64 {
let level = lev.unwrap_or(self.level());
(self.id >> (2 * (30 - level) + 1)) as f64
}
pub fn children(&self, orientation: Option<u8>) -> [S2CellId; 4] {
let mut childs = [self.child(0), self.child(3), self.child(2), self.child(1)];
if let Some(orientation) = orientation
&& orientation == 0
{
childs.swap(1, 3);
}
childs
}
pub fn children_ij(face: u8, level: u8, i: u32, j: u32) -> [S2CellId; 4] {
let i = i << 1;
let j = j << 1;
let level = level + 1;
[
S2CellId::from_face_ij(face, i, j, Some(level)),
S2CellId::from_face_ij(face, i + 1, j, Some(level)),
S2CellId::from_face_ij(face, i, j + 1, Some(level)),
S2CellId::from_face_ij(face, i + 1, j + 1, Some(level)),
]
}
pub fn parent(&self, level: Option<u8>) -> S2CellId {
let new_lsb = match level {
Some(level) => 1 << (2 * (K_MAX_LEVEL - level as u64)),
None => (self.id & (!self.id + 1)) << 2,
};
((self.id & (!new_lsb + 1)) | new_lsb).into()
}
pub fn range(&self) -> (Self, Self) {
let id = self.id;
let lsb = id & (!id + 1);
((id - (lsb - 1)).into(), (id + (lsb - 1)).into())
}
pub fn contains(&self, other: S2CellId) -> bool {
let (min, max) = self.range();
other >= min && other <= max
}
pub fn contains_s2point(&self, p: &S2Point) -> bool {
self.contains(p.into())
}
pub fn intersects(&self, other: S2CellId) -> bool {
let (min_self, max_self) = self.range();
let (min_other, max_other) = other.range();
min_other <= max_self && max_other >= min_self
}
pub fn lsb(&self) -> u64 {
self.id & (!(self.id) + 1)
}
pub fn next(&self) -> S2CellId {
let next = self.id.wrapping_add(self.lsb() << 1);
if next < K_WRAP_OFFSET { next.into() } else { (next - K_WRAP_OFFSET).into() }
}
pub fn prev(&self) -> S2CellId {
let prev = self.id.wrapping_sub(self.lsb() << 1);
if prev < K_WRAP_OFFSET { prev.into() } else { (prev.wrapping_add(K_WRAP_OFFSET)).into() }
}
pub fn center_st(&self) -> (u8, f64, f64) {
let id = self.id;
let (face, i, j, _or) = self.to_face_ij_orientation(None);
let delta = if (id & 1) != 0 {
1
} else if ((i as u64 ^ (id >> 2)) & 1) != 0 {
2
} else {
0
};
let si = 2 * i + delta;
let ti = 2 * j + delta;
(face, si_ti_to_st(si), si_ti_to_st(ti))
}
pub fn bounds_st(&self, level: Option<u8>) -> BBox {
let level = level.unwrap_or(self.level());
let (_face, s, t) = self.center_st();
let half_size = size_st(level) * 0.5;
BBox::new(s - half_size, t - half_size, s + half_size, t + half_size)
}
pub fn neighbors(&self) -> [S2CellId; 4] {
let level = self.level();
let size = size_ij(level) as i32;
let (face, i, j, _or) = self.to_face_ij_orientation(None);
let i = i as i32;
let j = j as i32;
let j_minus_size = j - size;
let i_minus_size = i - size;
[
S2CellId::from_ij_same(face, i, j_minus_size, j_minus_size >= 0).parent(Some(level)),
S2CellId::from_ij_same(face, i + size, j, i + size < K_MAX_SIZE as i32)
.parent(Some(level)),
S2CellId::from_ij_same(face, i, j + size, j + size < K_MAX_SIZE as i32)
.parent(Some(level)),
S2CellId::from_ij_same(face, i_minus_size, j, i_minus_size >= 0).parent(Some(level)),
]
}
pub fn neighbors_ij(face: u8, i: u32, j: u32, level: u8) -> [S2CellId; 4] {
let size = size_ij(level) as i32;
let i = i as i32;
let j = j as i32;
let j_minus_size: i32 = j - size;
let i_minus_size: i32 = i - size;
[
S2CellId::from_ij_same(face, i, j_minus_size, j_minus_size >= 0).parent(Some(level)),
S2CellId::from_ij_same(face, i + size, j, i + size < K_MAX_SIZE as i32)
.parent(Some(level)),
S2CellId::from_ij_same(face, i, j + size, j + size < K_MAX_SIZE as i32)
.parent(Some(level)),
S2CellId::from_ij_same(face, i_minus_size, j, i_minus_size >= 0).parent(Some(level)),
]
}
pub fn from_ij_same(face: u8, i: i32, j: i32, same_face: bool) -> S2CellId {
if same_face {
S2CellId::from_face_ij(face, i as u32, j as u32, None)
} else {
S2CellId::from_face_ij_wrap(face, i, j)
}
}
pub fn from_face_ij_wrap(face: u8, i: i32, j: i32) -> S2CellId {
let i = i.clamp(-1, K_MAX_SIZE as i32);
let j = j.clamp(-1, K_MAX_SIZE as i32);
let k_scale = 1. / K_MAX_SIZE as f64;
let k_limit = 1. + 2.220_446_049_250_313e-16;
let u = fmax(
-k_limit,
fmin(k_limit, k_scale * (2. * (i as f64 - (K_MAX_SIZE as f64) / 2.) + 1.)),
);
let v = fmax(
-k_limit,
fmin(k_limit, k_scale * (2. * (j as f64 - (K_MAX_SIZE as f64) / 2.) + 1.)),
);
let (n_face, nu, nv) = xyz_to_face_uv(&face_uv_to_xyz(face, u, v));
S2CellId::from_face_ij(n_face, st_to_ij(0.5 * (nu + 1.)), st_to_ij(0.5 * (nv + 1.)), None)
}
pub fn vertex_neighbors(&self, level: Option<u8>) -> Vec<S2CellId> {
let level = level.unwrap_or(self.level());
let mut neighbors: Vec<S2CellId> = Vec::new();
let (face, i, j, _or) = self.to_face_ij_orientation(None);
let halfsize = size_ij(level + 1) as i32;
let size = halfsize << 1;
let isame;
let jsame;
let ioffset;
let joffset;
let i = i as i32;
let j = j as i32;
if (i & halfsize) != 0 {
ioffset = size;
isame = i + size < K_MAX_SIZE as i32;
} else {
ioffset = -(size);
isame = i - size >= 0;
}
if (j & halfsize) != 0 {
joffset = size;
jsame = j + size < K_MAX_SIZE as i32;
} else {
joffset = -size;
jsame = j - size >= 0;
}
let i_new: i32 = i + ioffset;
let j_new: i32 = j + joffset;
neighbors.push(self.parent(Some(level)));
neighbors.push(S2CellId::from_ij_same(face, i_new, j, isame).parent(Some(level)));
neighbors.push(S2CellId::from_ij_same(face, i, j_new, jsame).parent(Some(level)));
if isame || jsame {
neighbors.push(
S2CellId::from_ij_same(face, i_new, j_new, isame && jsame).parent(Some(level)),
);
}
neighbors
}
pub fn low_bits(&self) -> u32 {
self.id as u32
}
pub fn high_bits(&self) -> u32 {
(self.id >> 32) as u32
}
pub fn is_out_of_bounds_wm(&self) -> bool {
self.face() != 0
}
}
impl fmt::Display for S2CellId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
if !self.is_valid() {
return f.write_str("Invalid");
}
write!(f, "{}/", self.face())?;
for level in 1..=self.level() {
write!(f, "{}", self.child_position(level))?;
}
Ok(())
}
}
impl From<S2CellId> for u64 {
fn from(val: S2CellId) -> Self {
val.id
}
}
impl From<&u64> for S2CellId {
fn from(value: &u64) -> Self {
S2CellId::new(*value)
}
}
impl From<u64> for S2CellId {
fn from(value: u64) -> Self {
S2CellId::new(value)
}
}
impl<M: Clone + Default> From<&LonLat<M>> for S2CellId {
fn from(value: &LonLat<M>) -> Self {
S2CellId::from_lon_lat(value)
}
}
impl From<&S2Point> for S2CellId {
fn from(value: &S2Point) -> Self {
S2CellId::from_s2_point(value)
}
}
impl From<String> for S2CellId {
fn from(s: String) -> Self {
S2CellId::from_string(&s)
}
}
impl From<&str> for S2CellId {
fn from(s: &str) -> Self {
S2CellId::from_string(s)
}
}
impl<M: Clone + Default> From<&VectorPoint<M>> for S2CellId {
fn from(value: &VectorPoint<M>) -> Self {
let point: S2Point = value.into();
(&point).into()
}
}
pub fn lsb_for_level(level: u8) -> u64 {
1_u64 << (2 * (K_MAX_LEVEL - (level as u64)))
}
pub fn size_st(level: u8) -> f64 {
ij_to_st(size_ij(level))
}
pub fn size_ij(level: u8) -> u32 {
1 << (K_MAX_CELL_LEVEL - level)
}
pub fn ij_level_to_bound_uv(i: u32, j: u32, level: u8) -> BBox {
let cell_size = size_ij(level) as i32;
let i_low = i as i32 & -cell_size;
let j_low = j as i32 & -cell_size;
let min_i = i_low;
let max_i = i_low + cell_size;
let min_j = j_low;
let max_j = j_low + cell_size;
BBox::new(
ST_TO_UV(ij_to_st(min_i as u32)),
ST_TO_UV(ij_to_st(min_j as u32)),
ST_TO_UV(ij_to_st(max_i as u32)),
ST_TO_UV(ij_to_st(max_j as u32)),
)
}