nodedb_graph/

csr.rs

//! Dense integer CSR adjacency index with interned node IDs and labels.
//!
//! Memory layout:
//! - Old-style `Vec<Vec<(String, u32)>>` ≈ 60 bytes/edge (heap String per edge)
//! - Dense CSR: contiguous `Vec<u32>` offsets + targets + `Vec<u16>` labels ≈ 10 bytes/edge
//!
//! Writes accumulate in a mutable buffer. Reads check both dense CSR and buffer.
//! `compact()` merges the buffer into dense arrays (double-buffered swap).

use std::collections::{HashMap, HashSet, hash_map::Entry};

// Re-export shared Direction from nodedb-types.
pub use nodedb_types::graph::Direction;

/// Dense integer CSR adjacency index.
pub struct CsrIndex {
    // ── Node interning ──
    pub(crate) node_to_id: HashMap<String, u32>,
    pub(crate) id_to_node: Vec<String>,

    // ── Label interning ──
    pub(crate) label_to_id: HashMap<String, u16>,
    pub(crate) id_to_label: Vec<String>,

    // ── Dense CSR (read-only between compactions) ──
    pub(crate) out_offsets: Vec<u32>,
    pub(crate) out_targets: Vec<u32>,
    pub(crate) out_labels: Vec<u16>,

    pub(crate) in_offsets: Vec<u32>,
    pub(crate) in_targets: Vec<u32>,
    pub(crate) in_labels: Vec<u16>,

    // ── Mutable write buffer ──
    pub(crate) buffer_out: Vec<Vec<(u16, u32)>>,
    pub(crate) buffer_in: Vec<Vec<(u16, u32)>>,

    /// Edges deleted since last compaction: `(src, label, dst)`.
    pub(crate) deleted_edges: HashSet<(u32, u16, u32)>,
}

impl Default for CsrIndex {
    fn default() -> Self {
        Self::new()
    }
}

impl CsrIndex {
    pub fn new() -> Self {
        Self {
            node_to_id: HashMap::new(),
            id_to_node: Vec::new(),
            label_to_id: HashMap::new(),
            id_to_label: Vec::new(),
            out_offsets: vec![0],
            out_targets: Vec::new(),
            out_labels: Vec::new(),
            in_offsets: vec![0],
            in_targets: Vec::new(),
            in_labels: Vec::new(),
            buffer_out: Vec::new(),
            buffer_in: Vec::new(),
            deleted_edges: HashSet::new(),
        }
    }

    /// Get or create a dense ID for a node.
    pub(crate) fn ensure_node(&mut self, node: &str) -> u32 {
        match self.node_to_id.entry(node.to_string()) {
            Entry::Occupied(e) => *e.get(),
            Entry::Vacant(e) => {
                let id = self.id_to_node.len() as u32;
                e.insert(id);
                self.id_to_node.push(node.to_string());
                self.out_offsets
                    .push(*self.out_offsets.last().unwrap_or(&0));
                self.in_offsets.push(*self.in_offsets.last().unwrap_or(&0));
                self.buffer_out.push(Vec::new());
                self.buffer_in.push(Vec::new());
                id
            }
        }
    }

    /// Get or create a dense ID for a label.
    fn ensure_label(&mut self, label: &str) -> u16 {
        match self.label_to_id.entry(label.to_string()) {
            Entry::Occupied(e) => *e.get(),
            Entry::Vacant(e) => {
                let id = self.id_to_label.len() as u16;
                e.insert(id);
                self.id_to_label.push(label.to_string());
                id
            }
        }
    }

    /// Add an edge (goes into mutable buffer). Deduplicates.
    pub fn add_edge(&mut self, src: &str, label: &str, dst: &str) {
        let src_id = self.ensure_node(src);
        let dst_id = self.ensure_node(dst);
        let label_id = self.ensure_label(label);

        // Deduplicate in buffer.
        if self.buffer_out[src_id as usize]
            .iter()
            .any(|&(l, d)| l == label_id && d == dst_id)
        {
            return;
        }
        // Deduplicate in dense CSR.
        if self.dense_has_edge(src_id, label_id, dst_id) {
            return;
        }

        self.buffer_out[src_id as usize].push((label_id, dst_id));
        self.buffer_in[dst_id as usize].push((label_id, src_id));
        self.deleted_edges.remove(&(src_id, label_id, dst_id));
    }

    /// Remove an edge.
    pub fn remove_edge(&mut self, src: &str, label: &str, dst: &str) {
        let (Some(&src_id), Some(&dst_id)) = (self.node_to_id.get(src), self.node_to_id.get(dst))
        else {
            return;
        };
        let Some(&label_id) = self.label_to_id.get(label) else {
            return;
        };

        self.buffer_out[src_id as usize].retain(|&(l, d)| !(l == label_id && d == dst_id));
        self.buffer_in[dst_id as usize].retain(|&(l, s)| !(l == label_id && s == src_id));

        if self.dense_has_edge(src_id, label_id, dst_id) {
            self.deleted_edges.insert((src_id, label_id, dst_id));
        }
    }

