open-standards/map/map-rs/crates/map-index/src/cell_id.rs

// S2CellId — 64-bit hierarchical cell identifier.
//
// Bit layout: 3 bits face | 60 bits position | 1-bit sentinel
// The sentinel bit is at position 2*(30-level) from the LSB. All bits below
// the sentinel are zero. This encoding allows the level to be determined from
// the cell ID itself.

use crate::lat_lng::S2LatLng;
use crate::point::S2Point;
use crate::projection::{
    face_st_to_uv, face_uv_to_xyz, st_to_ij, uv_to_st, xyz_to_face_uv, MAX_LEVEL,
};
use serde::{Deserialize, Serialize};

const MAX_SIZE: u32 = 1u32 << MAX_LEVEL; // 2^30

// Hilbert curve orientation masks.
const SWAP_MASK: u32 = 0x01;
const INVERT_MASK: u32 = 0x02;

/// A 64-bit S2 cell identifier.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
pub struct S2CellId(pub u64);

impl S2CellId {
    /// The sentinel cell (invalid).
    pub const SENTINEL: S2CellId = S2CellId(0);

    /// Create a cell ID from raw bits.
    #[inline]
    pub const fn from_raw(id: u64) -> Self {
        Self(id)
    }

    /// Get the raw 64-bit representation.
    #[inline]
    pub const fn raw(self) -> u64 {
        self.0
    }

    /// Check if this cell ID is valid.
    #[inline]
    pub fn is_valid(self) -> bool {
        let face = self.face();
        if face >= 6 {
            return false;
        }
        // A valid S2 cell has exactly one sentinel bit set in the lower 61 bits.
        // The sentinel bit is at position 2*(30-level), and all bits below it are 0.
        let lsb = self.0 & self.0.wrapping_neg(); // lowest set bit
        // The sentinel bit must be at an even position (0, 2, 4, ..., 60) within the pos bits.
        // Also, there must be at least one set bit in the lower 61 bits.
        let pos = self.pos();
        pos != 0 && (lsb & pos) == lsb
    }

    /// The face (0-5) this cell belongs to.
    #[inline]
    pub fn face(self) -> u8 {
        (self.0 >> 61) as u8
    }

    /// The position within the face (the 61 lower bits).
    #[inline]
    pub fn pos(self) -> u64 {
        self.0 & ((1u64 << 61) - 1)
    }

    /// The level of this cell (0-30).
    #[inline]
    pub fn level(self) -> u8 {
        if self.is_leaf() {
            return MAX_LEVEL;
        }
        // The sentinel bit is at position 2*(30-level).
        // trailing_zeros gives us the sentinel position directly.
        let tz = self.0.trailing_zeros();
        (MAX_LEVEL as u32 - tz / 2) as u8
    }

    /// Whether this is a leaf cell (level 30).
    #[inline]
    pub fn is_leaf(self) -> bool {
        (self.0 & 1) != 0
    }

    /// Whether this is a face cell (level 0).
    #[inline]
    pub fn is_face(self) -> bool {
        let level = self.level();
        level == 0
    }

    /// Construct a cell ID from face, IJ coordinates, and level.
    pub fn from_face_ij_level(face: u8, i: u32, j: u32, level: u8) -> S2CellId {
        // Interleave IJ bits via Hilbert curve, processing one level at a time.
        let mut pos: u64 = 0;
        let mut orientation: u32 = 0;

        // Initial orientation from face.
        // In S2, face 0,2,4 have orientation 0; face 1,3,5 have SWAP_MASK.
        if face & 1 != 0 {
            orientation = SWAP_MASK;
        }

        // Process each level from 0 to 29 (MSB to LSB).
        for lev in 0..MAX_LEVEL as u32 {
            let shift = MAX_LEVEL as u32 - 1 - lev;
            let bi = (i >> shift) & 1;
            let bj = (j >> shift) & 1;

            // Combine into a 2-bit value.
            let mut bits_val = (bi << 1) | bj;

            // Apply orientation transform.
            if orientation & SWAP_MASK != 0 {
                bits_val = ((bits_val >> 1) | ((bits_val & 1) << 1)) & 3;
            }
            if orientation & INVERT_MASK != 0 {
                bits_val ^= 3;
            }

            // Map (i_bit, j_bit) after orientation to Hilbert position.
            let hilbert_pos = IJ_TO_HILBERT[bits_val as usize];

            pos = (pos << 2) | hilbert_pos as u64;

            // Update orientation for next level.
            orientation ^= HILBERT_ORIENTATION_DELTA[hilbert_pos as usize];
        }

        // Construct the cell ID.
        let id = ((face as u64) << 61) | (pos << 1) | 1; // leaf cell

        let cell = S2CellId(id);
        cell.parent_at_level(level)
    }

