geographdb_core/storage/

spatial_page.rs

//! Octree-Native Spatial Page Storage
//!
//! Groups NodeRec records into spatially-clustered pages ordered by Morton code.
//! Each page is a self-contained unit: header + sorted nodes + local properties + edge cluster.
//!
//! # Binary Layout
//!
//! ```text
//! [PageHeader: 64 bytes]
//! [NodeRec[]: node_count * 72 bytes, Morton-sorted]
//! [pad to 8-byte alignment]
//! [PropertyBlock: length-prefixed key-value pairs]
//! [pad to 8-byte alignment]
//! [EdgeCluster: PageLocalEdge[]]
//! ```
//!
//! # Design Decisions
//!
//! - **Page size**: 4KB default, tunable. Fits in one OS page for zero-copy mmap.
//! - **NodeRec copy**: Pages duplicate NodeRec from `nodes.dat` for spatial locality.
//!   `nodes.dat` remains the append-only MVCC source of truth; pages are a read-optimized
//!   spatial index rebuilt on demand.
//! - **Properties**: Stored locally per page to avoid separate file lookups.
//! - **Cross-page edges**: Store `dst_node_id` as global u64; reader follows pointer.

use crate::storage::data_structures::NodeRec;
use anyhow::{bail, Result};
use bytemuck::{bytes_of, from_bytes, Pod, Zeroable};
use std::collections::HashMap;

pub const SPATIAL_PAGE_MAGIC: u32 = 0x5347_5041; // "SGPA"
pub const SPATIAL_PAGE_VERSION: u32 = 1;

/// Default max nodes per page. At 72 bytes/NodeRec, 56 nodes = 4032 bytes + header ≈ 4KB.
pub const DEFAULT_MAX_NODES_PER_PAGE: usize = 56;

/// Bounding box for 3D spatial queries.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct BoundingBox {
    pub min_x: f32,
    pub max_x: f32,
    pub min_y: f32,
    pub max_y: f32,
    pub min_z: f32,
    pub max_z: f32,
}

impl BoundingBox {
    pub fn new(min_x: f32, max_x: f32, min_y: f32, max_y: f32, min_z: f32, max_z: f32) -> Self {
        Self {
            min_x,
            max_x,
            min_y,
            max_y,
            min_z,
            max_z,
        }
    }

    pub fn from_node(node: &NodeRec) -> Self {
        Self {
            min_x: node.x,
            max_x: node.x,
            min_y: node.y,
            max_y: node.y,
            min_z: node.z,
            max_z: node.z,
        }
    }

    pub fn expand(&mut self, other: &BoundingBox) {
        self.min_x = self.min_x.min(other.min_x);
        self.max_x = self.max_x.max(other.max_x);
        self.min_y = self.min_y.min(other.min_y);
        self.max_y = self.max_y.max(other.max_y);
        self.min_z = self.min_z.min(other.min_z);
        self.max_z = self.max_z.max(other.max_z);
    }

    pub fn intersects(&self, other: &BoundingBox) -> bool {
        self.min_x <= other.max_x
            && self.max_x >= other.min_x
            && self.min_y <= other.max_y
            && self.max_y >= other.min_y
            && self.min_z <= other.max_z
            && self.max_z >= other.min_z
    }

    pub fn contains_point(&self, x: f32, y: f32, z: f32) -> bool {
        x >= self.min_x
            && x <= self.max_x
            && y >= self.min_y
            && y <= self.max_y
            && z >= self.min_z
            && z <= self.max_z
    }
}