    /// Get immediate neighbors.
    pub fn neighbors(
        &self,
        node: &str,
        label_filter: Option<&str>,
        direction: Direction,
    ) -> Vec<(String, String)> {
        match label_filter {
            Some(l) => self.neighbors_multi(node, &[l], direction),
            None => self.neighbors_multi(node, &[], direction),
        }
    }

    /// Get neighbors with multi-label filter. Empty labels = all edges.
    pub fn neighbors_multi(
        &self,
        node: &str,
        label_filters: &[&str],
        direction: Direction,
    ) -> Vec<(String, String)> {
        let Some(&node_id) = self.node_to_id.get(node) else {
            return Vec::new();
        };
        let label_ids: Vec<u16> = label_filters
            .iter()
            .filter_map(|l| self.label_to_id.get(*l).copied())
            .collect();
        let match_label = |lid: u16| label_ids.is_empty() || label_ids.contains(&lid);

        let mut result = Vec::new();

        if matches!(direction, Direction::Out | Direction::Both) {
            for (lid, dst) in self.iter_out_edges(node_id) {
                if match_label(lid) {
                    result.push((
                        self.id_to_label[lid as usize].clone(),
                        self.id_to_node[dst as usize].clone(),
                    ));
                }
            }
        }
        if matches!(direction, Direction::In | Direction::Both) {
            for (lid, src) in self.iter_in_edges(node_id) {
                if match_label(lid) {
                    result.push((
                        self.id_to_label[lid as usize].clone(),
                        self.id_to_node[src as usize].clone(),
                    ));
                }
            }
        }

        result
    }

    pub fn node_count(&self) -> usize {
        self.id_to_node.len()
    }

    pub fn contains_node(&self, node: &str) -> bool {
        self.node_to_id.contains_key(node)
    }

    /// Total edge count (dense + buffer - deleted).
    pub fn edge_count(&self) -> usize {
        let dense_out = self.out_targets.len();
        let buffer_out: usize = self.buffer_out.iter().map(|b| b.len()).sum();
        let deleted = self.deleted_edges.len();
        dense_out + buffer_out - deleted
    }

    /// Estimated memory usage in bytes.
    pub fn estimated_memory_bytes(&self) -> usize {
        let offsets = (self.out_offsets.len() + self.in_offsets.len()) * 4;
        let targets = (self.out_targets.len() + self.in_targets.len()) * 4;
        let labels = (self.out_labels.len() + self.in_labels.len()) * 2;
        let buffer: usize = self
            .buffer_out
            .iter()
            .chain(self.buffer_in.iter())
            .map(|b| b.len() * 6)
            .sum();
        let interning = self.id_to_node.iter().map(|s| s.len() + 24).sum::<usize>()
            + self.id_to_label.iter().map(|s| s.len() + 24).sum::<usize>();
        offsets + targets + labels + buffer + interning
    }

    fn dense_has_edge(&self, src: u32, label: u16, dst: u32) -> bool {
        for (lid, target) in self.dense_out_edges(src) {
            if lid == label && target == dst {
                return true;
            }
        }
        false
    }

