core_storage/
topology.rs

1use crate::pack::{push_u32, push_u32s, push_u64, read_u32, read_u32s, read_u64};
2use crate::types::Result as StoreResult;
3use serde::{Deserialize, Deserializer, Serialize, Serializer};
4use std::borrow::Cow;
5use std::cmp::Ordering;
6use std::collections::HashMap;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
9pub enum Direction {
10    Out,
11    In,
12}
13
14/// Flush the per-vertex insert buffer into the frozen block once it exceeds this.
15const INSERT_BUFFER: usize = 32;
16
17/// Sorted frozen neighbors plus an unsorted insert buffer. Serializes as `Vec<u32>`
18/// (merged, sorted, unique) so V6 snapshots stay HashMap-of-HashMap-of-Vec.
19#[derive(Debug, Default, Clone)]
20struct AdjList {
21    frozen: Vec<u32>,
22    delta: Vec<u32>,
23}
24
25impl AdjList {
26    fn contains(&self, id: u32) -> bool {
27        self.frozen.binary_search(&id).is_ok() || self.delta.contains(&id)
28    }
29
30    fn push(&mut self, id: u32) {
31        self.delta.push(id);
32        if self.delta.len() > INSERT_BUFFER {
33            self.flush();
34        }
35    }
36
37    fn flush(&mut self) {
38        if self.delta.is_empty() {
39            return;
40        }
41        if self.frozen.is_empty() {
42            self.delta.sort_unstable();
43            self.delta.dedup();
44            std::mem::swap(&mut self.frozen, &mut self.delta);
45            return;
46        }
47        self.frozen = merge_sorted_unique(&self.frozen, &self.delta);
48        self.delta.clear();
49    }
50
51    fn merged(&self) -> Vec<u32> {
52        merge_sorted_unique(&self.frozen, &self.delta)
53    }
54
55    fn remove(&mut self, id: u32) -> bool {
56        if let Ok(pos) = self.frozen.binary_search(&id) {
57            self.frozen.remove(pos);
58            return true;
59        }
60        if let Some(pos) = self.delta.iter().position(|&x| x == id) {
61            self.delta.swap_remove(pos);
62            return true;
63        }
64        false
65    }
66}
67
68impl Serialize for AdjList {
69    fn serialize<S: Serializer>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error> {
70        if self.delta.is_empty() {
71            self.frozen.serialize(serializer)
72        } else {
73            self.merged().serialize(serializer)
74        }
75    }
76}
77
78impl<'de> Deserialize<'de> for AdjList {
79    fn deserialize<D: Deserializer<'de>>(deserializer: D) -> std::result::Result<Self, D::Error> {
80        Ok(Self {
81            frozen: Vec::<u32>::deserialize(deserializer)?,
82            delta: Vec::new(),
83        })
84    }
85}
86
87fn merge_sorted_unique(frozen: &[u32], delta: &[u32]) -> Vec<u32> {
88    let mut extra: Vec<u32> = delta.to_vec();
89    extra.sort_unstable();
90    extra.dedup();
91    if frozen.is_empty() {
92        return extra;
93    }
94    if extra.is_empty() {
95        return frozen.to_vec();
96    }
97    let mut out = Vec::with_capacity(frozen.len() + extra.len());
98    let mut i = 0;
99    let mut j = 0;
100    while i < frozen.len() && j < extra.len() {
101        match frozen[i].cmp(&extra[j]) {
102            Ordering::Less => {
103                out.push(frozen[i]);
104                i += 1;
105            }
106            Ordering::Greater => {
107                out.push(extra[j]);
108                j += 1;
109            }
110            Ordering::Equal => {
111                out.push(frozen[i]);
112                i += 1;
113                j += 1;
114            }
115        }
116    }
117    out.extend_from_slice(&frozen[i..]);
118    out.extend_from_slice(&extra[j..]);
119    out
120}
121
122#[derive(Debug, Default, Clone, Serialize, Deserialize)]
123struct TypedAdjacency {
124    out: HashMap<u32, AdjList>,
125    inn: HashMap<u32, AdjList>,
126}
127
128/// Typed adjacency: per `(etype, dir, vertex)` a frozen sorted block plus an
129/// unsorted insert buffer. `neighbors` borrows the frozen block when the buffer
130/// is empty; otherwise it returns a sorted-unique merge.
131///
132/// On-disk (V6) shape is still `HashMap<u32, {out, inn: HashMap<u32, Vec<u32>>}>`.
133#[derive(Debug, Default, Clone, Serialize, Deserialize)]
134pub struct Topology {
135    by_type: HashMap<u32, TypedAdjacency>,
136    edge_count: u64,
137}
138
139impl Topology {
140    pub fn new() -> Self {
141        Self::default()
142    }
143
144    pub fn add_edge(&mut self, etype: u32, src: u32, dst: u32) -> bool {
145        let adj = self.by_type.entry(etype).or_default();
146        let dsts = adj.out.entry(src).or_default();
147        if dsts.contains(dst) {
148            return false;
149        }
150        dsts.push(dst);
151        let srcs = adj.inn.entry(dst).or_default();
152        assert!(
153            !srcs.contains(src),
154            "invariant: inn must not contain src as neighbor of dst when out lacks dst"
155        );
156        srcs.push(src);
157        self.edge_count += 1;
158        true
159    }
160
161    pub fn neighbors(&self, etype: u32, dir: Direction, v: u32) -> Cow<'_, [u32]> {
162        match self.adj_list(etype, dir, v) {
163            None => Cow::Borrowed(&[]),
164            Some(n) if n.delta.is_empty() => Cow::Borrowed(&n.frozen),
165            Some(n) => Cow::Owned(n.merged()),
166        }
167    }
168
169    pub fn degree(&self, etype: u32, dir: Direction, v: u32) -> usize {
170        self.neighbors(etype, dir, v).as_ref().len()
171    }
172
173    pub fn edge_count(&self) -> u64 {
174        self.edge_count
175    }
176
177    /// Edge-type ids present in the graph, sorted ascending.
178    /// `by_type` stays a `HashMap` (smaller than a BTreeMap migrate; snapshot encoding unchanged);
179    /// this method collect+sorts keys so iteration is deterministic.
180    pub fn etypes(&self) -> impl Iterator<Item = u32> + '_ {
181        let mut ids: Vec<u32> = self.by_type.keys().copied().collect();
182        ids.sort_unstable();
183        ids.into_iter()
184    }
185
186    pub fn remove_edge(&mut self, etype: u32, src: u32, dst: u32) -> bool {
187        let Some(adj) = self.by_type.get_mut(&etype) else {
188            return false;
189        };
190        let Some(dsts) = adj.out.get_mut(&src) else {
191            return false;
192        };
193        if !dsts.remove(dst) {
194            return false;
195        }
196        let srcs = adj
197            .inn
198            .get_mut(&dst)
199            .expect("invariant: inn bucket must exist when out contains dst");
200        assert!(
201            srcs.remove(src),
202            "invariant: inn must contain src when out contained dst"
203        );
204        self.edge_count -= 1;
205        true
206    }
207
208    fn adj_list(&self, etype: u32, dir: Direction, v: u32) -> Option<&AdjList> {
209        self.by_type.get(&etype).and_then(|adj| match dir {
210            Direction::Out => adj.out.get(&v),
211            Direction::In => adj.inn.get(&v),
212        })
213    }
214
215    /// V7 packed CSR: etype count, then per etype (id, out map, in map), then edge_count.
216    /// Each adjacency map is vertex-count + (vertex, length-prefixed frozen neighbor array).
217    /// Deltas are merged into frozen on pack; unpack leaves deltas empty.
218    pub(crate) fn pack(&self, out: &mut Vec<u8>) {
219        let mut etypes: Vec<u32> = self.by_type.keys().copied().collect();
220        etypes.sort_unstable();
221        push_u32(out, etypes.len() as u32);
222        for et in etypes {
223            push_u32(out, et);
224            let adj = &self.by_type[&et];
225            pack_adj_map(out, &adj.out);
226            pack_adj_map(out, &adj.inn);
227        }
228        push_u64(out, self.edge_count);
229    }
230
231    pub(crate) fn unpack(src: &[u8]) -> StoreResult<(Self, usize)> {
232        let mut pos = 0usize;
233        let n_etypes = read_u32(src, &mut pos)? as usize;
234        let mut by_type = HashMap::with_capacity(n_etypes);
235        for _ in 0..n_etypes {
236            let et = read_u32(src, &mut pos)?;
237            let out = unpack_adj_map(src, &mut pos)?;
238            let inn = unpack_adj_map(src, &mut pos)?;
239            by_type.insert(et, TypedAdjacency { out, inn });
240        }
241        let edge_count = read_u64(src, &mut pos)?;
242        Ok((
243            Self {
244                by_type,
245                edge_count,
246            },
247            pos,
248        ))
249    }
250}
251
252fn pack_adj_map(out: &mut Vec<u8>, map: &HashMap<u32, AdjList>) {
253    let mut verts: Vec<u32> = map.keys().copied().collect();
254    verts.sort_unstable();
255    push_u32(out, verts.len() as u32);
256    for v in verts {
257        push_u32(out, v);
258        let list = &map[&v];
259        if list.delta.is_empty() {
260            push_u32s(out, &list.frozen);
261        } else {
262            let merged = list.merged();
263            push_u32s(out, &merged);
264        }
265    }
266}
267
268fn unpack_adj_map(src: &[u8], pos: &mut usize) -> StoreResult<HashMap<u32, AdjList>> {
269    let n = read_u32(src, pos)? as usize;
270    let mut map = HashMap::with_capacity(n);
271    for _ in 0..n {
272        let v = read_u32(src, pos)?;
273        let frozen = read_u32s(src, pos)?;
274        map.insert(
275            v,
276            AdjList {
277                frozen,
278                delta: Vec::new(),
279            },
280        );
281    }
282    Ok(map)
283}
284
285#[cfg(test)]
286mod tests {
287    use super::*;
288
289    #[test]
290    fn edges_are_typed_directed_sorted_deduped() {
291        let mut t = Topology::new();
292        assert!(t.add_edge(0, 5, 9));
293        assert!(t.add_edge(0, 5, 3));
294        assert!(!t.add_edge(0, 5, 9)); // duplicate
295        assert!(t.add_edge(1, 5, 9)); // same pair, different type: distinct edge
296        assert_eq!(t.neighbors(0, Direction::Out, 5).as_ref(), &[3, 9]); // sorted
297        assert_eq!(t.neighbors(0, Direction::In, 9).as_ref(), &[5]);
298        assert_eq!(t.neighbors(0, Direction::Out, 999).as_ref(), &[] as &[u32]);
299        assert_eq!(t.degree(0, Direction::Out, 5), 2);
300        assert_eq!(t.edge_count(), 3);
301    }
302
303    #[test]
304    fn remove_edge_updates_both_sides_and_count() {
305        let mut t = Topology::new();
306        t.add_edge(0, 1, 2);
307        t.add_edge(0, 1, 3);
308        assert!(t.remove_edge(0, 1, 2));
309        assert!(!t.remove_edge(0, 1, 2)); // idempotent-false
310        assert!(!t.remove_edge(9, 1, 2)); // unknown type
311        assert_eq!(t.neighbors(0, Direction::Out, 1).as_ref(), &[3]);
312        assert_eq!(t.neighbors(0, Direction::In, 2).as_ref(), &[] as &[u32]);
313        assert_eq!(t.edge_count(), 1);
314        // re-add after remove works
315        assert!(t.add_edge(0, 1, 2));
316        assert_eq!(t.edge_count(), 2);
317    }
318
319    #[test]
320    fn etypes_empty_multiple_and_sorted() {
321        let empty = Topology::new();
322        assert_eq!(empty.etypes().collect::<Vec<_>>(), Vec::<u32>::new());
323
324        let mut t = Topology::new();
325        t.add_edge(3, 0, 1);
326        t.add_edge(1, 0, 1);
327        t.add_edge(3, 1, 2); // existing type, must not duplicate
328        t.add_edge(2, 0, 2);
329        assert_eq!(t.etypes().collect::<Vec<_>>(), vec![1, 2, 3]);
330    }
331
332    #[test]
333    fn insert_buffer_defers_sort_until_threshold_then_neighbors_sorted_unique() {
334        let mut t = Topology::new();
335        // Reverse inserts: a full sort on every add would keep the scan path borrowed.
336        for dst in (0u32..32).rev() {
337            assert!(t.add_edge(0, 0, dst));
338            let nbrs = t.neighbors(0, Direction::Out, 0);
339            assert!(
340                matches!(nbrs, Cow::Owned(_)),
341                "delta still dirty at {} edges; neighbors must take the owned merge path",
342                32 - dst
343            );
344            assert!(
345                nbrs.windows(2).all(|w| w[0] < w[1]),
346                "dirty merge must be sorted unique, got {nbrs:?}"
347            );
348            assert_eq!(nbrs.len(), (32 - dst) as usize);
349            let n = t.adj_list(0, Direction::Out, 0).unwrap();
350            assert!(
351                n.frozen.is_empty(),
352                "must not flush frozen before threshold"
353            );
354            assert_eq!(n.delta.len(), (32 - dst) as usize);
355        }
356        assert_eq!(t.degree(0, Direction::Out, 0), 32);
357        assert_eq!(t.neighbors(0, Direction::In, 31).as_ref(), &[0]);
358
359        assert!(t.add_edge(0, 0, 32)); // buffer len > 32 → merge into frozen
360        let nbrs = t.neighbors(0, Direction::Out, 0);
361        assert!(
362            matches!(nbrs, Cow::Borrowed(_)),
363            "after threshold flush, scan path must borrow the frozen block"
364        );
365        let expected: Vec<u32> = (0..33).collect();
366        assert_eq!(nbrs.as_ref(), expected.as_slice());
367        let n = t.adj_list(0, Direction::Out, 0).unwrap();
368        assert!(n.delta.is_empty());
369        assert_eq!(n.frozen, expected);
370        assert_eq!(t.edge_count(), 33);
371        assert_eq!(t.neighbors(0, Direction::In, 32).as_ref(), &[0]);
372        assert!(!t.add_edge(0, 0, 7));
373    }
374
375    #[test]
376    fn remove_edge_from_delta_and_from_frozen() {
377        let mut t = Topology::new();
378        for dst in 0u32..10 {
379            assert!(t.add_edge(0, 1, dst));
380        }
381        assert!(matches!(t.neighbors(0, Direction::Out, 1), Cow::Owned(_)));
382        assert!(t.remove_edge(0, 1, 7));
383        assert!(!t.remove_edge(0, 1, 7));
384        assert_eq!(t.neighbors(0, Direction::In, 7).as_ref(), &[] as &[u32]);
385        assert_eq!(
386            t.neighbors(0, Direction::Out, 1).as_ref(),
387            &[0, 1, 2, 3, 4, 5, 6, 8, 9]
388        );
389        assert_eq!(t.edge_count(), 9);
390        assert!(t.add_edge(0, 1, 7));
391        assert_eq!(t.edge_count(), 10);
392
393        let mut t = Topology::new();
394        for dst in 0u32..33 {
395            assert!(t.add_edge(0, 1, dst));
396        }
397        assert!(matches!(
398            t.neighbors(0, Direction::Out, 1),
399            Cow::Borrowed(_)
400        ));
401        assert!(t.remove_edge(0, 1, 0));
402        assert!(t.remove_edge(0, 1, 32));
403        assert_eq!(t.edge_count(), 31);
404        assert_eq!(t.neighbors(0, Direction::In, 0).as_ref(), &[] as &[u32]);
405        assert_eq!(t.neighbors(0, Direction::In, 16).as_ref(), &[1]);
406        let expected: Vec<u32> = (1..32).collect();
407        assert_eq!(
408            t.neighbors(0, Direction::Out, 1).as_ref(),
409            expected.as_slice()
410        );
411        assert!(t.add_edge(0, 1, 0));
412        assert_eq!(t.edge_count(), 32);
413    }
414
415    #[test]
416    fn serde_wire_is_hashmap_of_hashmap_of_vec() {
417        #[derive(Serialize, Deserialize, PartialEq, Debug)]
418        struct WireAdj {
419            out: HashMap<u32, Vec<u32>>,
420            inn: HashMap<u32, Vec<u32>>,
421        }
422        #[derive(Serialize, Deserialize, PartialEq, Debug)]
423        struct Wire {
424            by_type: HashMap<u32, WireAdj>,
425            edge_count: u64,
426        }
427
428        let mut by_type = HashMap::new();
429        by_type.insert(
430            0,
431            WireAdj {
432                out: HashMap::from([(5, vec![3, 9])]),
433                inn: HashMap::from([(3, vec![5]), (9, vec![5])]),
434            },
435        );
436        let wire = Wire {
437            by_type,
438            edge_count: 2,
439        };
440        let encoded = bincode::serialize(&wire).unwrap();
441        let t: Topology = bincode::deserialize(&encoded).unwrap();
442        assert_eq!(t.neighbors(0, Direction::Out, 5).as_ref(), &[3, 9]);
443        assert_eq!(t.neighbors(0, Direction::In, 3).as_ref(), &[5]);
444        assert_eq!(t.neighbors(0, Direction::In, 9).as_ref(), &[5]);
445        assert_eq!(t.edge_count(), 2);
446        assert!(matches!(
447            t.neighbors(0, Direction::Out, 5),
448            Cow::Borrowed(_)
449        ));
450
451        let roundtrip: Wire = bincode::deserialize(&bincode::serialize(&t).unwrap()).unwrap();
452        assert_eq!(roundtrip.edge_count, 2);
453        assert_eq!(roundtrip.by_type[&0].out[&5], vec![3, 9]);
454        assert_eq!(roundtrip.by_type[&0].inn[&3], vec![5]);
455        assert_eq!(roundtrip.by_type[&0].inn[&9], vec![5]);
456
457        // Dirty delta still encodes as a sorted unique Vec.
458        let mut dirty = Topology::new();
459        for dst in (0u32..10).rev() {
460            dirty.add_edge(1, 0, dst);
461        }
462        assert!(matches!(
463            dirty.neighbors(1, Direction::Out, 0),
464            Cow::Owned(_)
465        ));
466        let dirty_wire: Wire = bincode::deserialize(&bincode::serialize(&dirty).unwrap()).unwrap();
467        assert_eq!(dirty_wire.edge_count, 10);
468        assert_eq!(
469            dirty_wire.by_type[&1].out[&0],
470            (0..10).collect::<Vec<u32>>()
471        );
472        for dst in 0..10 {
473            assert_eq!(dirty_wire.by_type[&1].inn[&dst], vec![0]);
474        }
475    }
476
477    #[test]
478    fn pack_roundtrip_merges_delta_and_restores_frozen() {
479        let mut t = Topology::new();
480        for dst in (0u32..10).rev() {
481            t.add_edge(2, 1, dst);
482        }
483        t.add_edge(0, 5, 9);
484        assert!(matches!(t.neighbors(2, Direction::Out, 1), Cow::Owned(_)));
485        let mut buf = Vec::new();
486        t.pack(&mut buf);
487        let (back, consumed) = Topology::unpack(&buf).unwrap();
488        assert_eq!(consumed, buf.len());
489        assert_eq!(back.edge_count(), 11);
490        let expected: Vec<u32> = (0..10).collect();
491        assert_eq!(
492            back.neighbors(2, Direction::Out, 1).as_ref(),
493            expected.as_slice()
494        );
495        assert!(matches!(
496            back.neighbors(2, Direction::Out, 1),
497            Cow::Borrowed(_)
498        ));
499        assert_eq!(back.neighbors(0, Direction::Out, 5).as_ref(), &[9]);
500        assert_eq!(back.neighbors(0, Direction::In, 9).as_ref(), &[5]);
501        assert_eq!(back.etypes().collect::<Vec<_>>(), vec![0, 2]);
502    }
503}
core_storage/topology.rs

core_storage/
topology.rs
