nodedb_lite/engine/spatial/

index.rs

//! Per-collection spatial index manager for Lite.
//!
//! Wraps `nodedb_spatial::RTree` with:
//! - Incremental insert/delete on document put/delete
//! - Geometry extraction from document fields
//! - Checkpoint/restore via MessagePack + CRC32C (same pattern as HNSW/CSR)
//! - Spatial query execution (range search, nearest neighbor)

use std::collections::HashMap;

type CheckpointData = (
    Vec<(String, String, Vec<u8>)>,
    HashMap<(String, String), u64>,
    u64,
);

use nodedb_mem::ScopedMemory;
use nodedb_spatial::rtree::{RTree, RTreeEntry};
use nodedb_types::BoundingBox;
use nodedb_types::geometry::Geometry;

/// Manages per-collection R-tree spatial indexes.
///
/// Each collection that has geometry fields gets its own R-tree.
/// The field name is stored alongside so we know which document field
/// to extract geometry from.
pub struct SpatialIndexManager {
    /// (collection_name, field_name) → R-tree.
    indices: HashMap<(String, String), RTree>,
    /// Document ID → entry ID mapping for deletion.
    /// Key: (collection, doc_id), Value: entry_id in R-tree.
    doc_to_entry: HashMap<(String, String), u64>,
    /// Inverse map: entry_id → (collection, doc_id) for scan result resolution.
    entry_to_doc: HashMap<u64, (String, String)>,
    /// Next entry ID (monotonically increasing).
    next_id: u64,
    /// Governor handle bound to the spatial engine budget. Cloned into every
    /// R-tree this manager creates.
    memory: ScopedMemory,
}

impl SpatialIndexManager {
    pub fn new(memory: ScopedMemory) -> Self {
        Self {
            indices: HashMap::new(),
            doc_to_entry: HashMap::new(),
            entry_to_doc: HashMap::new(),
            next_id: 1,
            memory,
        }
    }

    /// Resolve an R-tree entry ID to its document ID within a collection.
    pub fn doc_id_for_entry(&self, entry_id: u64) -> Option<&str> {
        self.entry_to_doc
            .get(&entry_id)
            .map(|(_, doc_id)| doc_id.as_str())
    }

    /// Index a geometry from a document. If the document already has an entry,
    /// it is removed first (upsert semantics).
    pub fn index_document(
        &mut self,
        collection: &str,
        field: &str,
        doc_id: &str,
        geometry: &Geometry,
    ) {
        let key = (collection.to_string(), field.to_string());
        let doc_key = (collection.to_string(), doc_id.to_string());

        // Remove old entry if this document was previously indexed.
        if let Some(old_id) = self.doc_to_entry.remove(&doc_key) {
            self.entry_to_doc.remove(&old_id);
            if let Some(tree) = self.indices.get_mut(&key) {
                tree.delete(old_id);
            }
        }

        let bbox = nodedb_types::geometry_bbox(geometry);
        let entry_id = self.next_id;
        self.next_id += 1;

        let memory = self.memory.clone();
        let tree = self
            .indices
            .entry(key)
            .or_insert_with(|| RTree::new(memory));
        tree.insert(RTreeEntry { id: entry_id, bbox });
        self.doc_to_entry.insert(doc_key, entry_id);
        self.entry_to_doc
            .insert(entry_id, (collection.to_string(), doc_id.to_string()));
    }

    /// Remove a document's geometry from the index.
    pub fn remove_document(&mut self, collection: &str, field: &str, doc_id: &str) {
        let key = (collection.to_string(), field.to_string());
        let doc_key = (collection.to_string(), doc_id.to_string());

        if let Some(entry_id) = self.doc_to_entry.remove(&doc_key) {
            self.entry_to_doc.remove(&entry_id);
            if let Some(tree) = self.indices.get_mut(&key) {
                tree.delete(entry_id);
            }
        }
    }

