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

khive_runtime/
graph_traversal.rs
