leyline-engine 0.1.0

//! Segment encode/decode — the immutable unit of storage.
//!
//! Format spec: docs/SEGMENTS.md. Everything is sorted the way it is read:
//! forward by (src, label, dst), reverse by (dst, label, src), keys by
//! bytes. The v1 reader decodes the whole file into memory — that decoded
//! form *is* the cache tier.

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

const MAGIC: &[u8; 8] = b"LEYSEG1\0";
const VERSION: u32 = 1;
const HEADER_LEN: usize = 8 + 4 + 8 + 8 + 8;
const EDGE_RECORD_LEN: usize = 8 + 4 + 8;

type Edge = (NodeId, EdgeLabel, NodeId);

/// A decoded segment. Immutable once built; lookups are binary searches.
pub struct Segment {
    /// (src, label, dst), sorted.
    fwd: Vec<Edge>,
    /// (dst, label, src), sorted.
    rev: Vec<Edge>,
    /// (key, node), sorted by key.
    keys: Vec<(String, NodeId)>,
}

impl Segment {
    /// Encode from already-sorted parts. Callers hand over iterators in
    /// segment order (the memtable's natural order), so encoding is one
    /// pass, no sort.
    pub fn encode(
        fwd: impl ExactSizeIterator<Item = Edge>,
        rev: impl ExactSizeIterator<Item = Edge>,
        keys: &[(String, NodeId)],
    ) -> Vec<u8> {
        let mut buf = Vec::with_capacity(
            HEADER_LEN + (fwd.len() + rev.len()) * EDGE_RECORD_LEN + keys.len() * 16,
        );
        buf.extend_from_slice(MAGIC);
        buf.extend_from_slice(&VERSION.to_le_bytes());
        buf.extend_from_slice(&(fwd.len() as u64).to_le_bytes());
        buf.extend_from_slice(&(rev.len() as u64).to_le_bytes());
        buf.extend_from_slice(&(keys.len() as u64).to_le_bytes());

        for (a, label, b) in fwd.chain(rev) {
            buf.extend_from_slice(&a.0.to_le_bytes());
            buf.extend_from_slice(&label.0.to_le_bytes());
            buf.extend_from_slice(&b.0.to_le_bytes());
        }
        for (key, node) in keys {
            buf.extend_from_slice(&node.0.to_le_bytes());
            buf.extend_from_slice(&(key.len() as u16).to_le_bytes());
            buf.extend_from_slice(key.as_bytes());
        }

        let crc = crc32fast::hash(&buf);
        buf.extend_from_slice(&crc.to_le_bytes());
        buf
    }

    pub fn decode(bytes: &[u8]) -> Result<Self> {
        let corrupt = |why: &str| Error::CorruptSegment(why.to_owned());

        if bytes.len() < HEADER_LEN + 4 {
            return Err(corrupt("shorter than header + crc"));
        }
        let (body, crc_bytes) = bytes.split_at(bytes.len() - 4);
        let stored = u32::from_le_bytes(crc_bytes.try_into().unwrap());
        if crc32fast::hash(body) != stored {
            return Err(corrupt("crc mismatch"));
        }
        if &body[..8] != MAGIC {
            return Err(corrupt("bad magic"));
        }
        let version = u32::from_le_bytes(body[8..12].try_into().unwrap());
        if version != VERSION {
            return Err(Error::CorruptSegment(format!("unknown version {version}")));
        }

        let read_u64 = |at: usize| u64::from_le_bytes(body[at..at + 8].try_into().unwrap());
        let fwd_count = read_u64(12) as usize;
        let rev_count = read_u64(20) as usize;
        let key_count = read_u64(28) as usize;

        let mut at = HEADER_LEN;
        let mut read_edges = |count: usize| -> Result<Vec<Edge>> {
            let mut edges = Vec::with_capacity(count);
            for _ in 0..count {
                if at + EDGE_RECORD_LEN > body.len() {
                    return Err(Error::CorruptSegment("edge section truncated".into()));
                }
                let a = NodeId(u64::from_le_bytes(body[at..at + 8].try_into().unwrap()));
                let label =
                    EdgeLabel(u32::from_le_bytes(body[at + 8..at + 12].try_into().unwrap()));
                let b = NodeId(u64::from_le_bytes(body[at + 12..at + 20].try_into().unwrap()));
                edges.push((a, label, b));
                at += EDGE_RECORD_LEN;
            }
            Ok(edges)
        };
        let fwd = read_edges(fwd_count)?;
        let rev = read_edges(rev_count)?;

        let mut keys = Vec::with_capacity(key_count);
        for _ in 0..key_count {
            if at + 10 > body.len() {
                return Err(corrupt("key section truncated"));
            }
            let node = NodeId(u64::from_le_bytes(body[at..at + 8].try_into().unwrap()));
            let len = u16::from_le_bytes(body[at + 8..at + 10].try_into().unwrap()) as usize;
            at += 10;
            if at + len > body.len() {
                return Err(corrupt("key bytes truncated"));
            }
            let key = std::str::from_utf8(&body[at..at + len])
                .map_err(|_| corrupt("key not utf-8"))?
                .to_owned();
            keys.push((key, node));
            at += len;
        }
        if at != body.len() {
            return Err(corrupt("trailing bytes"));
        }

        Ok(Self { fwd, rev, keys })
    }