    pub(crate) fn dense_out_edges(&self, node: u32) -> impl Iterator<Item = (u16, u32)> + '_ {
        let idx = node as usize;
        if idx + 1 >= self.out_offsets.len() {
            return Vec::new().into_iter();
        }
        let start = self.out_offsets[idx] as usize;
        let end = self.out_offsets[idx + 1] as usize;
        (start..end)
            .map(move |i| (self.out_labels[i], self.out_targets[i]))
            .collect::<Vec<_>>()
            .into_iter()
    }

    pub(crate) fn dense_in_edges(&self, node: u32) -> impl Iterator<Item = (u16, u32)> + '_ {
        let idx = node as usize;
        if idx + 1 >= self.in_offsets.len() {
            return Vec::new().into_iter();
        }
        let start = self.in_offsets[idx] as usize;
        let end = self.in_offsets[idx + 1] as usize;
        (start..end)
            .map(move |i| (self.in_labels[i], self.in_targets[i]))
            .collect::<Vec<_>>()
            .into_iter()
    }

    pub(crate) fn iter_out_edges(&self, node: u32) -> impl Iterator<Item = (u16, u32)> + '_ {
        let idx = node as usize;
        let dense = self
            .dense_out_edges(node)
            .filter(move |&(lid, dst)| !self.deleted_edges.contains(&(node, lid, dst)));
        let buffer = if idx < self.buffer_out.len() {
            self.buffer_out[idx].to_vec()
        } else {
            Vec::new()
        };
        dense.chain(buffer)
    }

    pub(crate) fn iter_in_edges(&self, node: u32) -> impl Iterator<Item = (u16, u32)> + '_ {
        let idx = node as usize;
        let dense = self
            .dense_in_edges(node)
            .filter(move |&(lid, src)| !self.deleted_edges.contains(&(src, lid, node)));
        let buffer = if idx < self.buffer_in.len() {
            self.buffer_in[idx].to_vec()
        } else {
            Vec::new()
        };
        dense.chain(buffer)
    }

    /// Resolve a node string to its internal ID.
    pub(crate) fn node_id(&self, node: &str) -> Option<u32> {
        self.node_to_id.get(node).copied()
    }

    /// Resolve an internal ID to its node string.
    pub(crate) fn node_name(&self, id: u32) -> &str {
        self.id_to_node
            .get(id as usize)
            .map(|s| s.as_str())
            .unwrap_or("")
    }

    /// Resolve a label ID to its label string.
    pub(crate) fn label_name(&self, id: u16) -> &str {
        self.id_to_label
            .get(id as usize)
            .map(|s| s.as_str())
            .unwrap_or("")
    }

    /// Resolve a label string to its internal ID.
    pub(crate) fn label_id(&self, label: &str) -> Option<u16> {
        self.label_to_id.get(label).copied()
    }
}

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

    fn make_csr() -> CsrIndex {
        let mut csr = CsrIndex::new();
        csr.add_edge("a", "KNOWS", "b");
        csr.add_edge("b", "KNOWS", "c");
        csr.add_edge("c", "KNOWS", "d");
        csr.add_edge("a", "WORKS", "e");
        csr
    }

    #[test]
    fn neighbors_out() {
        let csr = make_csr();
        let n = csr.neighbors("a", None, Direction::Out);
        assert_eq!(n.len(), 2);
        let dsts: Vec<&str> = n.iter().map(|(_, d)| d.as_str()).collect();
        assert!(dsts.contains(&"b"));
        assert!(dsts.contains(&"e"));
    }

    #[test]
    fn neighbors_filtered() {
        let csr = make_csr();
        let n = csr.neighbors("a", Some("KNOWS"), Direction::Out);
        assert_eq!(n.len(), 1);
        assert_eq!(n[0].1, "b");
    }

    #[test]
    fn neighbors_in() {
        let csr = make_csr();
        let n = csr.neighbors("b", None, Direction::In);
        assert_eq!(n.len(), 1);
        assert_eq!(n[0].1, "a");
    }

    #[test]
    fn incremental_remove() {
        let mut csr = make_csr();
        assert_eq!(csr.neighbors("a", Some("KNOWS"), Direction::Out).len(), 1);
        csr.remove_edge("a", "KNOWS", "b");
        assert_eq!(csr.neighbors("a", Some("KNOWS"), Direction::Out).len(), 0);
    }

    #[test]
    fn duplicate_add_is_idempotent() {
        let mut csr = CsrIndex::new();
        csr.add_edge("a", "L", "b");
        csr.add_edge("a", "L", "b");
        assert_eq!(csr.neighbors("a", None, Direction::Out).len(), 1);
    }

    #[test]
    fn compact_merges_buffer_into_dense() {
        let mut csr = CsrIndex::new();
        csr.add_edge("a", "L", "b");
        csr.add_edge("b", "L", "c");
        assert_eq!(csr.neighbors("a", None, Direction::Out).len(), 1);

        csr.compact();
        assert!(csr.buffer_out.iter().all(|b| b.is_empty()));
        assert_eq!(csr.neighbors("a", None, Direction::Out).len(), 1);
        assert_eq!(csr.neighbors("b", None, Direction::Out).len(), 1);
    }

    #[test]
    fn compact_handles_deletes() {
        let mut csr = CsrIndex::new();
        csr.add_edge("a", "L", "b");
        csr.add_edge("a", "L", "c");
        csr.compact();

        csr.remove_edge("a", "L", "b");
        assert_eq!(csr.neighbors("a", None, Direction::Out).len(), 1);

        csr.compact();
        assert_eq!(csr.neighbors("a", None, Direction::Out).len(), 1);
        assert_eq!(csr.neighbors("a", None, Direction::Out)[0].1, "c");
    }

    #[test]
    fn label_interning() {
        let mut csr = CsrIndex::new();
        for i in 0..100 {
            csr.add_edge(&format!("n{i}"), "FOLLOWS", &format!("n{}", i + 1));
        }
        assert_eq!(csr.id_to_label.len(), 1);
    }

    #[test]
    fn edge_count() {
        let csr = make_csr();
        assert_eq!(csr.edge_count(), 4);
    }

    #[test]
    fn checkpoint_roundtrip() {
        let mut csr = make_csr();
        csr.compact();

        let bytes = csr.checkpoint_to_bytes();
        assert!(!bytes.is_empty());

        let restored = CsrIndex::from_checkpoint(&bytes).unwrap();
        assert_eq!(restored.node_count(), csr.node_count());
        assert_eq!(restored.edge_count(), csr.edge_count());

        let n = restored.neighbors("a", Some("KNOWS"), Direction::Out);
        assert_eq!(n.len(), 1);
        assert_eq!(n[0].1, "b");
    }

    #[test]
    fn memory_estimation() {
        let csr = make_csr();
        let mem = csr.estimated_memory_bytes();
        assert!(mem > 0);
    }
}