    /// Construct a leaf cell from a point.
    pub fn from_point(point: &S2Point) -> S2CellId {
        let (face, u, v) = xyz_to_face_uv(point);
        let s = uv_to_st(u);
        let t = uv_to_st(v);
        let i = st_to_ij(s, MAX_LEVEL);
        let j = st_to_ij(t, MAX_LEVEL);
        S2CellId::from_face_ij_level(face, i, j, MAX_LEVEL)
    }

    /// Convenience: construct from lat/lng in degrees.
    pub fn from_lat_lng(lat_deg: f64, lng_deg: f64) -> S2CellId {
        let ll = S2LatLng::from_degrees(lat_deg, lng_deg);
        S2CellId::from_point(&ll.to_point())
    }

    /// The parent cell at the given level. Panics if level > self.level().
    #[inline]
    pub fn parent_at_level(self, level: u8) -> S2CellId {
        assert!(level <= self.level(), "level must be <= current level");
        let new_lsb = self.lsb_for_level(level);
        // Clear bits below the sentinel position, keep bits at and above, then set sentinel.
        S2CellId((self.0 & new_lsb.wrapping_neg()) | new_lsb)
    }

    /// The immediate parent.
    #[inline]
    pub fn parent(self) -> S2CellId {
        assert!(!self.is_face(), "face cells have no parent");
        self.parent_at_level(self.level() - 1)
    }

    /// One of the four children (position 0-3).
    #[inline]
    pub fn child(self, position: u8) -> S2CellId {
        assert!(!self.is_leaf(), "leaf cells have no children");
        assert!(position < 4, "position must be 0-3");
        let old_lsb = self.lsb();
        let new_lsb = old_lsb >> 2;
        // The child ID has the parent's bits with 2 new position bits.
        S2CellId(self.0 - old_lsb + (2 * (position as u64) + 1) * new_lsb)
    }

    /// All four children.
    #[inline]
    pub fn children(self) -> [S2CellId; 4] {
        [
            self.child(0),
            self.child(1),
            self.child(2),
            self.child(3),
        ]
    }

    /// The lowest leaf cell contained by this cell.
    #[inline]
    pub fn range_min(self) -> S2CellId {
        S2CellId(self.0 - self.lsb() + 1)
    }

    /// The highest leaf cell contained by this cell.
    #[inline]
    pub fn range_max(self) -> S2CellId {
        S2CellId(self.0 + self.lsb() - 1)
    }

    /// Whether this cell contains another cell (hierarchically).
    #[inline]
    pub fn contains(self, other: &S2CellId) -> bool {
        other.0 >= self.range_min().0 && other.0 <= self.range_max().0
    }

    /// Get the center of this cell as an S2Point.
    pub fn to_point(self) -> S2Point {
        let (face, si, ti) = self.get_center_si_ti();
        let s = si_ti_to_st(si);
        let t = si_ti_to_st(ti);
        let (u, v) = face_st_to_uv(s, t);
        face_uv_to_xyz(face, u, v).normalize()
    }

    /// Get the center of this cell as (lat_deg, lng_deg).
    pub fn to_lat_lng(self) -> (f64, f64) {
        let p = self.to_point();
        let ll = S2LatLng::from_point(p);
        ll.to_degrees()
    }

    /// The 4 edge-adjacent neighbor cells at the same level.
    pub fn neighbors(self) -> [S2CellId; 4] {
        let level = self.level();
        let (face, i, j) = self.to_face_ij();
        let size = size_ij(level) as i64;

        [
            Self::from_face_ij_wrap(face, i as i64 - size, j as i64, level),
            Self::from_face_ij_wrap(face, i as i64 + size, j as i64, level),
            Self::from_face_ij_wrap(face, i as i64, j as i64 - size, level),
            Self::from_face_ij_wrap(face, i as i64, j as i64 + size, level),
        ]
    }

    // ── Internal helpers ──

    /// The lowest set bit (the sentinel).
    #[inline]
    fn lsb(self) -> u64 {
        self.0 & self.0.wrapping_neg()
    }

    /// The sentinel bit position for a given level.
    #[inline]
    fn lsb_for_level(self, level: u8) -> u64 {
        1u64 << (2 * (MAX_LEVEL as u32 - level as u32))
    }