    pub fn out(&self, src: NodeId, label: EdgeLabel) -> impl Iterator<Item = NodeId> + '_ {
        Self::scan(&self.fwd, src, label)
    }

    pub fn r#in(&self, dst: NodeId, label: EdgeLabel) -> impl Iterator<Item = NodeId> + '_ {
        Self::scan(&self.rev, dst, label)
    }

    /// Contiguous range for (a, label) in a sorted edge slice.
    fn scan(edges: &[Edge], a: NodeId, label: EdgeLabel) -> impl Iterator<Item = NodeId> + '_ {
        let lo = edges.partition_point(|&(x, l, _)| (x, l) < (a, label));
        let hi = edges.partition_point(|&(x, l, _)| (x, l) <= (a, label));
        edges[lo..hi].iter().map(|&(_, _, b)| b)
    }

    pub fn keys(&self) -> impl Iterator<Item = (&str, NodeId)> + '_ {
        self.keys.iter().map(|(k, n)| (k.as_str(), *n))
    }

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

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

    fn sample() -> Vec<u8> {
        let fwd = vec![
            (NodeId(0), EdgeLabel(0), NodeId(1)),
            (NodeId(0), EdgeLabel(0), NodeId(2)),
            (NodeId(1), EdgeLabel(1), NodeId(2)),
        ];
        let rev: Vec<Edge> = vec![
            (NodeId(1), EdgeLabel(0), NodeId(0)),
            (NodeId(2), EdgeLabel(0), NodeId(0)),
            (NodeId(2), EdgeLabel(1), NodeId(1)),
        ];
        let keys = vec![
            ("alice".to_owned(), NodeId(0)),
            ("bob".to_owned(), NodeId(1)),
            ("carol".to_owned(), NodeId(2)),
        ];
        Segment::encode(fwd.into_iter(), rev.into_iter(), &keys)
    }

    #[test]
    fn roundtrip() {
        let seg = Segment::decode(&sample()).unwrap();
        assert_eq!(
            seg.out(NodeId(0), EdgeLabel(0)).collect::<Vec<_>>(),
            vec![NodeId(1), NodeId(2)]
        );
        assert_eq!(
            seg.r#in(NodeId(2), EdgeLabel(1)).collect::<Vec<_>>(),
            vec![NodeId(1)]
        );
        assert_eq!(seg.out(NodeId(0), EdgeLabel(1)).count(), 0);
        assert_eq!(seg.keys().count(), 3);
    }

    #[test]
    fn corruption_detected() {
        let mut bytes = sample();
        let mid = bytes.len() / 2;
        bytes[mid] ^= 0xff;
        assert!(matches!(
            Segment::decode(&bytes),
            Err(Error::CorruptSegment(_))
        ));
    }

    #[test]
    fn truncation_detected() {
        let bytes = sample();
        assert!(Segment::decode(&bytes[..bytes.len() - 1]).is_err());
    }
}