leyline-engine 0.1.0

use std::collections::BTreeSet;

use crate::ids::{EdgeLabel, NodeId};
use crate::Result;

/// In-memory write buffer, drained into an immutable segment at flush.
///
/// Edges live in two BTreeSets sorted exactly the way segments are:
/// `(src, label, dst)` forward and `(dst, label, src)` reverse. Flush is a
/// straight iteration into CSR runs, and a neighborhood is one contiguous
/// range scan — the same access shape the segment reader will have.
pub struct Memtable {
    nodes: BTreeSet<NodeId>,
    fwd: BTreeSet<(NodeId, EdgeLabel, NodeId)>,
    rev: BTreeSet<(NodeId, EdgeLabel, NodeId)>,
}

impl Memtable {
    pub fn new() -> Self {
        Self {
            nodes: BTreeSet::new(),
            fwd: BTreeSet::new(),
            rev: BTreeSet::new(),
        }
    }

    pub fn insert_node(&mut self, id: NodeId) {
        self.nodes.insert(id);
    }

    pub fn contains_node(&self, id: NodeId) -> bool {
        self.nodes.contains(&id)
    }

    pub fn insert_edge(&mut self, src: NodeId, label: EdgeLabel, dst: NodeId) -> Result<()> {
        self.fwd.insert((src, label, dst));
        self.rev.insert((dst, label, src));
        Ok(())
    }

    /// Outbound neighbors of `src` under `label`, ascending.
    pub fn out(&self, src: NodeId, label: EdgeLabel) -> impl Iterator<Item = NodeId> + '_ {
        self.fwd
            .range((src, label, NodeId(u64::MIN))..=(src, label, NodeId(u64::MAX)))
            .map(|&(_, _, dst)| dst)
    }

    /// Inbound neighbors of `dst` under `label`, ascending.
    pub fn r#in(&self, dst: NodeId, label: EdgeLabel) -> impl Iterator<Item = NodeId> + '_ {
        self.rev
            .range((dst, label, NodeId(u64::MIN))..=(dst, label, NodeId(u64::MAX)))
            .map(|&(_, _, src)| src)
    }

    pub fn edge_count(&self) -> usize {
        self.fwd.len()
    }

    pub fn is_empty(&self) -> bool {
        self.nodes.is_empty() && self.fwd.is_empty()
    }

    /// Flush order: already (src, label, dst)-sorted, one pass to encode.
    pub fn iter_fwd(&self) -> impl ExactSizeIterator<Item = (NodeId, EdgeLabel, NodeId)> + '_ {
        self.fwd.iter().copied()
    }

    pub fn iter_rev(&self) -> impl ExactSizeIterator<Item = (NodeId, EdgeLabel, NodeId)> + '_ {
        self.rev.iter().copied()
    }

    pub fn clear(&mut self) {
        self.nodes.clear();
        self.fwd.clear();
        self.rev.clear();
    }
}

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