    /// Decompose back to face and IJ coordinates.
    fn to_face_ij(self) -> (u8, u32, u32) {
        let face = self.face();

        // Extract position bits (between face bits and sentinel).
        let pos_bits = (self.0 >> 1) & ((1u64 << (2 * MAX_LEVEL as u32)) - 1);

        // Reverse the Hilbert curve to recover I, J.
        let mut i: u32 = 0;
        let mut j: u32 = 0;
        let mut orientation: u32 = if face & 1 != 0 { SWAP_MASK } else { 0 };

        for lev in 0..MAX_LEVEL as u32 {
            let shift = 2 * (MAX_LEVEL as u32 - 1 - lev);
            let hilbert_pos = ((pos_bits >> shift) & 3) as u32;

            // Reverse orientation to recover original bits.
            let mut bits_val = HILBERT_TO_IJ[hilbert_pos as usize];

            if orientation & INVERT_MASK != 0 {
                bits_val ^= 3;
            }
            if orientation & SWAP_MASK != 0 {
                bits_val = ((bits_val >> 1) | ((bits_val & 1) << 1)) & 3;
            }

            let bi = (bits_val >> 1) & 1;
            let bj = bits_val & 1;

            let target_shift = MAX_LEVEL as u32 - 1 - lev;
            i |= bi << target_shift;
            j |= bj << target_shift;

            // Update orientation for next level.
            orientation ^= HILBERT_ORIENTATION_DELTA[hilbert_pos as usize];
        }

        (face, i, j)
    }

    /// Get center SI/TI coordinates. SI/TI are in [0, 2*MAX_SIZE].
    fn get_center_si_ti(self) -> (u8, u32, u32) {
        let (face, i, j) = self.to_face_ij();
        let level = self.level();

        if self.is_leaf() {
            // Leaf cell: center = 2*i + 1 (in SI/TI space which is [0, 2*MAX_SIZE])
            return (face, 2 * i + 1, 2 * j + 1);
        }

        // Non-leaf: center of the IJ range at this level.
        let size = size_ij(level);
        let si = 2 * (i & (!(size - 1))) + size;
        let ti = 2 * (j & (!(size - 1))) + size;
        (face, si, ti)
    }

    /// Construct cell ID from face and IJ with wrapping across faces.
    fn from_face_ij_wrap(face: u8, i: i64, j: i64, level: u8) -> S2CellId {
        let max = MAX_SIZE as i64;

        // If both coordinates are within range, use the same face.
        if i >= 0 && i < max && j >= 0 && j < max {
            return S2CellId::from_face_ij_level(face, i as u32, j as u32, level);
        }

        // Convert IJ to ST (allowing out-of-range [0,1] values).
        let scale = 1.0 / max as f64;
        let s = (i as f64 + 0.5) * scale;
        let t = (j as f64 + 0.5) * scale;
        // Convert to UV (quadratic projection, works outside [0,1]).
        let u = crate::projection::st_to_uv(s);
        let v = crate::projection::st_to_uv(t);
        // Project to XYZ (point will be off the face if coordinates are out-of-range).
        let p = face_uv_to_xyz(face, u, v).normalize();
        S2CellId::from_point(&p).parent_at_level(level)
    }
}

/// Convert SI/TI to ST. SI/TI is in [0, 2*MAX_SIZE], ST is in [0, 1].
#[inline]
fn si_ti_to_st(si: u32) -> f64 {
    si as f64 / (2 * MAX_SIZE) as f64
}

/// Number of IJ units spanned by a cell at the given level.
#[inline]
fn size_ij(level: u8) -> u32 {
    1u32 << (MAX_LEVEL as u32 - level as u32)
}

// ── Hilbert curve lookup tables ──
// These are the standard 2x2 Hilbert curve tables.

/// Map (i_bit, j_bit) after orientation transform -> Hilbert position [0..3].
/// Index is (i_bit << 1 | j_bit).
const IJ_TO_HILBERT: [u32; 4] = [0, 1, 3, 2];

/// Reverse map: Hilbert position -> (i_bit << 1 | j_bit).
const HILBERT_TO_IJ: [u32; 4] = [0, 1, 3, 2];

