leyline_engine/
segment.rs1use 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
18pub struct Segment {
20 fwd: Vec<Edge>,
22 rev: Vec<Edge>,
24 keys: Vec<(String, NodeId)>,
26}
27
28impl Segment {
29 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 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}