Skip to main content

lago_knowledge/
traversal.rs

1//! BFS graph traversal over wikilink edges.
2
3use std::collections::{HashSet, VecDeque};
4
5use serde::Serialize;
6
7use crate::index::KnowledgeIndex;
8
9/// A note returned from graph traversal, with its depth from the start.
10#[derive(Debug, Clone, Serialize)]
11pub struct TraversalResult {
12    /// Relative path.
13    pub path: String,
14    /// Note name.
15    pub name: String,
16    /// Depth from the start note (0 = the start note itself).
17    pub depth: usize,
18    /// Outgoing wikilinks.
19    pub links: Vec<String>,
20}
21
22impl KnowledgeIndex {
23    /// BFS traversal from a starting note, up to `depth` hops and
24    /// `max_notes` total results.
25    ///
26    /// The start note is resolved via wikilink resolution (name or path).
27    /// Returns notes in BFS order with their depth from the start.
28    pub fn traverse(&self, start: &str, depth: usize, max_notes: usize) -> Vec<TraversalResult> {
29        let mut visited = HashSet::new();
30        let mut results = Vec::new();
31        let mut queue = VecDeque::new();
32
33        // Resolve the start note
34        let start_note = match self.resolve_wikilink(start) {
35            Some(note) => note,
36            None => return results,
37        };
38
39        queue.push_back((start_note.path.clone(), 0usize));
40
41        while let Some((path, level)) = queue.pop_front() {
42            if results.len() >= max_notes {
43                break;
44            }
45
46            if !visited.insert(path.clone()) {
47                continue;
48            }
49
50            let note = match self.notes.get(&path) {
51                Some(n) => n,
52                None => continue,
53            };
54
55            results.push(TraversalResult {
56                path: note.path.clone(),
57                name: note.name.clone(),
58                depth: level,
59                links: note.links.clone(),
60            });
61
62            // Enqueue linked notes if within depth
63            if level < depth {
64                for link_target in &note.links {
65                    if let Some(linked_note) = self.resolve_wikilink(link_target) {
66                        if !visited.contains(&linked_note.path) {
67                            queue.push_back((linked_note.path.clone(), level + 1));
68                        }
69                    }
70                }
71            }
72        }
73
74        results
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use crate::index::KnowledgeIndex;
81    use lago_core::ManifestEntry;
82    use lago_store::BlobStore;
83    use tempfile::TempDir;
84
85    fn build_index(files: &[(&str, &str)]) -> (TempDir, KnowledgeIndex) {
86        let tmp = TempDir::new().unwrap();
87        let store = BlobStore::open(tmp.path()).unwrap();
88        let mut entries = Vec::new();
89
90        for (path, content) in files {
91            let hash = store.put(content.as_bytes()).unwrap();
92            entries.push(ManifestEntry {
93                path: path.to_string(),
94                blob_hash: hash,
95                size_bytes: content.len() as u64,
96                content_type: Some("text/markdown".to_string()),
97                updated_at: 0,
98            });
99        }
100
101        let index = KnowledgeIndex::build(&entries, &store).unwrap();
102        (tmp, index)
103    }
104
105    #[test]
106    fn linear_chain() {
107        // A → B → C
108        let (_tmp, index) = build_index(&[
109            ("/a.md", "# A\n\nSee [[B]]."),
110            ("/b.md", "# B\n\nSee [[C]]."),
111            ("/c.md", "# C\n\nEnd."),
112        ]);
113
114        let results = index.traverse("A", 2, 10);
115        assert_eq!(results.len(), 3);
116        assert_eq!(results[0].name, "a");
117        assert_eq!(results[0].depth, 0);
118        assert_eq!(results[1].name, "b");
119        assert_eq!(results[1].depth, 1);
120        assert_eq!(results[2].name, "c");
121        assert_eq!(results[2].depth, 2);
122    }
123
124    #[test]
125    fn branching_graph() {
126        // A → B, A → C
127        let (_tmp, index) = build_index(&[
128            ("/a.md", "# A\n\nSee [[B]] and [[C]]."),
129            ("/b.md", "# B\n\nLeaf."),
130            ("/c.md", "# C\n\nLeaf."),
131        ]);
132
133        let results = index.traverse("A", 1, 10);
134        assert_eq!(results.len(), 3);
135        assert_eq!(results[0].name, "a");
136    }
137
138    #[test]
139    fn cycle_handling() {
140        // A → B → A (cycle)
141        let (_tmp, index) = build_index(&[
142            ("/a.md", "# A\n\nSee [[B]]."),
143            ("/b.md", "# B\n\nSee [[A]]."),
144        ]);
145
146        let results = index.traverse("A", 5, 10);
147        assert_eq!(results.len(), 2); // Should not loop
148    }
149
150    #[test]
151    fn max_depth_zero() {
152        let (_tmp, index) =
153            build_index(&[("/a.md", "# A\n\nSee [[B]]."), ("/b.md", "# B\n\nEnd.")]);
154
155        let results = index.traverse("A", 0, 10);
156        assert_eq!(results.len(), 1); // Only the start note
157        assert_eq!(results[0].name, "a");
158    }
159
160    #[test]
161    fn max_notes_limit() {
162        let (_tmp, index) = build_index(&[
163            ("/a.md", "# A\n\nSee [[B]] and [[C]] and [[D]]."),
164            ("/b.md", "# B\n\nEnd."),
165            ("/c.md", "# C\n\nEnd."),
166            ("/d.md", "# D\n\nEnd."),
167        ]);
168
169        let results = index.traverse("A", 1, 2);
170        assert_eq!(results.len(), 2);
171    }
172
173    #[test]
174    fn missing_start_note() {
175        let (_tmp, index) = build_index(&[("/a.md", "# A")]);
176        let results = index.traverse("nonexistent", 1, 10);
177        assert!(results.is_empty());
178    }
179}