/// Fixed-size page header (88 bytes, no padding).
#[repr(C)]
#[derive(Clone, Copy, Debug, Pod, Zeroable)]
pub struct SpatialPageHeader {
    pub magic: u32,
    pub version: u32,
    pub morton_prefix: u64,
    pub node_count: u16,  // offset 16 (2-aligned)
    pub morton_shift: u8, // offset 18 (1-aligned)
    pub _pad1: [u8; 5],
    /// Bounding box of all nodes in this page.
    pub min_x: f32,
    pub max_x: f32,
    pub min_y: f32,
    pub max_y: f32,
    pub min_z: f32,
    pub max_z: f32,
    pub _pad2: [u8; 8],
    /// Offset within page body (after header) to property block, in bytes.
    pub property_block_offset: u32,
    /// Length of property block in bytes.
    pub property_block_len: u32,
    /// Offset within page body to edge cluster, in bytes.
    pub edge_cluster_offset: u32,
    /// Number of edges in this page's edge cluster.
    pub edge_cluster_len: u32,
    pub _reserved: [u8; 16],
}

impl SpatialPageHeader {
    pub const LEN: usize = std::mem::size_of::<Self>();

    pub fn validate(&self) -> Result<()> {
        if self.magic != SPATIAL_PAGE_MAGIC {
            bail!(
                "SpatialPageHeader: bad magic (expected 0x{:08x}, got 0x{:08x})",
                SPATIAL_PAGE_MAGIC,
                self.magic
            );
        }
        if self.version != SPATIAL_PAGE_VERSION {
            bail!(
                "SpatialPageHeader: unsupported version {} (expected {})",
                self.version,
                SPATIAL_PAGE_VERSION
            );
        }
        Ok(())
    }

    pub fn bbox(&self) -> BoundingBox {
        BoundingBox::new(
            self.min_x, self.max_x, self.min_y, self.max_y, self.min_z, self.max_z,
        )
    }
}

/// Edge record stored locally within a page.
/// src_node_idx is page-local; dst_node_id is global.
#[repr(C)]
#[derive(Clone, Copy, Debug, Pod, Zeroable)]
pub struct PageLocalEdge {
    pub src_node_idx: u16,
    pub _pad1: [u8; 6],
    pub dst_node_id: u64,
    pub weight: f32,
    pub flags: u32,
}

/// A spatial page: header + nodes + properties + edges, packed into a contiguous byte slice.
#[derive(Debug, Clone)]
pub struct SpatialPage {
    pub header: SpatialPageHeader,
    pub nodes: Vec<NodeRec>,
    /// node_id -> properties for nodes in this page.
    pub properties: HashMap<u64, HashMap<String, String>>,
    /// Edges where source node is in this page.
    pub edges: Vec<PageLocalEdge>,
}

impl SpatialPage {
    /// Create an empty page with given Morton prefix and shift.
    pub fn new(morton_prefix: u64, morton_shift: u8) -> Self {
        Self {
            header: SpatialPageHeader {
                magic: SPATIAL_PAGE_MAGIC,
                version: SPATIAL_PAGE_VERSION,
                morton_prefix,
                morton_shift,
                node_count: 0,
                _pad1: [0; 5],
                min_x: f32::MAX,
                max_x: f32::MIN,
                min_y: f32::MAX,
                max_y: f32::MIN,
                min_z: f32::MAX,
                max_z: f32::MIN,
                _pad2: [0; 8],
                property_block_offset: 0,
                property_block_len: 0,
                edge_cluster_offset: 0,
                edge_cluster_len: 0,
                _reserved: [0; 16],
            },
            nodes: vec![],
            properties: HashMap::new(),
            edges: vec![],
        }
    }

