core_query/
traverse.rs

1use crate::view::GraphView;
2use core_storage::Direction;
3use std::collections::{BTreeSet, VecDeque};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq)]
6pub enum Dir {
7    Out,
8    In,
9    Both,
10}
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
13pub struct EdgeRef {
14    pub etype: u32,
15    pub src: u32,
16    pub dst: u32,
17}
18
19/// Typed 1-hop expansion. `etypes=None` → all edge types in the graph (sorted).
20/// Deterministic: etype asc, then neighbor asc; `Dir::Both` = Out then In.
21/// Dedupe only identical `EdgeRef` triples — Out and In of a pair are distinct.
22pub fn expand(view: &GraphView, id: u32, etypes: Option<&[u32]>, dir: Dir) -> Vec<EdgeRef> {
23    let types: Vec<u32> = match etypes {
24        Some(ts) => {
25            let mut v = ts.to_vec();
26            v.sort_unstable();
27            v.dedup();
28            v
29        }
30        None => view.topo.etypes().collect(),
31    };
32    let mut out = Vec::new();
33    let mut seen = BTreeSet::new();
34    for etype in types {
35        if matches!(dir, Dir::Out | Dir::Both) {
36            for &dst in view.topo.neighbors(etype, Direction::Out, id).as_ref() {
37                push_unique(
38                    &mut out,
39                    &mut seen,
40                    EdgeRef {
41                        etype,
42                        src: id,
43                        dst,
44                    },
45                );
46            }
47        }
48        if matches!(dir, Dir::In | Dir::Both) {
49            for &src in view.topo.neighbors(etype, Direction::In, id).as_ref() {
50                push_unique(
51                    &mut out,
52                    &mut seen,
53                    EdgeRef {
54                        etype,
55                        src,
56                        dst: id,
57                    },
58                );
59            }
60        }
61    }
62    out
63}
64
65#[derive(Debug, Clone, PartialEq)]
66pub struct Neighborhood {
67    pub nodes: Vec<(u32, u32)>, // (node id, depth), BFS order, start excluded, first-seen depth
68    pub edges: Vec<EdgeRef>,    // every edge traversed to reach a node (deduped, sorted)
69}
70
71pub fn neighborhood(
72    view: &GraphView,
73    start: u32,
74    depth: u32,
75    etypes: Option<&[u32]>,
76    dir: Dir,
77) -> Neighborhood {
78    let mut visited = BTreeSet::new();
79    visited.insert(start);
80    let mut queue = VecDeque::new();
81    queue.push_back((start, 0u32));
82    let mut nodes = Vec::new();
83    let mut edges = BTreeSet::new();
84
85    while let Some((id, d)) = queue.pop_front() {
86        if d >= depth {
87            continue;
88        }
89        for e in expand(view, id, etypes, dir) {
90            edges.insert(e);
91            let nbr = if e.src == id { e.dst } else { e.src };
92            if visited.insert(nbr) {
93                let nd = d + 1;
94                nodes.push((nbr, nd));
95                queue.push_back((nbr, nd));
96            }
97        }
98    }
99
100    Neighborhood {
101        nodes,
102        edges: edges.into_iter().collect(),
103    }
104}
105
106fn push_unique(out: &mut Vec<EdgeRef>, seen: &mut BTreeSet<EdgeRef>, e: EdgeRef) {
107    if seen.insert(e) {
108        out.push(e);
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::{expand, neighborhood, Dir, EdgeRef};
115    use crate::view::GraphView;
116    use core_storage::{ColumnStore, EdgeProps, IdMap, Interner, Topology};
117
118    struct Fx {
119        ids: IdMap,
120        syms: Interner,
121        labels: Vec<u32>,
122        props: ColumnStore,
123        topo: Topology,
124        eprops: EdgeProps,
125    }
126
127    impl Fx {
128        fn new() -> Self {
129            Fx {
130                ids: IdMap::new(),
131                syms: Interner::new(),
132                labels: vec![],
133                props: ColumnStore::new(),
134                topo: Topology::new(),
135                eprops: EdgeProps::new(),
136            }
137        }
138
139        fn add(&mut self, label: &str, key: &str) -> u32 {
140            let id = self.ids.get_or_insert(key);
141            let sym = self.syms.intern(label);
142            self.labels.resize(id as usize + 1, u32::MAX);
143            self.labels[id as usize] = sym;
144            id
145        }
146
147        fn view(&self) -> GraphView<'_> {
148            GraphView {
149                ids: &self.ids,
150                syms: &self.syms,
151                labels: &self.labels,
152                props: &self.props,
153                topo: &self.topo,
154                edge_props: &self.eprops,
155                mask: None,
156            }
157        }
158    }
159
160    /// Diamond + shortcut:
161    ///   A -KNOWS-> B -KNOWS-> D
162    ///   A -KNOWS-> C -KNOWS-> D
163    ///   A -LIKES-> D
164    struct Diamond {
165        fx: Fx,
166        a: u32,
167        b: u32,
168        c: u32,
169        d: u32,
170        knows: u32,
171        likes: u32,
172    }
173
174    fn diamond() -> Diamond {
175        let mut fx = Fx::new();
176        let a = fx.add("Person", "a");
177        let b = fx.add("Person", "b");
178        let c = fx.add("Person", "c");
179        let d = fx.add("Person", "d");
180        let knows = fx.syms.intern("KNOWS");
181        let likes = fx.syms.intern("LIKES");
182        fx.topo.add_edge(knows, a, b);
183        fx.topo.add_edge(knows, a, c);
184        fx.topo.add_edge(knows, b, d);
185        fx.topo.add_edge(knows, c, d);
186        fx.topo.add_edge(likes, a, d);
187        Diamond {
188            fx,
189            a,
190            b,
191            c,
192            d,
193            knows,
194            likes,
195        }
196    }
197
198    fn e(etype: u32, src: u32, dst: u32) -> EdgeRef {
199        EdgeRef { etype, src, dst }
200    }
201
202    #[test]
203    fn expand_etype_then_neighbor_order_and_both_is_out_then_in() {
204        let g = diamond();
205        let v = g.fx.view();
206        assert_eq!(
207            expand(&v, g.a, None, Dir::Out),
208            vec![
209                e(g.knows, g.a, g.b),
210                e(g.knows, g.a, g.c),
211                e(g.likes, g.a, g.d)
212            ]
213        );
214        // caller etype order is ignored; result is still etype asc
215        assert_eq!(
216            expand(&v, g.a, Some(&[g.likes, g.knows]), Dir::Out),
217            vec![
218                e(g.knows, g.a, g.b),
219                e(g.knows, g.a, g.c),
220                e(g.likes, g.a, g.d)
221            ]
222        );
223        assert_eq!(
224            expand(&v, g.d, Some(&[g.knows]), Dir::In),
225            vec![e(g.knows, g.b, g.d), e(g.knows, g.c, g.d)]
226        );
227        assert_eq!(
228            expand(&v, g.b, Some(&[g.knows]), Dir::Both),
229            vec![e(g.knows, g.b, g.d), e(g.knows, g.a, g.b)]
230        );
231        // The same directed triple is visible as Out from src and In from dst.
232        assert!(expand(&v, g.a, Some(&[g.knows]), Dir::Out).contains(&e(g.knows, g.a, g.b)));
233        assert_eq!(
234            expand(&v, g.b, Some(&[g.knows]), Dir::In),
235            vec![e(g.knows, g.a, g.b)]
236        );
237    }
238
239    #[test]
240    fn expand_none_uses_topo_etypes_not_interned_symbols() {
241        let mut fx = Fx::new();
242        // Labels/fields interned between etypes: symbol space ≠ topology etypes.
243        let a = fx.add("Person", "a");
244        let b = fx.add("Person", "b");
245        let knows = fx.syms.intern("KNOWS");
246        let _age = fx.syms.intern("age");
247        let _company = fx.syms.intern("Company");
248        let likes = fx.syms.intern("LIKES");
249        fx.topo.add_edge(knows, a, b);
250        fx.topo.add_edge(likes, a, b);
251        assert!(fx.syms.get("Person").unwrap() < knows);
252        assert!(knows < fx.syms.get("age").unwrap());
253        assert!(fx.syms.get("age").unwrap() < likes);
254        let v = fx.view();
255        assert_eq!(
256            expand(&v, a, None, Dir::Out),
257            vec![e(knows, a, b), e(likes, a, b)]
258        );
259    }
260
261    #[test]
262    fn expand_dedupes_only_identical_triples() {
263        let mut g = diamond();
264        g.fx.topo.add_edge(g.knows, g.a, g.a); // self-loop: Out and In are the same triple
265        let v = g.fx.view();
266        let both = expand(&v, g.a, Some(&[g.knows, g.knows]), Dir::Both);
267        let self_loop_hits = both.iter().filter(|x| *x == &e(g.knows, g.a, g.a)).count();
268        assert_eq!(self_loop_hits, 1);
269        assert!(both.contains(&e(g.knows, g.a, g.b)));
270        assert!(both.contains(&e(g.knows, g.a, g.c)));
271    }
272
273    #[test]
274    fn neighborhood_depth0_empty() {
275        let g = diamond();
276        let v = g.fx.view();
277        let n = neighborhood(&v, g.a, 0, None, Dir::Out);
278        assert!(n.nodes.is_empty());
279        assert!(n.edges.is_empty());
280    }
281
282    #[test]
283    fn neighborhood_depth1_vs_depth2_and_first_seen() {
284        let g = diamond();
285        let v = g.fx.view();
286
287        let d1 = neighborhood(&v, g.a, 1, Some(&[g.knows]), Dir::Out);
288        assert_eq!(d1.nodes, vec![(g.b, 1), (g.c, 1)]);
289        assert_eq!(d1.edges, vec![e(g.knows, g.a, g.b), e(g.knows, g.a, g.c)]);
290
291        let d2 = neighborhood(&v, g.a, 2, Some(&[g.knows]), Dir::Out);
292        assert_eq!(d2.nodes, vec![(g.b, 1), (g.c, 1), (g.d, 2)]);
293        // every traversed edge, including both diamond legs to D
294        assert_eq!(
295            d2.edges,
296            vec![
297                e(g.knows, g.a, g.b),
298                e(g.knows, g.a, g.c),
299                e(g.knows, g.b, g.d),
300                e(g.knows, g.c, g.d),
301            ]
302        );
303
304        // LIKES shortcut: D is first seen at depth 1, not 2
305        let all = neighborhood(&v, g.a, 2, None, Dir::Out);
306        assert_eq!(all.nodes, vec![(g.b, 1), (g.c, 1), (g.d, 1)]);
307        assert_eq!(
308            all.edges,
309            vec![
310                e(g.knows, g.a, g.b),
311                e(g.knows, g.a, g.c),
312                e(g.knows, g.b, g.d),
313                e(g.knows, g.c, g.d),
314                e(g.likes, g.a, g.d),
315            ]
316        );
317    }
318
319    #[test]
320    fn neighborhood_dir_in_out_both_and_etype_filter() {
321        let g = diamond();
322        let v = g.fx.view();
323
324        let inn = neighborhood(&v, g.d, 1, None, Dir::In);
325        assert_eq!(inn.nodes, vec![(g.b, 1), (g.c, 1), (g.a, 1)]);
326        assert_eq!(
327            inn.edges,
328            vec![
329                e(g.knows, g.b, g.d),
330                e(g.knows, g.c, g.d),
331                e(g.likes, g.a, g.d),
332            ]
333        );
334
335        let likes_only = neighborhood(&v, g.a, 2, Some(&[g.likes]), Dir::Out);
336        assert_eq!(likes_only.nodes, vec![(g.d, 1)]);
337        assert_eq!(likes_only.edges, vec![e(g.likes, g.a, g.d)]);
338
339        let both = neighborhood(&v, g.b, 1, Some(&[g.knows]), Dir::Both);
340        assert_eq!(both.nodes, vec![(g.d, 1), (g.a, 1)]);
341        assert_eq!(both.edges, vec![e(g.knows, g.a, g.b), e(g.knows, g.b, g.d)]);
342    }
343
344    #[test]
345    fn neighborhood_is_deterministic() {
346        let g = diamond();
347        let v = g.fx.view();
348        let x = neighborhood(&v, g.a, 2, None, Dir::Both);
349        let y = neighborhood(&v, g.a, 2, None, Dir::Both);
350        assert_eq!(x, y);
351        assert_eq!(
352            expand(&v, g.a, None, Dir::Both),
353            expand(&v, g.a, None, Dir::Both)
354        );
355    }
356}
core_query/traverse.rs

core_query/
traverse.rs