/// Orientation delta after visiting a child at each Hilbert position.
const HILBERT_ORIENTATION_DELTA: [u32; 4] = [
    SWAP_MASK,                     // pos 0: swap
    0,                             // pos 1: no change
    0,                             // pos 2: no change
    SWAP_MASK | INVERT_MASK,       // pos 3: swap + invert
];

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

    #[test]
    fn face_cells_valid() {
        for face in 0..6u8 {
            let cell = S2CellId::from_face_ij_level(face, 0, 0, 0);
            assert!(cell.is_valid(), "face {face} cell should be valid");
            assert_eq!(cell.face(), face);
            assert_eq!(cell.level(), 0);
            assert!(cell.is_face());
        }
    }

    #[test]
    fn leaf_cell_is_leaf() {
        let cell = S2CellId::from_face_ij_level(0, 0, 0, MAX_LEVEL);
        assert!(cell.is_leaf());
        assert_eq!(cell.level(), MAX_LEVEL);
    }

    #[test]
    fn parent_child_consistency() {
        let leaf = S2CellId::from_face_ij_level(0, 100, 200, MAX_LEVEL);
        let parent = leaf.parent();
        assert_eq!(parent.level(), MAX_LEVEL - 1);
        assert!(parent.contains(&leaf));

        // One of the children should contain the original leaf.
        let children = parent.children();
        let mut found = false;
        for c in &children {
            if c.contains(&leaf) {
                found = true;
                break;
            }
        }
        assert!(found, "parent's children should contain the original leaf");
    }

    #[test]
    fn parent_at_level() {
        let leaf = S2CellId::from_face_ij_level(2, 500, 300, MAX_LEVEL);
        for level in 0..MAX_LEVEL {
            let ancestor = leaf.parent_at_level(level);
            assert_eq!(ancestor.level(), level);
            assert!(ancestor.contains(&leaf));
        }
    }

    #[test]
    fn contains_transitive() {
        let leaf = S2CellId::from_face_ij_level(1, 1000, 2000, MAX_LEVEL);
        let l10 = leaf.parent_at_level(10);
        let l5 = leaf.parent_at_level(5);
        let l0 = leaf.parent_at_level(0);

        assert!(l0.contains(&l5));
        assert!(l5.contains(&l10));
        assert!(l0.contains(&l10));
        assert!(l10.contains(&leaf));
        assert!(l0.contains(&leaf));
    }

    #[test]
    fn range_min_max() {
        let cell = S2CellId::from_face_ij_level(0, 0, 0, 5);
        let rmin = cell.range_min();
        let rmax = cell.range_max();
        assert!(rmin.0 < rmax.0);
        assert!(rmin.is_leaf());
        assert!(rmax.is_leaf());
    }

    #[test]
    fn from_point_to_point_round_trip() {
        // Test several known points.
        let test_points = [
            (0.0_f64, 0.0_f64),   // equator / prime meridian
            (45.0, 90.0),         // mid-latitude
            (-33.8688, 151.2093), // Sydney
            (90.0, 0.0),          // north pole
            (-90.0, 0.0),         // south pole
        ];

        for (lat, lng) in &test_points {
            let cell = S2CellId::from_lat_lng(*lat, *lng);
            let (lat2, lng2) = cell.to_lat_lng();

            // At leaf level, the cell is very small. Center should be very close.
            let dist = S2LatLng::from_degrees(*lat, *lng)
                .get_distance(S2LatLng::from_degrees(lat2, lng2));
            // At level 30, cells are ~1 cm, so distance should be tiny.
            assert!(
                dist < 1e-7,
                "round trip too far for ({lat}, {lng}): dist = {dist}"
            );
        }
    }

    #[test]
    fn lat0_lon0_face0() {
        let cell = S2CellId::from_lat_lng(0.0, 0.0);
        assert_eq!(cell.face(), 0, "lat=0 lon=0 should be on face 0");
    }

    #[test]
    fn neighbors_count() {
        let cell = S2CellId::from_face_ij_level(0, 500, 500, 10);
        let nbrs = cell.neighbors();
        // All neighbors should be at the same level.
        for n in &nbrs {
            assert_eq!(n.level(), 10);
        }
        // Neighbors should be distinct from the cell and from each other.
        for n in &nbrs {
            assert_ne!(n.0, cell.0);
        }
    }

    #[test]
    fn children_are_four() {
        let cell = S2CellId::from_face_ij_level(0, 0, 0, 10);
        let children = cell.children();
        assert_eq!(children.len(), 4);
        for (i, c) in children.iter().enumerate() {
            assert_eq!(c.level(), 11, "child {i} should be at level 11");
            assert!(cell.contains(c), "parent should contain child {i}");
        }
        // Children should be distinct.
        for i in 0..4 {
            for j in (i + 1)..4 {
                assert_ne!(children[i].0, children[j].0);
            }
        }
    }

    #[test]
    fn sentinel_is_invalid() {
        assert!(!S2CellId::SENTINEL.is_valid());
    }

    #[test]
    fn from_face_ij_level_level_preserved() {
        for level in 0..=MAX_LEVEL {
            let cell = S2CellId::from_face_ij_level(0, 0, 0, level);
            assert_eq!(cell.level(), level, "level mismatch at level {level}");
        }
    }
}