leyline_engine/

graph.rs

use std::collections::{BTreeSet, HashMap};
use std::path::Path;

use crate::ids::{EdgeLabel, NodeId};
use crate::manifest::{Manifest, SegmentMeta};
use crate::memtable::Memtable;
use crate::segment::Segment;
use crate::store::LocalStore;
use crate::{Error, Result};

/// Top-level graph handle: memtable over immutable segments, LSM-style.
///
/// Reads merge the memtable with every open segment; `flush` drains the
/// memtable into a new segment and commits a new manifest generation.
/// `Graph::new` is in-memory only (no flush); `Graph::open` is the embedded
/// persistent mode.
pub struct Graph {
    memtable: Memtable,
    segments: Vec<Segment>,
    store: Option<LocalStore>,
    /// All keys ever, loaded from segments on open. Fits memory until it
    /// doesn't — see docs/SEGMENTS.md, planned evolution.
    key_index: HashMap<String, NodeId>,
    /// Keys created since the last flush; what the next segment's key
    /// section will contain.
    pending_keys: Vec<(String, NodeId)>,
    labels: Vec<String>,
    next_id: u64,
    generation: u64,
    committed_edges: u64,
}

impl Graph {
    /// In-memory only; `flush` returns `Error::NoStore`.
    pub fn new() -> Self {
        Self {
            memtable: Memtable::new(),
            segments: Vec::new(),
            store: None,
            key_index: HashMap::new(),
            pending_keys: Vec::new(),
            labels: Vec::new(),
            next_id: 0,
            generation: 0,
            committed_edges: 0,
        }
    }

    /// Open (or create) a graph rooted at `dir`.
    pub fn open(dir: impl AsRef<Path>) -> Result<Self> {
        let store = LocalStore::open(dir)?;
        let manifest = store.load_manifest()?.unwrap_or_else(Manifest::empty);

        let mut segments = Vec::with_capacity(manifest.segments.len());
        let mut key_index = HashMap::new();
        for meta in &manifest.segments {
            let segment = Segment::decode(&store.read_segment(&meta.name)?)?;
            for (key, node) in segment.keys() {
                key_index.insert(key.to_owned(), node);
            }
            segments.push(segment);
        }

        Ok(Self {
            memtable: Memtable::new(),
            segments,
            store: Some(store),
            key_index,
            pending_keys: Vec::new(),
            labels: manifest.labels,
            next_id: manifest.next_node_id,
            generation: manifest.generation,
            committed_edges: manifest.edge_count,
        })
    }

    /// Create a node with a user-facing key, or return the existing id.
    pub fn upsert_node(&mut self, key: &str) -> NodeId {
        if let Some(&id) = self.key_index.get(key) {
            return id;
        }
        let id = NodeId(self.next_id);
        self.next_id += 1;
        self.key_index.insert(key.to_owned(), id);
        self.pending_keys.push((key.to_owned(), id));
        self.memtable.insert_node(id);
        id
    }

    pub fn node(&self, key: &str) -> Option<NodeId> {
        self.key_index.get(key).copied()
    }

    /// Intern a label string, returning its id. Label indices are identity
    /// forever — the manifest's label list is append-only.
    pub fn label(&mut self, name: &str) -> EdgeLabel {
        if let Some(pos) = self.labels.iter().position(|l| l == name) {
            return EdgeLabel(pos as u32);
        }
        self.labels.push(name.to_owned());
        EdgeLabel((self.labels.len() - 1) as u32)
    }

    pub fn add_edge(&mut self, src: NodeId, label: EdgeLabel, dst: NodeId) -> Result<()> {
        // Edges are a set. The memtable dedups itself; this check keeps a
        // re-added committed edge from inflating counters or segments.
        if self
            .segments
            .iter()
            .any(|s| s.out(src, label).any(|n| n == dst))
        {
            return Ok(());
        }
        self.memtable.insert_edge(src, label, dst)
    }

    /// Outbound neighbors of `src` under `label`, merged across memtable
    /// and segments, ascending and deduplicated.
    pub fn out(&self, src: NodeId, label: EdgeLabel) -> impl Iterator<Item = NodeId> {
        let mut merged: BTreeSet<NodeId> = self.memtable.out(src, label).collect();
        for segment in &self.segments {
            merged.extend(segment.out(src, label));
        }
        merged.into_iter()
    }

    /// Inbound neighbors of `dst` under `label`, merged the same way.
    pub fn r#in(&self, dst: NodeId, label: EdgeLabel) -> impl Iterator<Item = NodeId> {
        let mut merged: BTreeSet<NodeId> = self.memtable.r#in(dst, label).collect();
        for segment in &self.segments {
            merged.extend(segment.r#in(dst, label));
        }
        merged.into_iter()
    }