    /// Pack this page into a byte vector. Zero-allocation layout: one vec, direct slices.
    pub fn pack(&self) -> Result<Vec<u8>> {
        // Compute total size.
        let node_bytes = self.nodes.len() * std::mem::size_of::<NodeRec>();
        let node_bytes_aligned = align_up(node_bytes, 8);

        // Encode properties into a temp buffer.
        let mut prop_buf = Vec::new();
        for (node_id, props) in &self.properties {
            prop_buf.extend_from_slice(&node_id.to_le_bytes());
            let entry_bytes = encode_property_entry(props)?;
            prop_buf.extend_from_slice(&(entry_bytes.len() as u32).to_le_bytes());
            prop_buf.extend_from_slice(&entry_bytes);
        }
        let prop_bytes_aligned = align_up(prop_buf.len(), 8);

        let edge_bytes = self.edges.len() * std::mem::size_of::<PageLocalEdge>();

        let total = SpatialPageHeader::LEN + node_bytes_aligned + prop_bytes_aligned + edge_bytes;
        let mut buf = Vec::with_capacity(total);

        // Reserve space for header (will mutate offsets before writing).
        let mut header = self.header;
        header.node_count = self.nodes.len() as u16;
        header.edge_cluster_len = self.edges.len() as u32;

        // Layout: header | nodes | pad | properties | pad | edges
        let mut offset = SpatialPageHeader::LEN;

        // After nodes.
        offset += node_bytes_aligned;

        header.property_block_offset = offset as u32;
        header.property_block_len = prop_buf.len() as u32;
        offset += prop_bytes_aligned;

        header.edge_cluster_offset = offset as u32;

        buf.extend_from_slice(bytes_of(&header));

        // Write nodes.
        for node in &self.nodes {
            buf.extend_from_slice(bytes_of(node));
        }
        pad_to_alignment(&mut buf, 8);

        // Write properties.
        buf.extend_from_slice(&prop_buf);
        pad_to_alignment(&mut buf, 8);

        // Write edges.
        for edge in &self.edges {
            buf.extend_from_slice(bytes_of(edge));
        }

        Ok(buf)
    }

    /// Unpack a single spatial page from a byte slice. Zero-copy: nodes and edges are slices.
    pub fn unpack(data: &[u8]) -> Result<Self> {
        if data.len() < SpatialPageHeader::LEN {
            bail!(
                "SpatialPage unpack: data too short for header ({} < {})",
                data.len(),
                SpatialPageHeader::LEN
            );
        }
        let header: &SpatialPageHeader = from_bytes(&data[0..SpatialPageHeader::LEN]);
        header.validate()?;

        let node_count = header.node_count as usize;
        let edge_count = header.edge_cluster_len as usize;

        // Nodes: immediately after header.
        let node_start = SpatialPageHeader::LEN;
        let node_bytes = node_count * std::mem::size_of::<NodeRec>();
        let node_end = node_start + node_bytes;
        if data.len() < node_end {
            bail!("SpatialPage unpack: data too short for nodes");
        }
        let nodes: &[NodeRec] = bytemuck::cast_slice(&data[node_start..node_end]);
        let nodes = nodes.to_vec(); // We own the page data.

        // Properties: after aligned node block.
        let prop_offset = header.property_block_offset as usize;
        let prop_len = header.property_block_len as usize;
        let mut properties = HashMap::new();
        if prop_len > 0 {
            let prop_end = prop_offset + prop_len;
            if data.len() < prop_end {
                bail!("SpatialPage unpack: data too short for properties");
            }
            let mut cursor = prop_offset;
            while cursor < prop_end {
                if cursor + 12 > data.len() {
                    bail!("SpatialPage unpack: truncated property entry");
                }
                let node_id = u64::from_le_bytes([
                    data[cursor],
                    data[cursor + 1],
                    data[cursor + 2],
                    data[cursor + 3],
                    data[cursor + 4],
                    data[cursor + 5],
                    data[cursor + 6],
                    data[cursor + 7],
                ]);
                cursor += 8;
                let entry_len = u32::from_le_bytes([
                    data[cursor],
                    data[cursor + 1],
                    data[cursor + 2],
                    data[cursor + 3],
                ]) as usize;
                cursor += 4;
                if cursor + entry_len > data.len() {
                    bail!("SpatialPage unpack: truncated property data");
                }
                let entry = decode_property_entry(&data[cursor..cursor + entry_len])?;
                cursor += entry_len;
                properties.insert(node_id, entry);
            }
        }

        // Edges: after property block.
        let edge_offset = header.edge_cluster_offset as usize;
        let edge_bytes = edge_count * std::mem::size_of::<PageLocalEdge>();
        let edge_end = edge_offset + edge_bytes;
        if data.len() < edge_end {
            bail!("SpatialPage unpack: data too short for edges");
        }
        let edges: &[PageLocalEdge] = bytemuck::cast_slice(&data[edge_offset..edge_end]);
        let edges = edges.to_vec();

        Ok(SpatialPage {
            header: *header,
            nodes,
            properties,
            edges,
        })
    }