    /// Remove every entry of every R-tree `collection` owns. Each tree stays
    /// registered under its `(collection, field)` key, so the next checkpoint
    /// writes it out empty instead of leaving a stale blob behind. Returns
    /// the `(field, doc_id)` of every entry removed.
    pub fn truncate_collection(&mut self, collection: &str) -> Vec<(String, String)> {
        let fields: Vec<String> = self
            .indices
            .keys()
            .filter(|(coll, _)| coll == collection)
            .map(|(_, field)| field.clone())
            .collect();
        let mut removed = Vec::new();
        for field in fields {
            let key = (collection.to_string(), field.clone());
            let Some(tree) = self.indices.get_mut(&key) else {
                continue;
            };
            for entry in tree.entries() {
                if let Some((_, doc_id)) = self.entry_to_doc.get(&entry.id) {
                    removed.push((field.clone(), doc_id.clone()));
                }
            }
            *tree = RTree::new(self.memory.clone());
        }
        self.doc_to_entry.retain(|(coll, _), entry_id| {
            let owned = coll == collection;
            if owned {
                self.entry_to_doc.remove(entry_id);
            }
            !owned
        });
        removed
    }

    /// Range search: find all document entry IDs whose bbox intersects the query.
    pub fn search(&self, collection: &str, field: &str, query: &BoundingBox) -> Vec<&RTreeEntry> {
        let key = (collection.to_string(), field.to_string());
        match self.indices.get(&key) {
            Some(tree) => tree.search(query),
            None => Vec::new(),
        }
    }

    /// Nearest-neighbor search.
    pub fn nearest(
        &self,
        collection: &str,
        field: &str,
        lng: f64,
        lat: f64,
        k: usize,
    ) -> Vec<nodedb_spatial::rtree::NnResult> {
        let key = (collection.to_string(), field.to_string());
        match self.indices.get(&key) {
            Some(tree) => tree.nearest(lng, lat, k),
            None => Vec::new(),
        }
    }

    /// Number of indexed entries across all collections.
    /// Whether no spatial indices exist.
    pub fn is_empty(&self) -> bool {
        self.indices.is_empty()
    }

    pub fn total_entries(&self) -> usize {
        self.indices.values().map(|t| t.len()).sum()
    }

    /// Number of indexed collections.
    pub fn collection_count(&self) -> usize {
        self.indices.len()
    }

    /// Checkpoint all R-trees to bytes for persistence.
    ///
    /// Returns a vec of `(collection, field, rtree_bytes)`.
    pub fn checkpoint_all(&self) -> Vec<(String, String, Vec<u8>)> {
        let mut results = Vec::new();
        for ((collection, field), tree) in &self.indices {
            match tree.checkpoint_to_bytes(None) {
                Ok(bytes) => results.push((collection.clone(), field.clone(), bytes)),
                Err(e) => {
                    tracing::error!(
                        collection = %collection,
                        field = %field,
                        error = %e,
                        "spatial index checkpoint failed"
                    );
                }
            }
        }
        results
    }

    /// Return the data needed by the checkpoint module to persist full state.
    ///
    /// Returns `(rtree_checkpoints, doc_to_entry, next_id)`.
    pub fn checkpoint_data(&self) -> CheckpointData {
        (
            self.checkpoint_all(),
            self.doc_to_entry.clone(),
            self.next_id,
        )
    }

    /// Load a fully-restored checkpoint.
    ///
    /// Replaces the current in-memory state with the provided R-tree
    /// checkpoints and the exact `doc_id → entry_id` mapping that was
    /// serialised at flush time.
    pub fn load_checkpoint(
        &mut self,
        checkpoints: &[(String, String, Vec<u8>)],
        doc_to_entry: HashMap<(String, String), u64>,
        next_id: u64,
    ) {
        // Rebuild inverse map from the restored forward map.
        self.entry_to_doc = doc_to_entry
            .iter()
            .map(|((col, doc_id), &eid)| (eid, (col.clone(), doc_id.clone())))
            .collect();
        self.doc_to_entry = doc_to_entry;
        self.next_id = next_id;
        for (collection, field, bytes) in checkpoints {
            match RTree::from_checkpoint(bytes, None, self.memory.clone()) {
                Ok(tree) => {
                    self.indices
                        .insert((collection.clone(), field.clone()), tree);
                }
                Err(e) => {
                    tracing::warn!(
                        collection = %collection,
                        field = %field,
                        error = %e,
                        "spatial R-tree restore failed; collection will be empty until rebuilt"
                    );
                }
            }
        }
    }

