khive_runtime/
graph_traversal.rs

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

khive_runtime/
graph_traversal.rs