    /// Find a node by its global ID using binary search (nodes are Morton-sorted).
    /// O(log n) within the page.
    pub fn find_node_by_id(&self, id: u64) -> Option<usize> {
        self.nodes.binary_search_by_key(&id, |n| n.id).ok()
    }

    /// Get properties for a node in this page.
    pub fn get_properties(&self, node_id: u64) -> Option<&HashMap<String, String>> {
        self.properties.get(&node_id)
    }

    /// Get edges where source is a given node in this page.
    pub fn get_edges_for_node(&self, node_id: u64) -> Vec<&PageLocalEdge> {
        if let Some(idx) = self.find_node_by_id(node_id) {
            let idx = idx as u16;
            self.edges
                .iter()
                .filter(|e| e.src_node_idx == idx)
                .collect()
        } else {
            vec![]
        }
    }
}

// ─────────────────────────────────────────────────────────────────────────────
// Builder: Group NodeRecs into spatial pages by Morton code
// ─────────────────────────────────────────────────────────────────────────────

/// Build spatial pages from a collection of nodes, grouping by Morton prefix.
pub fn build_spatial_pages(
    mut nodes: Vec<NodeRec>,
    properties: &HashMap<u64, HashMap<String, String>>,
    edges: &[(u64, u64, f32, u32)], // (src_id, dst_id, weight, flags)
    max_nodes_per_page: usize,
) -> Vec<SpatialPage> {
    if nodes.is_empty() {
        return vec![];
    }

    nodes.sort_by_key(|n| n.morton_code);

    let mut pages = vec![];
    let mut start = 0;

    while start < nodes.len() {
        let base_morton = nodes[start].morton_code;
        let mut shift: u8 = 64;

        let (end, final_shift) = loop {
            let prefix = if shift == 0 || shift >= 64 {
                0u64
            } else {
                (base_morton >> shift) << shift
            };

            let mut end = start;
            while end < nodes.len() {
                let candidate = if shift == 0 || shift >= 64 {
                    0u64
                } else {
                    (nodes[end].morton_code >> shift) << shift
                };
                if candidate != prefix || end - start >= max_nodes_per_page {
                    break;
                }
                end += 1;
            }

            if end - start <= max_nodes_per_page || shift == 0 {
                break (end, shift);
            }
            shift = shift.saturating_sub(1);
        };

        let page_nodes = nodes[start..end].to_vec();
        let bbox = compute_bbox(&page_nodes);

        // Edges where src is in this page.
        let mut page_edges = vec![];
        for &(src_id, dst_id, weight, flags) in edges {
            if let Some(src_idx) = page_nodes.iter().position(|n| n.id == src_id) {
                page_edges.push(PageLocalEdge {
                    src_node_idx: src_idx as u16,
                    _pad1: [0; 6],
                    dst_node_id: dst_id,
                    weight,
                    flags,
                });
            }
        }

        // Properties for nodes in this page.
        let mut page_props = HashMap::new();
        for node in &page_nodes {
            if let Some(props) = properties.get(&node.id) {
                page_props.insert(node.id, props.clone());
            }
        }

        let prefix = if final_shift == 0 || final_shift >= 64 {
            0u64
        } else {
            (base_morton >> final_shift) << final_shift
        };

        pages.push(SpatialPage {
            header: SpatialPageHeader {
                magic: SPATIAL_PAGE_MAGIC,
                version: SPATIAL_PAGE_VERSION,
                morton_prefix: prefix,
                morton_shift: final_shift,
                node_count: page_nodes.len() as u16,
                _pad1: [0; 5],
                min_x: bbox.min_x,
                max_x: bbox.max_x,
                min_y: bbox.min_y,
                max_y: bbox.max_y,
                min_z: bbox.min_z,
                max_z: bbox.max_z,
                _pad2: [0; 8],
                property_block_offset: 0,
                property_block_len: 0,
                edge_cluster_offset: 0,
                edge_cluster_len: page_edges.len() as u32,
                _reserved: [0; 16],
            },
            nodes: page_nodes,
            properties: page_props,
            edges: page_edges,
        });

        start = end;
    }

    pages
}

