Skip to main content

leyline_engine/
segment.rs

1//! Segment encode/decode — the immutable unit of storage.
2//!
3//! Format spec: docs/SEGMENTS.md. Everything is sorted the way it is read:
4//! forward by (src, label, dst), reverse by (dst, label, src), keys by
5//! bytes. The v1 reader decodes the whole file into memory — that decoded
6//! form *is* the cache tier.
7
8use crate::ids::{EdgeLabel, NodeId};
9use crate::{Error, Result};
10
11const MAGIC: &[u8; 8] = b"LEYSEG1\0";
12const VERSION: u32 = 1;
13const HEADER_LEN: usize = 8 + 4 + 8 + 8 + 8;
14const EDGE_RECORD_LEN: usize = 8 + 4 + 8;
15
16type Edge = (NodeId, EdgeLabel, NodeId);
17
18/// A decoded segment. Immutable once built; lookups are binary searches.
19pub struct Segment {
20    /// (src, label, dst), sorted.
21    fwd: Vec<Edge>,
22    /// (dst, label, src), sorted.
23    rev: Vec<Edge>,
24    /// (key, node), sorted by key.
25    keys: Vec<(String, NodeId)>,
26}
27
28impl Segment {
29    /// Encode from already-sorted parts. Callers hand over iterators in
30    /// segment order (the memtable's natural order), so encoding is one
31    /// pass, no sort.
32    pub fn encode(
33        fwd: impl ExactSizeIterator<Item = Edge>,
34        rev: impl ExactSizeIterator<Item = Edge>,
35        keys: &[(String, NodeId)],
36    ) -> Vec<u8> {
37        let mut buf = Vec::with_capacity(
38            HEADER_LEN + (fwd.len() + rev.len()) * EDGE_RECORD_LEN + keys.len() * 16,
39        );
40        buf.extend_from_slice(MAGIC);
41        buf.extend_from_slice(&VERSION.to_le_bytes());
42        buf.extend_from_slice(&(fwd.len() as u64).to_le_bytes());
43        buf.extend_from_slice(&(rev.len() as u64).to_le_bytes());
44        buf.extend_from_slice(&(keys.len() as u64).to_le_bytes());
45
46        for (a, label, b) in fwd.chain(rev) {
47            buf.extend_from_slice(&a.0.to_le_bytes());
48            buf.extend_from_slice(&label.0.to_le_bytes());
49            buf.extend_from_slice(&b.0.to_le_bytes());
50        }
51        for (key, node) in keys {
52            buf.extend_from_slice(&node.0.to_le_bytes());
53            buf.extend_from_slice(&(key.len() as u16).to_le_bytes());
54            buf.extend_from_slice(key.as_bytes());
55        }
56
57        let crc = crc32fast::hash(&buf);
58        buf.extend_from_slice(&crc.to_le_bytes());
59        buf
60    }
61
62    pub fn decode(bytes: &[u8]) -> Result<Self> {
63        let corrupt = |why: &str| Error::CorruptSegment(why.to_owned());
64
65        if bytes.len() < HEADER_LEN + 4 {
66            return Err(corrupt("shorter than header + crc"));
67        }
68        let (body, crc_bytes) = bytes.split_at(bytes.len() - 4);
69        let stored = u32::from_le_bytes(crc_bytes.try_into().unwrap());
70        if crc32fast::hash(body) != stored {
71            return Err(corrupt("crc mismatch"));
72        }
73        if &body[..8] != MAGIC {
74            return Err(corrupt("bad magic"));
75        }
76        let version = u32::from_le_bytes(body[8..12].try_into().unwrap());
77        if version != VERSION {
78            return Err(Error::CorruptSegment(format!("unknown version {version}")));
79        }
80
81        let read_u64 = |at: usize| u64::from_le_bytes(body[at..at + 8].try_into().unwrap());
82        let fwd_count = read_u64(12) as usize;
83        let rev_count = read_u64(20) as usize;
84        let key_count = read_u64(28) as usize;
85
86        let mut at = HEADER_LEN;
87        let mut read_edges = |count: usize| -> Result<Vec<Edge>> {
88            let mut edges = Vec::with_capacity(count);
89            for _ in 0..count {
90                if at + EDGE_RECORD_LEN > body.len() {
91                    return Err(Error::CorruptSegment("edge section truncated".into()));
92                }
93                let a = NodeId(u64::from_le_bytes(body[at..at + 8].try_into().unwrap()));
94                let label =
95                    EdgeLabel(u32::from_le_bytes(body[at + 8..at + 12].try_into().unwrap()));
96                let b = NodeId(u64::from_le_bytes(body[at + 12..at + 20].try_into().unwrap()));
97                edges.push((a, label, b));
98                at += EDGE_RECORD_LEN;
99            }
100            Ok(edges)
101        };
102        let fwd = read_edges(fwd_count)?;
103        let rev = read_edges(rev_count)?;
104
105        let mut keys = Vec::with_capacity(key_count);
106        for _ in 0..key_count {
107            if at + 10 > body.len() {
108                return Err(corrupt("key section truncated"));
109            }
110            let node = NodeId(u64::from_le_bytes(body[at..at + 8].try_into().unwrap()));
111            let len = u16::from_le_bytes(body[at + 8..at + 10].try_into().unwrap()) as usize;
112            at += 10;
113            if at + len > body.len() {
114                return Err(corrupt("key bytes truncated"));
115            }
116            let key = std::str::from_utf8(&body[at..at + len])
117                .map_err(|_| corrupt("key not utf-8"))?
118                .to_owned();
119            keys.push((key, node));
120            at += len;
121        }
122        if at != body.len() {
123            return Err(corrupt("trailing bytes"));
124        }
125
126        Ok(Self { fwd, rev, keys })
127    }
128
129    pub fn out(&self, src: NodeId, label: EdgeLabel) -> impl Iterator<Item = NodeId> + '_ {
130        Self::scan(&self.fwd, src, label)
131    }
132
133    pub fn r#in(&self, dst: NodeId, label: EdgeLabel) -> impl Iterator<Item = NodeId> + '_ {
134        Self::scan(&self.rev, dst, label)
135    }
136
137    /// Contiguous range for (a, label) in a sorted edge slice.
138    fn scan(edges: &[Edge], a: NodeId, label: EdgeLabel) -> impl Iterator<Item = NodeId> + '_ {
139        let lo = edges.partition_point(|&(x, l, _)| (x, l) < (a, label));
140        let hi = edges.partition_point(|&(x, l, _)| (x, l) <= (a, label));
141        edges[lo..hi].iter().map(|&(_, _, b)| b)
142    }
143
144    pub fn keys(&self) -> impl Iterator<Item = (&str, NodeId)> + '_ {
145        self.keys.iter().map(|(k, n)| (k.as_str(), *n))
146    }
147
148    pub fn edge_count(&self) -> usize {
149        self.fwd.len()
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use super::*;
156
157    fn sample() -> Vec<u8> {
158        let fwd = vec![
159            (NodeId(0), EdgeLabel(0), NodeId(1)),
160            (NodeId(0), EdgeLabel(0), NodeId(2)),
161            (NodeId(1), EdgeLabel(1), NodeId(2)),
162        ];
163        let rev: Vec<Edge> = vec![
164            (NodeId(1), EdgeLabel(0), NodeId(0)),
165            (NodeId(2), EdgeLabel(0), NodeId(0)),
166            (NodeId(2), EdgeLabel(1), NodeId(1)),
167        ];
168        let keys = vec![
169            ("alice".to_owned(), NodeId(0)),
170            ("bob".to_owned(), NodeId(1)),
171            ("carol".to_owned(), NodeId(2)),
172        ];
173        Segment::encode(fwd.into_iter(), rev.into_iter(), &keys)
174    }
175
176    #[test]
177    fn roundtrip() {
178        let seg = Segment::decode(&sample()).unwrap();
179        assert_eq!(
180            seg.out(NodeId(0), EdgeLabel(0)).collect::<Vec<_>>(),
181            vec![NodeId(1), NodeId(2)]
182        );
183        assert_eq!(
184            seg.r#in(NodeId(2), EdgeLabel(1)).collect::<Vec<_>>(),
185            vec![NodeId(1)]
186        );
187        assert_eq!(seg.out(NodeId(0), EdgeLabel(1)).count(), 0);
188        assert_eq!(seg.keys().count(), 3);
189    }
190
191    #[test]
192    fn corruption_detected() {
193        let mut bytes = sample();
194        let mid = bytes.len() / 2;
195        bytes[mid] ^= 0xff;
196        assert!(matches!(
197            Segment::decode(&bytes),
198            Err(Error::CorruptSegment(_))
199        ));
200    }
201
202    #[test]
203    fn truncation_detected() {
204        let bytes = sample();
205        assert!(Segment::decode(&bytes[..bytes.len() - 1]).is_err());
206    }
207}