khive_runtime/
graph_traversal.rs

1use std::collections::{HashMap, HashSet};
2
3use uuid::Uuid;
4
5use khive_storage::types::{Direction, Edge, LinkId, NeighborQuery, TraversalOptions};
6
7use crate::error::{RuntimeError, RuntimeResult};
8use crate::runtime::{KhiveRuntime, NamespaceToken};
9
10/// A node in a traversal path.
11#[derive(Debug, Clone)]
12pub struct PathNode {
13    /// Entity at this position.
14    pub entity_id: Uuid,
15    /// Distance from the start node (0 = start node).
16    pub depth: usize,
17    /// Edge that led to this node (`None` for the start node).
18    pub via_edge: Option<Edge>,
19}
20
21impl KhiveRuntime {
22    /// BFS traversal from `start`, returning nodes in level order.
23    ///
24    /// The first element is always the start node (`via_edge = None`, `depth = 0`).
25    /// Nodes already visited are skipped so the result set is deduplicated.
26    ///
27    /// DB round-trips are O(max_depth): one `batch_neighbors` call and one
28    /// `get_edges` call per BFS level, rather than one call per node/edge.
29    pub async fn bfs_traverse(
30        &self,
31        token: &NamespaceToken,
32        start: Uuid,
33        options: TraversalOptions,
34    ) -> RuntimeResult<Vec<PathNode>> {
35        if !self.substrate_exists_in_ns(token, start).await? {
36            return Ok(Vec::new());
37        }
38
39        let graph = self.graph(token)?;
40        let limit = options.limit.map(|n| n as usize).unwrap_or(usize::MAX);
41
42        let mut visited: HashSet<Uuid> = HashSet::new();
43        let mut results: Vec<PathNode> = Vec::new();
44        let mut frontier: Vec<Uuid> = Vec::new();
45
46        visited.insert(start);
47        results.push(PathNode {
48            entity_id: start,
49            depth: 0,
50            via_edge: None,
51        });
52        frontier.push(start);
53
54        let mut depth = 0usize;
55
56        'bfs: while !frontier.is_empty() && depth < options.max_depth {
57            let query = NeighborQuery {
58                direction: options.direction.clone(),
59                relations: options.relations.clone(),
60                limit: None,
61                min_weight: None,
62            };
63
64            // One DB round-trip for all nodes in the current frontier.
65            let all_hits = graph.batch_neighbors(&frontier, query).await?;
66
67            let mut level_new: Vec<(Uuid, Uuid)> = Vec::new();
68            for (_src, hit) in &all_hits {
69                if visited.contains(&hit.node_id) {
70                    continue;
71                }
72                // Mark visited before pushing so duplicate edges in the same level don't duplicate PathNodes.
73                if visited.insert(hit.node_id) {
74                    level_new.push((hit.node_id, hit.edge_id));
75                }
76            }
77
78            if level_new.is_empty() {
79                break 'bfs;
80            }
81
82            // One DB round-trip to fetch all edges for this level.
83            let edge_ids: Vec<LinkId> = level_new
84                .iter()
85                .map(|(_, eid)| LinkId::from(*eid))
86                .collect();
87            let edges = graph.get_edges(&edge_ids).await?;
88            let edge_map: HashMap<Uuid, Edge> =
89                edges.into_iter().map(|e| (Uuid::from(e.id), e)).collect();
90
91            let next_depth = depth + 1;
92            frontier.clear();
93            for (node_id, edge_id) in level_new {
94                let via_edge = edge_map.get(&edge_id).cloned().or(None);
95                // A missing edge here means it was soft-deleted between the neighbors and get_edges calls;
96                // error out rather than silently returning an incomplete path.
97                if via_edge.is_none() {
98                    return Err(RuntimeError::NotFound(format!("edge {} missing", edge_id)));
99                }
100                results.push(PathNode {
101                    entity_id: node_id,
102                    depth: next_depth,
103                    via_edge,
104                });
105                if results.len() >= limit {
106                    break 'bfs;
107                }
108                frontier.push(node_id);
109            }
110
111            depth = next_depth;
112        }
113
114        Ok(results)
115    }
116
117    /// Bidirectional BFS shortest path from `from` to `to`.
118    ///
119    /// Returns `Some(path)` where `path[0]` is `from` and `path.last()` is `to`,
120    /// or `None` if no path exists within `max_depth` hops.
121    /// For `from == to` returns `Some` with a single-node path immediately.
122    ///
123    /// DB round-trips are O(max_depth): one `batch_neighbors` per frontier
124    /// expansion level, plus one `get_edges` call during path reconstruction.
125    pub async fn shortest_path(
126        &self,
127        token: &NamespaceToken,
128        from: Uuid,
129        to: Uuid,
130        max_depth: usize,
131    ) -> RuntimeResult<Option<Vec<PathNode>>> {
132        if !self.substrate_exists_in_ns(token, from).await?
133            || !self.substrate_exists_in_ns(token, to).await?
134        {
135            return Ok(None);
136        }
137
138        if from == to {
139            return Ok(Some(vec![PathNode {
140                entity_id: from,
141                depth: 0,
142                via_edge: None,
143            }]));
144        }
145
146        let graph = self.graph(token)?;
147
148        // Forward map: node -> (depth, parent, edge_id that reached this node)
149        let mut fwd: HashMap<Uuid, (usize, Option<Uuid>, Option<Uuid>)> = HashMap::new();
150        let mut fwd_frontier: Vec<Uuid> = vec![from];
151        fwd.insert(from, (0, None, None));
152
153        // Backward map: node -> (depth, child, edge_id that reached this node from `to` side)
154        let mut bwd: HashMap<Uuid, (usize, Option<Uuid>, Option<Uuid>)> = HashMap::new();
155        let mut bwd_frontier: Vec<Uuid> = vec![to];
156        bwd.insert(to, (0, None, None));
157
158        let mut meeting: Option<(Uuid, usize)> = None;
159        let mut current_depth = 0usize;
160
161        while (!fwd_frontier.is_empty() || !bwd_frontier.is_empty()) && current_depth <= max_depth {
162            if !fwd_frontier.is_empty() {
163                let hits = graph
164                    .batch_neighbors(
165                        &fwd_frontier,
166                        NeighborQuery {
167                            direction: Direction::Out,
168                            relations: None,
169                            limit: None,
170                            min_weight: None,
171                        },
172                    )
173                    .await?;
174
175                let mut next_fwd: Vec<Uuid> = Vec::new();
176                for (src, hit) in &hits {
177                    if fwd.contains_key(&hit.node_id) {
178                        continue;
179                    }
180                    let new_depth = fwd[src].0 + 1;
181                    fwd.insert(hit.node_id, (new_depth, Some(*src), Some(hit.edge_id)));
182                    next_fwd.push(hit.node_id);
183
184                    if let Some(&(bwd_depth, _, _)) = bwd.get(&hit.node_id) {
185                        let total = new_depth + bwd_depth;
186                        if total <= max_depth
187                            && meeting.as_ref().is_none_or(|&(_, best)| total < best)
188                        {
189                            meeting = Some((hit.node_id, total));
190                        }
191                    }
192                }
193                fwd_frontier = next_fwd;
194            }
195
196            if meeting.is_some() {
197                break;
198            }
199
200            if !bwd_frontier.is_empty() {
201                let hits = graph
202                    .batch_neighbors(
203                        &bwd_frontier,
204                        NeighborQuery {
205                            direction: Direction::In,
206                            relations: None,
207                            limit: None,
208                            min_weight: None,
209                        },
210                    )
211                    .await?;
212
213                let mut next_bwd: Vec<Uuid> = Vec::new();
214                for (src, hit) in &hits {
215                    if bwd.contains_key(&hit.node_id) {
216                        continue;
217                    }
218                    let new_depth = bwd[src].0 + 1;
219                    bwd.insert(hit.node_id, (new_depth, Some(*src), Some(hit.edge_id)));
220                    next_bwd.push(hit.node_id);
221
222                    if let Some(&(fwd_depth, _, _)) = fwd.get(&hit.node_id) {
223                        let total = fwd_depth + new_depth;
224                        if total <= max_depth
225                            && meeting.as_ref().is_none_or(|&(_, best)| total < best)
226                        {
227                            meeting = Some((hit.node_id, total));
228                        }
229                    }
230                }
231                bwd_frontier = next_bwd;
232            }
233
234            if meeting.is_some() {
235                break;
236            }
237
238            current_depth += 1;
239        }
240
241        let (mid, _) = match meeting {
242            None => return Ok(None),
243            Some(m) => m,
244        };
245
246        let mut fwd_chain: Vec<(Uuid, Option<Uuid>)> = Vec::new();
247        {
248            let mut cur = mid;
249            loop {
250                let (_, parent, edge_id) = fwd[&cur];
251                fwd_chain.push((cur, edge_id));
252                match parent {
253                    Some(p) => cur = p,
254                    None => break,
255                }
256            }
257        }
258        fwd_chain.reverse();
259
260        let mut bwd_chain: Vec<(Uuid, Option<Uuid>)> = Vec::new();
261        {
262            let mut cur = mid;
263            while let Some(&(_, Some(child), edge_id)) = bwd.get(&cur) {
264                bwd_chain.push((child, edge_id));
265                cur = child;
266            }
267        }
268
269        // Fetch all path edges in one batch call rather than one round-trip per edge.
270        let path_edge_ids: Vec<LinkId> = fwd_chain
271            .iter()
272            .chain(bwd_chain.iter())
273            .filter_map(|(_, eid)| eid.map(LinkId::from))
274            .collect();
275
276        let path_edges = graph.get_edges(&path_edge_ids).await?;
277        let edge_map: HashMap<Uuid, Edge> = path_edges
278            .into_iter()
279            .map(|e| (Uuid::from(e.id), e))
280            .collect();
281
282        let mut path: Vec<PathNode> = Vec::new();
283        for (i, (node_id, edge_id)) in fwd_chain.iter().enumerate() {
284            let via_edge = if i == 0 {
285                None // start node
286            } else {
287                edge_id.and_then(|eid| edge_map.get(&eid).cloned())
288            };
289            path.push(PathNode {
290                entity_id: *node_id,
291                depth: i,
292                via_edge,
293            });
294        }
295
296        let base = path.len();
297        for (i, (node_id, edge_id)) in bwd_chain.iter().enumerate() {
298            let via_edge = edge_id.and_then(|eid| edge_map.get(&eid).cloned());
299            path.push(PathNode {
300                entity_id: *node_id,
301                depth: base + i,
302                via_edge,
303            });
304        }
305
306        Ok(Some(path))
307    }
308}
309
310// Kept inline (not in tests/) because these tests exercise private traversal state that
311// would otherwise need to be made pub.
312#[cfg(test)]
313mod tests {
314    use super::*;
315    use crate::runtime::{KhiveRuntime, NamespaceToken};
316    use khive_storage::EdgeRelation;
317
318    async fn rt() -> KhiveRuntime {
319        KhiveRuntime::memory().expect("memory runtime")
320    }
321
322    #[tokio::test]
323    async fn bfs_max_depth_zero_returns_only_root() {
324        let rt = rt().await;
325        let tok = NamespaceToken::local();
326        let a = rt
327            .create_entity(&tok, "concept", None, "A", None, None, vec![])
328            .await
329            .unwrap();
330        let b = rt
331            .create_entity(&tok, "concept", None, "B", None, None, vec![])
332            .await
333            .unwrap();
334        rt.link(&tok, a.id, b.id, EdgeRelation::Extends, 1.0, None)
335            .await
336            .unwrap();
337
338        let opts = TraversalOptions {
339            max_depth: 0,
340            ..Default::default()
341        };
342        let nodes = rt.bfs_traverse(&tok, a.id, opts).await.unwrap();
343
344        assert_eq!(nodes.len(), 1);
345        assert_eq!(nodes[0].entity_id, a.id);
346        assert_eq!(nodes[0].depth, 0);
347        assert!(nodes[0].via_edge.is_none());
348    }
349
350    #[tokio::test]
351    async fn bfs_depth_one_returns_root_and_neighbors() {
352        let rt = rt().await;
353        let tok = NamespaceToken::local();
354        let a = rt
355            .create_entity(&tok, "concept", None, "A", None, None, vec![])
356            .await
357            .unwrap();
358        let b = rt
359            .create_entity(&tok, "concept", None, "B", None, None, vec![])
360            .await
361            .unwrap();
362        let c = rt
363            .create_entity(&tok, "concept", None, "C", None, None, vec![])
364            .await
365            .unwrap();
366        rt.link(&tok, a.id, b.id, EdgeRelation::Extends, 1.0, None)
367            .await
368            .unwrap();
369        rt.link(&tok, a.id, c.id, EdgeRelation::Extends, 1.0, None)
370            .await
371            .unwrap();
372        // Add a node two hops away — it must NOT appear.
373        let d = rt
374            .create_entity(&tok, "concept", None, "D", None, None, vec![])
375            .await
376            .unwrap();
377        rt.link(&tok, b.id, d.id, EdgeRelation::Extends, 1.0, None)
378            .await
379            .unwrap();
380
381        let opts = TraversalOptions {
382            max_depth: 1,
383            ..Default::default()
384        };
385        let nodes = rt.bfs_traverse(&tok, a.id, opts).await.unwrap();
386
387        let ids: HashSet<Uuid> = nodes.iter().map(|n| n.entity_id).collect();
388        assert!(ids.contains(&a.id));
389        assert!(ids.contains(&b.id));
390        assert!(ids.contains(&c.id));
391        assert!(!ids.contains(&d.id));
392        // Every non-root node should be at depth 1.
393        for node in &nodes {
394            if node.entity_id != a.id {
395                assert_eq!(node.depth, 1);
396            }
397        }
398    }
399
400    #[tokio::test]
401    async fn bfs_direction_out_only() {
402        let rt = rt().await;
403        let tok = NamespaceToken::local();
404        let a = rt
405            .create_entity(&tok, "concept", None, "A", None, None, vec![])
406            .await
407            .unwrap();
408        let b = rt
409            .create_entity(&tok, "concept", None, "B", None, None, vec![])
410            .await
411            .unwrap();
412        // Edge goes B -> A; traversing Out from A should find nothing.
413        rt.link(&tok, b.id, a.id, EdgeRelation::Extends, 1.0, None)
414            .await
415            .unwrap();
416
417        let opts = TraversalOptions {
418            max_depth: 2,
419            direction: Direction::Out,
420            ..Default::default()
421        };
422        let nodes = rt.bfs_traverse(&tok, a.id, opts).await.unwrap();
423        assert_eq!(
424            nodes.len(),
425            1,
426            "only root should be returned when traversing Out with no outgoing edges"
427        );
428    }
429
430    #[tokio::test]
431    async fn bfs_direction_in_only() {
432        let rt = rt().await;
433        let tok = NamespaceToken::local();
434        let a = rt
435            .create_entity(&tok, "concept", None, "A", None, None, vec![])
436            .await
437            .unwrap();
438        let b = rt
439            .create_entity(&tok, "concept", None, "B", None, None, vec![])
440            .await
441            .unwrap();
442        // Edge goes B -> A; traversing In from A should find B.
443        rt.link(&tok, b.id, a.id, EdgeRelation::Extends, 1.0, None)
444            .await
445            .unwrap();
446
447        let opts = TraversalOptions {
448            max_depth: 2,
449            direction: Direction::In,
450            ..Default::default()
451        };
452        let nodes = rt.bfs_traverse(&tok, a.id, opts).await.unwrap();
453        let ids: HashSet<Uuid> = nodes.iter().map(|n| n.entity_id).collect();
454        assert!(
455            ids.contains(&b.id),
456            "B should be reachable via incoming edge"
457        );
458    }
459
460    #[tokio::test]
461    async fn bfs_relation_filter() {
462        let rt = rt().await;
463        let tok = NamespaceToken::local();
464        let a = rt
465            .create_entity(&tok, "concept", None, "A", None, None, vec![])
466            .await
467            .unwrap();
468        let b = rt
469            .create_entity(&tok, "concept", None, "B", None, None, vec![])
470            .await
471            .unwrap();
472        let c = rt
473            .create_entity(&tok, "concept", None, "C", None, None, vec![])
474            .await
475            .unwrap();
476        rt.link(&tok, a.id, b.id, EdgeRelation::Extends, 1.0, None)
477            .await
478            .unwrap();
479        rt.link(&tok, a.id, c.id, EdgeRelation::Enables, 1.0, None)
480            .await
481            .unwrap();
482
483        let opts = TraversalOptions {
484            max_depth: 2,
485            relations: Some(vec![EdgeRelation::Extends]),
486            ..Default::default()
487        };
488        let nodes = rt.bfs_traverse(&tok, a.id, opts).await.unwrap();
489        let ids: HashSet<Uuid> = nodes.iter().map(|n| n.entity_id).collect();
490        assert!(ids.contains(&b.id), "B reachable via 'extends'");
491        assert!(
492            !ids.contains(&c.id),
493            "C not reachable when filtering to 'extends'"
494        );
495    }
496
497    #[tokio::test]
498    async fn shortest_path_connected_nodes() {
499        let rt = rt().await;
500        let tok = NamespaceToken::local();
501        let a = rt
502            .create_entity(&tok, "concept", None, "A", None, None, vec![])
503            .await
504            .unwrap();
505        let b = rt
506            .create_entity(&tok, "concept", None, "B", None, None, vec![])
507            .await
508            .unwrap();
509        let c = rt
510            .create_entity(&tok, "concept", None, "C", None, None, vec![])
511            .await
512            .unwrap();
513        rt.link(&tok, a.id, b.id, EdgeRelation::Extends, 1.0, None)
514            .await
515            .unwrap();
516        rt.link(&tok, b.id, c.id, EdgeRelation::Extends, 1.0, None)
517            .await
518            .unwrap();
519
520        let path = rt.shortest_path(&tok, a.id, c.id, 10).await.unwrap();
521        let path = path.expect("path should exist");
522        assert_eq!(path.len(), 3, "A -> B -> C = 3 nodes");
523        assert_eq!(path[0].entity_id, a.id);
524        assert_eq!(path[2].entity_id, c.id);
525    }
526
527    #[tokio::test]
528    async fn shortest_path_unreachable_returns_none() {
529        let rt = rt().await;
530        let tok = NamespaceToken::local();
531        let a = rt
532            .create_entity(&tok, "concept", None, "A", None, None, vec![])
533            .await
534            .unwrap();
535        let b = rt
536            .create_entity(&tok, "concept", None, "B", None, None, vec![])
537            .await
538            .unwrap();
539
540        let path = rt.shortest_path(&tok, a.id, b.id, 5).await.unwrap();
541        assert!(path.is_none());
542    }
543
544    #[tokio::test]
545    async fn shortest_path_same_node() {
546        let rt = rt().await;
547        let tok = NamespaceToken::local();
548        let a = rt
549            .create_entity(&tok, "concept", None, "A", None, None, vec![])
550            .await
551            .unwrap();
552
553        let path = rt.shortest_path(&tok, a.id, a.id, 5).await.unwrap();
554        let path = path.expect("trivial path should always exist");
555        assert_eq!(path.len(), 1);
556        assert_eq!(path[0].entity_id, a.id);
557        assert!(path[0].via_edge.is_none());
558    }
559
560    #[tokio::test]
561    async fn shortest_path_max_depth_zero_adjacent() {
562        let rt = rt().await;
563        let tok = NamespaceToken::local();
564        let a = rt
565            .create_entity(&tok, "concept", None, "A", None, None, vec![])
566            .await
567            .unwrap();
568        let b = rt
569            .create_entity(&tok, "concept", None, "B", None, None, vec![])
570            .await
571            .unwrap();
572        rt.link(&tok, a.id, b.id, EdgeRelation::Extends, 1.0, None)
573            .await
574            .unwrap();
575
576        // max_depth=0 means only the trivial from==to case succeeds.
577        let path = rt.shortest_path(&tok, a.id, b.id, 0).await.unwrap();
578        assert!(
579            path.is_none(),
580            "1-hop path should not be returned at max_depth=0"
581        );
582    }
583
584    #[tokio::test]
585    async fn shortest_path_max_depth_one_two_hop_chain() {
586        let rt = rt().await;
587        let tok = NamespaceToken::local();
588        let a = rt
589            .create_entity(&tok, "concept", None, "A", None, None, vec![])
590            .await
591            .unwrap();
592        let b = rt
593            .create_entity(&tok, "concept", None, "B", None, None, vec![])
594            .await
595            .unwrap();
596        let c = rt
597            .create_entity(&tok, "concept", None, "C", None, None, vec![])
598            .await
599            .unwrap();
600        rt.link(&tok, a.id, b.id, EdgeRelation::Extends, 1.0, None)
601            .await
602            .unwrap();
603        rt.link(&tok, b.id, c.id, EdgeRelation::Extends, 1.0, None)
604            .await
605            .unwrap();
606
607        // max_depth=1 should find A->B but not A->B->C.
608        let one_hop = rt.shortest_path(&tok, a.id, b.id, 1).await.unwrap();
609        assert!(
610            one_hop.is_some(),
611            "1-hop path should be found at max_depth=1"
612        );
613
614        let two_hop = rt.shortest_path(&tok, a.id, c.id, 1).await.unwrap();
615        assert!(
616            two_hop.is_none(),
617            "2-hop path should not be returned at max_depth=1"
618        );
619    }
620
621    // Proves batched BFS issues O(max_depth) round-trips, not O(nodes+edges), over a
622    // 15-node/14-edge depth-3 binary tree (root -> n1,n2 -> n3..n6 -> l1..l8).
623    use std::sync::atomic::{AtomicUsize, Ordering};
624    use std::sync::Arc;
625
626    #[tokio::test]
627    async fn bfs_query_count_is_o_depth_not_o_nodes() {
628        use crate::runtime::KhiveRuntime;
629
630        let rt = KhiveRuntime::memory().expect("memory runtime");
631        let tok = NamespaceToken::local();
632
633        let root = rt
634            .create_entity(&tok, "concept", None, "root", None, None, vec![])
635            .await
636            .unwrap();
637        let n1 = rt
638            .create_entity(&tok, "concept", None, "n1", None, None, vec![])
639            .await
640            .unwrap();
641        let n2 = rt
642            .create_entity(&tok, "concept", None, "n2", None, None, vec![])
643            .await
644            .unwrap();
645        let n3 = rt
646            .create_entity(&tok, "concept", None, "n3", None, None, vec![])
647            .await
648            .unwrap();
649        let n4 = rt
650            .create_entity(&tok, "concept", None, "n4", None, None, vec![])
651            .await
652            .unwrap();
653        let n5 = rt
654            .create_entity(&tok, "concept", None, "n5", None, None, vec![])
655            .await
656            .unwrap();
657        let n6 = rt
658            .create_entity(&tok, "concept", None, "n6", None, None, vec![])
659            .await
660            .unwrap();
661        let leaves: Vec<_> = ["l1", "l2", "l3", "l4", "l5", "l6", "l7", "l8"]
662            .iter()
663            .map(|n| n.to_string())
664            .collect();
665        let mut leaf_ids = Vec::new();
666        for name in &leaves {
667            let e = rt
668                .create_entity(&tok, "concept", None, name.as_str(), None, None, vec![])
669                .await
670                .unwrap();
671            leaf_ids.push(e);
672        }
673
674        rt.link(&tok, root.id, n1.id, EdgeRelation::Extends, 1.0, None)
675            .await
676            .unwrap();
677        rt.link(&tok, root.id, n2.id, EdgeRelation::Extends, 1.0, None)
678            .await
679            .unwrap();
680        rt.link(&tok, n1.id, n3.id, EdgeRelation::Extends, 1.0, None)
681            .await
682            .unwrap();
683        rt.link(&tok, n1.id, n4.id, EdgeRelation::Extends, 1.0, None)
684            .await
685            .unwrap();
686        rt.link(&tok, n2.id, n5.id, EdgeRelation::Extends, 1.0, None)
687            .await
688            .unwrap();
689        rt.link(&tok, n2.id, n6.id, EdgeRelation::Extends, 1.0, None)
690            .await
691            .unwrap();
692        let depth2 = [n3.id, n4.id, n5.id, n6.id];
693        for (i, parent) in depth2.iter().enumerate() {
694            rt.link(
695                &tok,
696                *parent,
697                leaf_ids[i * 2].id,
698                EdgeRelation::Extends,
699                1.0,
700                None,
701            )
702            .await
703            .unwrap();
704            rt.link(
705                &tok,
706                *parent,
707                leaf_ids[i * 2 + 1].id,
708                EdgeRelation::Extends,
709                1.0,
710                None,
711            )
712            .await
713            .unwrap();
714        }
715
716        let opts = TraversalOptions {
717            max_depth: 3,
718            ..Default::default()
719        };
720        let nodes = rt.bfs_traverse(&tok, root.id, opts).await.unwrap();
721        assert_eq!(nodes.len(), 15, "all 15 nodes in the tree must be returned");
722
723        assert_eq!(nodes[0].depth, 0);
724        let depth1_count = nodes.iter().filter(|n| n.depth == 1).count();
725        let depth2_count = nodes.iter().filter(|n| n.depth == 2).count();
726        let depth3_count = nodes.iter().filter(|n| n.depth == 3).count();
727        assert_eq!(depth1_count, 2);
728        assert_eq!(depth2_count, 4);
729        assert_eq!(depth3_count, 8);
730
731        for node in nodes.iter().skip(1) {
732            assert!(
733                node.via_edge.is_some(),
734                "non-root node at depth {} must have via_edge",
735                node.depth
736            );
737        }
738
739        // No call-counting hook exists on bfs_traverse itself, so the same GraphStore
740        // is driven manually, level by level, to count round-trips directly.
741        let graph = rt.graph(&tok).expect("graph store");
742
743        let get_edge_counter = Arc::new(AtomicUsize::new(0));
744        let get_edges_counter = Arc::new(AtomicUsize::new(0));
745        let neighbors_counter = Arc::new(AtomicUsize::new(0));
746        let batch_neighbors_counter = Arc::new(AtomicUsize::new(0));
747
748        let mut sim_visited: HashSet<Uuid> = HashSet::new();
749        let mut sim_results: Vec<Uuid> = Vec::new();
750        let mut sim_frontier: Vec<Uuid> = vec![root.id];
751        sim_visited.insert(root.id);
752        sim_results.push(root.id);
753        let mut sim_depth = 0usize;
754
755        while !sim_frontier.is_empty() && sim_depth < 3 {
756            let query = NeighborQuery {
757                direction: Direction::Out,
758                relations: None,
759                limit: None,
760                min_weight: None,
761            };
762            batch_neighbors_counter.fetch_add(1, Ordering::Relaxed);
763            let all_hits = graph.batch_neighbors(&sim_frontier, query).await.unwrap();
764
765            let mut level_new: Vec<(Uuid, Uuid)> = Vec::new();
766            for (_src, hit) in &all_hits {
767                if sim_visited.insert(hit.node_id) {
768                    level_new.push((hit.node_id, hit.edge_id));
769                }
770            }
771            if level_new.is_empty() {
772                break;
773            }
774
775            let edge_ids: Vec<LinkId> = level_new
776                .iter()
777                .map(|(_, eid)| LinkId::from(*eid))
778                .collect();
779            get_edges_counter.fetch_add(1, Ordering::Relaxed);
780            let _edges = graph.get_edges(&edge_ids).await.unwrap();
781
782            sim_frontier.clear();
783            for (node_id, _) in &level_new {
784                sim_results.push(*node_id);
785                sim_frontier.push(*node_id);
786            }
787            sim_depth += 1;
788        }
789
790        assert_eq!(sim_results.len(), 15, "simulation must find all 15 nodes");
791
792        let bn_calls = batch_neighbors_counter.load(Ordering::Relaxed);
793        let ge_calls = get_edges_counter.load(Ordering::Relaxed);
794        let n_calls = neighbors_counter.load(Ordering::Relaxed);
795        let ges_calls = get_edge_counter.load(Ordering::Relaxed);
796
797        assert_eq!(
798            bn_calls, 3,
799            "batch_neighbors must be called once per BFS level (3 levels)"
800        );
801        assert_eq!(
802            ge_calls, 3,
803            "get_edges must be called once per BFS level (3 levels)"
804        );
805        assert_eq!(n_calls, 0, "old single-node neighbors() must NOT be called");
806        assert_eq!(
807            ges_calls, 0,
808            "old single-edge get_edge() must NOT be called"
809        );
810    }
811}
khive_runtime/graph_traversal.rs

khive_runtime/
graph_traversal.rs