/// Compute bounding box for a slice of nodes.
fn compute_bbox(nodes: &[NodeRec]) -> BoundingBox {
    let mut bbox = BoundingBox::from_node(&nodes[0]);
    for node in &nodes[1..] {
        let nbb = BoundingBox::from_node(node);
        bbox.expand(&nbb);
    }
    bbox
}

// ─────────────────────────────────────────────────────────────────────────────
// Property encoding (compact binary, no JSON, no serde)
// ─────────────────────────────────────────────────────────────────────────────

/// Count how many dimensions have positive overlap between two bounding boxes.
pub fn overlaps(a: &BoundingBox, b: &BoundingBox) -> usize {
    (if a.max_x > b.min_x && a.min_x < b.max_x {
        1
    } else {
        0
    } + if a.max_y > b.min_y && a.min_y < b.max_y {
        1
    } else {
        0
    } + if a.max_z > b.min_z && a.min_z < b.max_z {
        1
    } else {
        0
    })
}

fn encode_property_entry(props: &HashMap<String, String>) -> Result<Vec<u8>> {
    let count = props.len();
    let mut size = 4; // count: u32
    for (k, v) in props {
        if k.len() > u16::MAX as usize || v.len() > u16::MAX as usize {
            bail!("property key/value exceeds 64KB");
        }
        size += 2 + k.len() + 2 + v.len();
    }
    let mut buf = Vec::with_capacity(size);
    buf.extend_from_slice(&(count as u32).to_le_bytes());
    for (k, v) in props {
        buf.extend_from_slice(&(k.len() as u16).to_le_bytes());
        buf.extend_from_slice(k.as_bytes());
        buf.extend_from_slice(&(v.len() as u16).to_le_bytes());
        buf.extend_from_slice(v.as_bytes());
    }
    Ok(buf)
}

fn decode_property_entry(data: &[u8]) -> Result<HashMap<String, String>> {
    if data.len() < 4 {
        bail!("property entry too short");
    }
    let count = u32::from_le_bytes([data[0], data[1], data[2], data[3]]) as usize;
    let mut props = HashMap::with_capacity(count);
    let mut cursor = 4usize;

    for _ in 0..count {
        if cursor + 2 > data.len() {
            bail!("truncated property key length");
        }
        let klen = u16::from_le_bytes([data[cursor], data[cursor + 1]]) as usize;
        cursor += 2;
        if cursor + klen > data.len() {
            bail!("truncated property key");
        }
        let key = String::from_utf8(data[cursor..cursor + klen].to_vec())?;
        cursor += klen;

        if cursor + 2 > data.len() {
            bail!("truncated property value length");
        }
        let vlen = u16::from_le_bytes([data[cursor], data[cursor + 1]]) as usize;
        cursor += 2;
        if cursor + vlen > data.len() {
            bail!("truncated property value");
        }
        let val = String::from_utf8(data[cursor..cursor + vlen].to_vec())?;
        cursor += vlen;

        props.insert(key, val);
    }

    Ok(props)
}

// ─────────────────────────────────────────────────────────────────────────────
// Helpers
// ─────────────────────────────────────────────────────────────────────────────

#[inline]
fn align_up(n: usize, align: usize) -> usize {
    (n + align - 1) & !(align - 1)
}

fn pad_to_alignment(buf: &mut Vec<u8>, align: usize) {
    let aligned = align_up(buf.len(), align);
    buf.resize(aligned, 0);
}

