the_code_graph_domain/
traversal.rs

1use crate::model::*;
2use std::collections::{HashMap, HashSet, VecDeque};
3
4// ---------------------------------------------------------------------------
5// InMemoryGraph
6// ---------------------------------------------------------------------------
7
8pub struct InMemoryGraph {
9    outgoing: HashMap<String, Vec<(String, EdgeKind)>>,
10    incoming: HashMap<String, Vec<(String, EdgeKind)>>,
11}
12
13impl Default for InMemoryGraph {
14    fn default() -> Self {
15        Self::new()
16    }
17}
18
19impl InMemoryGraph {
20    pub fn new() -> Self {
21        InMemoryGraph {
22            outgoing: HashMap::new(),
23            incoming: HashMap::new(),
24        }
25    }
26
27    pub fn add_edge(&mut self, edge: Edge) {
28        self.outgoing
29            .entry(edge.source.clone())
30            .or_default()
31            .push((edge.target.clone(), edge.kind));
32        self.incoming
33            .entry(edge.target)
34            .or_default()
35            .push((edge.source, edge.kind));
36    }
37
38    pub fn from_edges(edges: Vec<Edge>) -> Self {
39        let mut graph = Self::new();
40        for edge in edges {
41            graph.add_edge(edge);
42        }
43        graph
44    }
45
46    pub fn bfs(&self, start: &str, direction: Direction, max_depth: usize) -> Vec<TraversalResult> {
47        self.bfs_inner(start, direction, max_depth, |_| true)
48    }
49
50    pub fn bfs_filtered(
51        &self,
52        start: &str,
53        direction: Direction,
54        max_depth: usize,
55        min_confidence: Confidence,
56    ) -> Vec<TraversalResult> {
57        self.bfs_inner(start, direction, max_depth, move |kind| {
58            kind.confidence() >= min_confidence
59        })
60    }
61
62    fn bfs_inner<F>(
63        &self,
64        start: &str,
65        direction: Direction,
66        max_depth: usize,
67        edge_filter: F,
68    ) -> Vec<TraversalResult>
69    where
70        F: Fn(&EdgeKind) -> bool,
71    {
72        let adjacency = match direction {
73            Direction::Forward => &self.outgoing,
74            Direction::Backward => &self.incoming,
75        };
76
77        let mut results = Vec::new();
78        let mut visited: HashSet<String> = HashSet::new();
79        // Queue entries: (node, depth, path)
80        let mut queue: VecDeque<(String, usize, Vec<String>, EdgeKind)> = VecDeque::new();
81
82        visited.insert(start.to_string());
83
84        if let Some(neighbors) = adjacency.get(start) {
85            for (neighbor, kind) in neighbors {
86                if edge_filter(kind) && !visited.contains(neighbor.as_str()) {
87                    queue.push_back((neighbor.clone(), 1, vec![start.to_string()], *kind));
88                }
89            }
90        }
91
92        while let Some((node, depth, path, edge_kind)) = queue.pop_front() {
93            if visited.contains(&node) {
94                continue;
95            }
96            visited.insert(node.clone());
97
98            let mut new_path = path.clone();
99            new_path.push(node.clone());
100
101            results.push(TraversalResult {
102                node: node.clone(),
103                depth,
104                path: new_path.clone(),
105                edge_kind,
106            });
107
108            if depth < max_depth {
109                if let Some(neighbors) = adjacency.get(&node) {
110                    for (neighbor, kind) in neighbors {
111                        if edge_filter(kind) && !visited.contains(neighbor.as_str()) {
112                            queue.push_back((neighbor.clone(), depth + 1, new_path.clone(), *kind));
113                        }
114                    }
115                }
116            }
117        }
118
119        results
120    }
121
122    pub fn dfs(&self, start: &str, direction: Direction) -> Vec<TraversalResult> {
123        let adjacency = match direction {
124            Direction::Forward => &self.outgoing,
125            Direction::Backward => &self.incoming,
126        };
127
128        let mut results = Vec::new();
129        let mut visited: HashSet<String> = HashSet::new();
130        // Stack entries: (node, depth, path, edge_kind)
131        let mut stack: Vec<(String, usize, Vec<String>, EdgeKind)> = Vec::new();
132
133        visited.insert(start.to_string());
134
135        if let Some(neighbors) = adjacency.get(start) {
136            for (neighbor, kind) in neighbors.iter().rev() {
137                if !visited.contains(neighbor.as_str()) {
138                    stack.push((neighbor.clone(), 1, vec![start.to_string()], *kind));
139                }
140            }
141        }
142
143        while let Some((node, depth, path, edge_kind)) = stack.pop() {
144            if visited.contains(&node) {
145                continue;
146            }
147            visited.insert(node.clone());
148
149            let mut new_path = path.clone();
150            new_path.push(node.clone());
151
152            results.push(TraversalResult {
153                node: node.clone(),
154                depth,
155                path: new_path.clone(),
156                edge_kind,
157            });
158
159            if let Some(neighbors) = adjacency.get(&node) {
160                for (neighbor, kind) in neighbors.iter().rev() {
161                    if !visited.contains(neighbor.as_str()) {
162                        stack.push((neighbor.clone(), depth + 1, new_path.clone(), *kind));
163                    }
164                }
165            }
166        }
167
168        results
169    }
170}
171
172// ---------------------------------------------------------------------------
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    fn make_edge(src: &str, tgt: &str, kind: EdgeKind) -> Edge {
179        Edge {
180            kind,
181            source: src.into(),
182            target: tgt.into(),
183            metadata: None,
184        }
185    }
186
187    #[test]
188    fn bfs_returns_nodes_at_correct_depth() {
189        let edges = vec![
190            make_edge("A", "B", EdgeKind::Calls),
191            make_edge("B", "C", EdgeKind::Calls),
192            make_edge("C", "D", EdgeKind::Calls),
193        ];
194        let graph = InMemoryGraph::from_edges(edges);
195        let results = graph.bfs("A", Direction::Forward, 3);
196        assert_eq!(results.len(), 3);
197        assert!(results.iter().any(|r| r.node == "B" && r.depth == 1));
198        assert!(results.iter().any(|r| r.node == "C" && r.depth == 2));
199        assert!(results.iter().any(|r| r.node == "D" && r.depth == 3));
200    }
201
202    #[test]
203    fn bfs_respects_max_depth() {
204        let edges = vec![
205            make_edge("A", "B", EdgeKind::Calls),
206            make_edge("B", "C", EdgeKind::Calls),
207            make_edge("C", "D", EdgeKind::Calls),
208        ];
209        let graph = InMemoryGraph::from_edges(edges);
210        let results = graph.bfs("A", Direction::Forward, 2);
211        assert_eq!(results.len(), 2);
212    }
213
214    #[test]
215    fn bfs_filtered_excludes_low_confidence() {
216        let edges = vec![
217            make_edge("A", "B", EdgeKind::Calls),       // High
218            make_edge("A", "C", EdgeKind::DependsOn),   // Low
219            make_edge("B", "D", EdgeKind::ImportsFrom), // Medium
220        ];
221        let graph = InMemoryGraph::from_edges(edges);
222        let results = graph.bfs_filtered("A", Direction::Forward, 3, Confidence::High);
223        let names: Vec<&str> = results.iter().map(|r| r.node.as_str()).collect();
224        assert!(names.contains(&"B"));
225        assert!(!names.contains(&"C"));
226        assert!(!names.contains(&"D"));
227    }
228
229    #[test]
230    fn dfs_detects_cycle_without_infinite_loop() {
231        let edges = vec![
232            make_edge("A", "B", EdgeKind::Calls),
233            make_edge("B", "C", EdgeKind::Calls),
234            make_edge("C", "A", EdgeKind::Calls),
235        ];
236        let graph = InMemoryGraph::from_edges(edges);
237        let results = graph.dfs("A", Direction::Forward);
238        assert!(results.len() <= 3);
239        assert!(results.iter().filter(|r| r.node == "B").count() <= 1);
240    }
241
242    #[test]
243    fn bfs_empty_graph_returns_empty() {
244        let graph = InMemoryGraph::from_edges(vec![]);
245        let results = graph.bfs("nonexistent", Direction::Forward, 10);
246        assert!(results.is_empty());
247    }
248
249    #[test]
250    fn dfs_empty_graph_returns_empty() {
251        let graph = InMemoryGraph::from_edges(vec![]);
252        let results = graph.dfs("nonexistent", Direction::Forward);
253        assert!(results.is_empty());
254    }
255
256    #[test]
257    fn self_referential_edge_handled_gracefully() {
258        let edges = vec![make_edge("A", "A", EdgeKind::Calls)];
259        let graph = InMemoryGraph::from_edges(edges);
260        let results = graph.bfs("A", Direction::Forward, 5);
261        assert!(results.iter().filter(|r| r.node == "A").count() <= 1);
262    }
263
264    #[test]
265    fn in_memory_graph_incremental_construction() {
266        let mut graph = InMemoryGraph::new();
267        graph.add_edge(make_edge("a::foo", "b::bar", EdgeKind::Calls));
268        let results = graph.bfs("a::foo", Direction::Forward, 1);
269        assert_eq!(results.len(), 1);
270        assert_eq!(results[0].node, "b::bar");
271    }
272
273    #[test]
274    fn bfs_backward_finds_callers() {
275        let edges = vec![
276            make_edge("A", "C", EdgeKind::Calls),
277            make_edge("B", "C", EdgeKind::Calls),
278        ];
279        let graph = InMemoryGraph::from_edges(edges);
280        let results = graph.bfs("C", Direction::Backward, 1);
281        assert_eq!(results.len(), 2);
282        let names: Vec<&str> = results.iter().map(|r| r.node.as_str()).collect();
283        assert!(names.contains(&"A"));
284        assert!(names.contains(&"B"));
285    }
286}
the_code_graph_domain/traversal.rs

the_code_graph_domain/
traversal.rs
