Skip to main content

leyline_engine/
memtable.rs

1use std::collections::BTreeSet;
2
3use crate::ids::{EdgeLabel, NodeId};
4use crate::Result;
5
6/// In-memory write buffer, drained into an immutable segment at flush.
7///
8/// Edges live in two BTreeSets sorted exactly the way segments are:
9/// `(src, label, dst)` forward and `(dst, label, src)` reverse. Flush is a
10/// straight iteration into CSR runs, and a neighborhood is one contiguous
11/// range scan — the same access shape the segment reader will have.
12pub struct Memtable {
13    nodes: BTreeSet<NodeId>,
14    fwd: BTreeSet<(NodeId, EdgeLabel, NodeId)>,
15    rev: BTreeSet<(NodeId, EdgeLabel, NodeId)>,
16}
17
18impl Memtable {
19    pub fn new() -> Self {
20        Self {
21            nodes: BTreeSet::new(),
22            fwd: BTreeSet::new(),
23            rev: BTreeSet::new(),
24        }
25    }
26
27    pub fn insert_node(&mut self, id: NodeId) {
28        self.nodes.insert(id);
29    }
30
31    pub fn contains_node(&self, id: NodeId) -> bool {
32        self.nodes.contains(&id)
33    }
34
35    pub fn insert_edge(&mut self, src: NodeId, label: EdgeLabel, dst: NodeId) -> Result<()> {
36        self.fwd.insert((src, label, dst));
37        self.rev.insert((dst, label, src));
38        Ok(())
39    }
40
41    /// Outbound neighbors of `src` under `label`, ascending.
42    pub fn out(&self, src: NodeId, label: EdgeLabel) -> impl Iterator<Item = NodeId> + '_ {
43        self.fwd
44            .range((src, label, NodeId(u64::MIN))..=(src, label, NodeId(u64::MAX)))
45            .map(|&(_, _, dst)| dst)
46    }
47
48    /// Inbound neighbors of `dst` under `label`, ascending.
49    pub fn r#in(&self, dst: NodeId, label: EdgeLabel) -> impl Iterator<Item = NodeId> + '_ {
50        self.rev
51            .range((dst, label, NodeId(u64::MIN))..=(dst, label, NodeId(u64::MAX)))
52            .map(|&(_, _, src)| src)
53    }
54
55    pub fn edge_count(&self) -> usize {
56        self.fwd.len()
57    }
58
59    pub fn is_empty(&self) -> bool {
60        self.nodes.is_empty() && self.fwd.is_empty()
61    }
62
63    /// Flush order: already (src, label, dst)-sorted, one pass to encode.
64    pub fn iter_fwd(&self) -> impl ExactSizeIterator<Item = (NodeId, EdgeLabel, NodeId)> + '_ {
65        self.fwd.iter().copied()
66    }
67
68    pub fn iter_rev(&self) -> impl ExactSizeIterator<Item = (NodeId, EdgeLabel, NodeId)> + '_ {
69        self.rev.iter().copied()
70    }
71
72    pub fn clear(&mut self) {
73        self.nodes.clear();
74        self.fwd.clear();
75        self.rev.clear();
76    }
77}
78
79impl Default for Memtable {
80    fn default() -> Self {
81        Self::new()
82    }
83}