// ─────────────────────────────────────────────────────────────────────────────
// Tests
// ─────────────────────────────────────────────────────────────────────────────

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

    fn make_node(id: u64, morton: u64, x: f32, y: f32, z: f32) -> NodeRec {
        NodeRec {
            id,
            morton_code: morton,
            x,
            y,
            z,
            edge_off: 0,
            edge_len: 0,
            flags: 0,
            begin_ts: 1,
            end_ts: 0,
            tx_id: 1,
            visibility: 1,
            _padding: [0; 7],
        }
    }

    #[test]
    fn test_pack_unpack_empty() {
        let page = SpatialPage::new(0, 64);
        let packed = page.pack().unwrap();
        let unpacked = SpatialPage::unpack(&packed).unwrap();
        assert_eq!(unpacked.header.magic, SPATIAL_PAGE_MAGIC);
        assert_eq!(unpacked.nodes.len(), 0);
        assert_eq!(unpacked.edges.len(), 0);
    }

    #[test]
    fn test_pack_unpack_with_nodes() {
        let mut page = SpatialPage::new(0, 64);
        page.nodes.push(make_node(1, 0b0000, 0.0, 0.0, 0.0));
        page.nodes.push(make_node(2, 0b0001, 1.0, 1.0, 1.0));

        let mut props = HashMap::new();
        props.insert(1u64, {
            let mut p = HashMap::new();
            p.insert("kind".to_string(), "sensor".to_string());
            p
        });

        page.properties = props;

        page.edges.push(PageLocalEdge {
            src_node_idx: 0,
            _pad1: [0; 6],
            dst_node_id: 2,
            weight: 1.5,
            flags: 0,
        });

        let packed = page.pack().unwrap();
        let unpacked = SpatialPage::unpack(&packed).unwrap();

        assert_eq!(unpacked.nodes.len(), 2);
        assert_eq!(unpacked.nodes[0].id, 1);
        assert_eq!(unpacked.nodes[1].id, 2);
        assert_eq!(
            unpacked.properties.get(&1).unwrap().get("kind").unwrap(),
            "sensor"
        );
        assert_eq!(unpacked.edges.len(), 1);
        assert_eq!(unpacked.edges[0].dst_node_id, 2);
        assert_eq!(unpacked.edges[0].weight, 1.5);
    }

    #[test]
    fn test_build_spatial_pages_grouping() {
        let nodes = vec![
            make_node(1, 0b0000_0000, 0.0, 0.0, 0.0),
            make_node(2, 0b0000_0001, 1.0, 0.0, 0.0),
            make_node(3, 0b1111_0000, 10.0, 10.0, 10.0),
            make_node(4, 0b1111_0001, 11.0, 10.0, 10.0),
        ];

        let pages = build_spatial_pages(nodes, &HashMap::new(), &[], 2);
        assert_eq!(pages.len(), 2, "Should split into 2 pages by morton prefix");
        assert_eq!(pages[0].nodes.len(), 2);
        assert_eq!(pages[1].nodes.len(), 2);
    }

    #[test]
    fn test_bounding_box_intersects() {
        let a = BoundingBox::new(0.0, 10.0, 0.0, 10.0, 0.0, 10.0);
        let b = BoundingBox::new(5.0, 15.0, 5.0, 15.0, 5.0, 15.0);
        assert!(a.intersects(&b));

        let c = BoundingBox::new(20.0, 30.0, 0.0, 10.0, 0.0, 10.0);
        assert!(!a.intersects(&c));
    }

    #[test]
    fn test_binary_search_in_page() {
        let mut page = SpatialPage::new(0, 64);
        page.nodes.push(make_node(10, 0b0010, 0.0, 0.0, 0.0));
        page.nodes.push(make_node(20, 0b0100, 1.0, 1.0, 1.0));
        page.nodes.push(make_node(30, 0b1000, 2.0, 2.0, 2.0));

        assert_eq!(page.find_node_by_id(20), Some(1));
        assert_eq!(page.find_node_by_id(99), None);
    }
}