    /// Restore R-trees from raw checkpoint bytes only (no doc_to_entry).
    ///
    /// Used as a fallback when no docmap is available (e.g. legacy checkpoints
    /// written before this field was introduced). After restoration, upserts
    /// and deletes of already-indexed docs may not evict stale entries; a full
    /// rebuild from documents is the reliable recovery path in that case.
    pub fn restore_all(checkpoints: &[(String, String, Vec<u8>)], memory: ScopedMemory) -> Self {
        let mut manager = Self::new(memory);
        for (collection, field, bytes) in checkpoints {
            match RTree::from_checkpoint(bytes, None, manager.memory.clone()) {
                Ok(tree) => {
                    let max_id = tree.entries().iter().map(|e| e.id).max().unwrap_or(0);
                    if max_id >= manager.next_id {
                        manager.next_id = max_id + 1;
                    }
                    manager
                        .indices
                        .insert((collection.clone(), field.clone()), tree);
                }
                Err(e) => {
                    tracing::warn!(
                        collection = %collection,
                        field = %field,
                        error = %e,
                        "spatial index restore failed, will rebuild from documents"
                    );
                }
            }
        }
        manager
    }

    /// Rebuild spatial index from a collection of documents.
    ///
    /// Scans all documents, extracts geometry from the specified field,
    /// and builds the R-tree.
    pub fn rebuild_from_documents(
        &mut self,
        collection: &str,
        field: &str,
        documents: &[(String, Geometry)],
    ) {
        let entries: Vec<RTreeEntry> = documents
            .iter()
            .map(|(doc_id, geom)| {
                let id = self.next_id;
                self.next_id += 1;
                let doc_key = (collection.to_string(), doc_id.clone());
                self.doc_to_entry.insert(doc_key, id);
                self.entry_to_doc
                    .insert(id, (collection.to_string(), doc_id.clone()));
                RTreeEntry {
                    id,
                    bbox: nodedb_types::geometry_bbox(geom),
                }
            })
            .collect();

        let tree = RTree::bulk_load(entries, self.memory.clone());
        self.indices
            .insert((collection.to_string(), field.to_string()), tree);
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use nodedb_mem::{EngineId, EngineLimits, GovernorConfig, MemoryGovernor};
    use nodedb_types::{DatabaseId, TenantId};

    use super::*;

    /// Build a real, uncapped governor scoped to the spatial engine for tests.
    fn test_memory() -> ScopedMemory {
        let per_engine = usize::MAX / EngineId::ALL.len();
        let governor = Arc::new(
            MemoryGovernor::new(GovernorConfig {
                global_ceiling: per_engine * EngineId::ALL.len(),
                engine_limits: EngineLimits::uniform(per_engine),
            })
            .expect("test governor"),
        );
        ScopedMemory::new(
            governor,
            DatabaseId::DEFAULT,
            TenantId::new(0),
            EngineId::Spatial,
        )
    }

    #[test]
    fn index_and_search() {
        let mut mgr = SpatialIndexManager::new(test_memory());
        mgr.index_document("places", "location", "doc1", &Geometry::point(10.0, 20.0));
        mgr.index_document("places", "location", "doc2", &Geometry::point(11.0, 21.0));
        mgr.index_document("places", "location", "doc3", &Geometry::point(50.0, 50.0));

        let results = mgr.search(
            "places",
            "location",
            &BoundingBox::new(9.0, 19.0, 12.0, 22.0),
        );
        assert_eq!(results.len(), 2);
    }

    #[test]
    fn truncate_collection_empties_only_that_collection() {
        let mut mgr = SpatialIndexManager::new(test_memory());
        mgr.index_document("places", "loc", "doc1", &Geometry::point(10.0, 20.0));
        mgr.index_document("places", "loc", "doc2", &Geometry::point(11.0, 21.0));
        mgr.index_document("other", "loc", "doc9", &Geometry::point(10.0, 20.0));

        let mut removed = mgr.truncate_collection("places");
        removed.sort();
        assert_eq!(
            removed,
            vec![
                ("loc".to_string(), "doc1".to_string()),
                ("loc".to_string(), "doc2".to_string())
            ]
        );
        let bbox = BoundingBox::new(9.0, 19.0, 12.0, 22.0);
        assert!(mgr.search("places", "loc", &bbox).is_empty());
        assert_eq!(mgr.search("other", "loc", &bbox).len(), 1);
        assert_eq!(
            mgr.collection_count(),
            2,
            "the emptied tree stays registered"
        );
        assert!(mgr.doc_id_for_entry(1).is_none());

        mgr.index_document("places", "loc", "doc3", &Geometry::point(10.5, 20.5));
        assert_eq!(mgr.search("places", "loc", &bbox).len(), 1);
    }

    #[test]
    fn upsert_replaces_old_entry() {
        let mut mgr = SpatialIndexManager::new(test_memory());
        mgr.index_document("places", "loc", "doc1", &Geometry::point(10.0, 20.0));
        mgr.index_document("places", "loc", "doc1", &Geometry::point(50.0, 50.0));

        // Old location should not be found.
        let old = mgr.search("places", "loc", &BoundingBox::new(9.0, 19.0, 12.0, 22.0));
        assert!(old.is_empty());

        // New location should be found.
        let new = mgr.search("places", "loc", &BoundingBox::new(49.0, 49.0, 51.0, 51.0));
        assert_eq!(new.len(), 1);
    }

    #[test]
    fn remove_document() {
        let mut mgr = SpatialIndexManager::new(test_memory());
        mgr.index_document("places", "loc", "doc1", &Geometry::point(10.0, 20.0));
        mgr.remove_document("places", "loc", "doc1");

        let results = mgr.search("places", "loc", &BoundingBox::new(0.0, 0.0, 180.0, 90.0));
        assert!(results.is_empty());
    }

    #[test]
    fn checkpoint_restore_roundtrip() {
        let mut mgr = SpatialIndexManager::new(test_memory());
        for i in 0..50 {
            mgr.index_document(
                "buildings",
                "geom",
                &format!("b{i}"),
                &Geometry::point(i as f64 * 0.5, i as f64 * 0.3),
            );
        }

        let checkpoints = mgr.checkpoint_all();
        assert_eq!(checkpoints.len(), 1);

        let restored = SpatialIndexManager::restore_all(&checkpoints, test_memory());
        assert_eq!(restored.total_entries(), 50);

        let results = restored.search(
            "buildings",
            "geom",
            &BoundingBox::new(-180.0, -90.0, 180.0, 90.0),
        );
        assert_eq!(results.len(), 50);
    }

    #[test]
    fn nearest_neighbor() {
        let mut mgr = SpatialIndexManager::new(test_memory());
        mgr.index_document("pois", "loc", "a", &Geometry::point(0.0, 0.0));
        mgr.index_document("pois", "loc", "b", &Geometry::point(10.0, 10.0));
        mgr.index_document("pois", "loc", "c", &Geometry::point(1.0, 1.0));

        let nn = mgr.nearest("pois", "loc", 0.5, 0.5, 2);
        assert_eq!(nn.len(), 2);
    }

    #[test]
    fn rebuild_from_documents() {
        let mut mgr = SpatialIndexManager::new(test_memory());
        let docs: Vec<(String, Geometry)> = (0..100)
            .map(|i| {
                (
                    format!("d{i}"),
                    Geometry::point(i as f64 * 0.1, i as f64 * 0.1),
                )
            })
            .collect();
        mgr.rebuild_from_documents("col", "geom", &docs);
        assert_eq!(mgr.total_entries(), 100);
    }

    #[test]
    fn multiple_collections() {
        let mut mgr = SpatialIndexManager::new(test_memory());
        mgr.index_document("a", "loc", "d1", &Geometry::point(0.0, 0.0));
        mgr.index_document("b", "loc", "d1", &Geometry::point(50.0, 50.0));

        assert_eq!(mgr.collection_count(), 2);

        let a_results = mgr.search("a", "loc", &BoundingBox::new(-1.0, -1.0, 1.0, 1.0));
        assert_eq!(a_results.len(), 1);

        let b_results = mgr.search("b", "loc", &BoundingBox::new(-1.0, -1.0, 1.0, 1.0));
        assert!(b_results.is_empty());
    }
}