    /// Drain the memtable into a new segment and commit the next manifest
    /// generation. No-op (returns the current generation) when there is
    /// nothing to write.
    pub fn flush(&mut self) -> Result<u64> {
        let store = self.store.as_ref().ok_or(Error::NoStore)?;
        if self.memtable.is_empty() && self.pending_keys.is_empty() {
            return Ok(self.generation);
        }

        // Segment key sections are sorted by key bytes (docs/SEGMENTS.md).
        self.pending_keys.sort_by(|(a, _), (b, _)| a.cmp(b));
        let bytes = Segment::encode(
            self.memtable.iter_fwd(),
            self.memtable.iter_rev(),
            &self.pending_keys,
        );

        let generation = self.generation + 1;
        let name = format!("seg-{generation:08}.seg");
        // Segment before manifest: a crash here leaves an orphan file for
        // the sweep, never a manifest naming a segment that isn't there.
        store.write_segment(&name, &bytes)?;

        let edges = self.memtable.edge_count() as u64;
        let mut segments_meta: Vec<SegmentMeta> = self
            .manifest_segments()
            .unwrap_or_default();
        segments_meta.push(SegmentMeta {
            name: name.clone(),
            level: 0,
            edges,
            bytes: bytes.len() as u64,
        });

        let manifest = Manifest {
            format_version: crate::manifest::FORMAT_VERSION,
            generation,
            labels: self.labels.clone(),
            next_node_id: self.next_id,
            node_count: self.key_index.len() as u64,
            edge_count: self.committed_edges + edges,
            segments: segments_meta,
        };
        store.commit_manifest(&manifest)?;

        self.segments.push(Segment::decode(&bytes)?);
        self.memtable.clear();
        self.pending_keys.clear();
        self.generation = generation;
        self.committed_edges += edges;
        Ok(generation)
    }

    /// Re-read the committed segment list; the manifest is the truth for
    /// what was durable before this process.
    fn manifest_segments(&self) -> Option<Vec<SegmentMeta>> {
        let store = self.store.as_ref()?;
        store.load_manifest().ok().flatten().map(|m| m.segments)
    }

    pub fn generation(&self) -> u64 {
        self.generation
    }

    pub fn node_count(&self) -> u64 {
        self.key_index.len() as u64
    }

    pub fn edge_count(&self) -> u64 {
        self.committed_edges + self.memtable.edge_count() as u64
    }
}

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

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

    #[test]
    fn upsert_is_idempotent() {
        let mut g = Graph::new();
        let a = g.upsert_node("alice");
        assert_eq!(a, g.upsert_node("alice"));
        assert_ne!(a, g.upsert_node("bob"));
    }

    #[test]
    fn out_and_in_traversal() {
        let mut g = Graph::new();
        let follows = g.label("follows");
        let a = g.upsert_node("alice");
        let b = g.upsert_node("bob");
        let c = g.upsert_node("carol");
        g.add_edge(a, follows, b).unwrap();
        g.add_edge(a, follows, c).unwrap();
        g.add_edge(b, follows, c).unwrap();

        assert_eq!(g.out(a, follows).collect::<Vec<_>>(), vec![b, c]);
        assert_eq!(g.r#in(c, follows).collect::<Vec<_>>(), vec![a, b]);
    }

    #[test]
    fn in_memory_graph_cannot_flush() {
        let mut g = Graph::new();
        assert!(matches!(g.flush(), Err(Error::NoStore)));
    }

    #[test]
    fn flush_and_reopen_roundtrip() {
        let dir = tempfile::tempdir().unwrap();

        let (a, b, follows_idx) = {
            let mut g = Graph::open(dir.path()).unwrap();
            let follows = g.label("follows");
            let a = g.upsert_node("alice");
            let b = g.upsert_node("bob");
            g.add_edge(a, follows, b).unwrap();
            assert_eq!(g.flush().unwrap(), 1);
            (a, b, follows)
        };

        let mut g = Graph::open(dir.path()).unwrap();
        assert_eq!(g.generation(), 1);
        assert_eq!(g.node_count(), 2);
        assert_eq!(g.edge_count(), 1);
        // Same keys resolve to the same ids after reopen.
        assert_eq!(g.node("alice"), Some(a));
        assert_eq!(g.upsert_node("bob"), b);
        // Label interning must agree with what segments were written with.
        assert_eq!(g.label("follows"), follows_idx);
        assert_eq!(g.out(a, follows_idx).collect::<Vec<_>>(), vec![b]);
        assert_eq!(g.r#in(b, follows_idx).collect::<Vec<_>>(), vec![a]);
    }

    #[test]
    fn reads_merge_across_segments_and_memtable() {
        let dir = tempfile::tempdir().unwrap();
        let mut g = Graph::open(dir.path()).unwrap();
        let follows = g.label("follows");
        let a = g.upsert_node("alice");
        let b = g.upsert_node("bob");
        g.add_edge(a, follows, b).unwrap();
        g.flush().unwrap();

        let c = g.upsert_node("carol");
        g.add_edge(a, follows, c).unwrap();
        g.flush().unwrap();

        let d = g.upsert_node("dave");
        g.add_edge(a, follows, d).unwrap(); // memtable only

        // One neighborhood spread over two segments plus the memtable.
        assert_eq!(g.out(a, follows).collect::<Vec<_>>(), vec![b, c, d]);
        assert_eq!(g.generation(), 2);
        assert_eq!(g.edge_count(), 3);
    }

    #[test]
    fn flush_is_noop_when_clean() {
        let dir = tempfile::tempdir().unwrap();
        let mut g = Graph::open(dir.path()).unwrap();
        let follows = g.label("follows");
        let a = g.upsert_node("alice");
        let b = g.upsert_node("bob");
        g.add_edge(a, follows, b).unwrap();
        assert_eq!(g.flush().unwrap(), 1);
        assert_eq!(g.flush().unwrap(), 1); // nothing new to write

        // Re-adding a committed edge is absorbed, not re-counted.
        g.add_edge(a, follows, b).unwrap();
        assert_eq!(g.edge_count(), 1);
        assert_eq!(g.flush().unwrap(), 1);
    }
}