1use crate::analyze::FileAnalysisOutput;
6use petgraph::graph::{DiGraph, NodeIndex};
7use petgraph::visit::EdgeRef;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use std::collections::HashSet;
11
12#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
13pub enum SymbolKind {
14 Function,
15 Class,
16}
17
18#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
19pub enum Node {
20 File {
21 path: String,
22 },
23 Symbol {
24 name: String,
25 kind: SymbolKind,
26 file_path: String,
27 },
28 Module {
29 path: String,
30 },
31}
32
33#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
34pub enum Edge {
35 Contains,
36 Calls,
37 Imports,
38}
39
40#[derive(Debug, Clone, Serialize, Deserialize)]
41pub struct StructuralGraph {
42 pub graph: DiGraph<Node, Edge>,
43 #[serde(skip)]
44 symbol_index: HashMap<String, NodeIndex>,
45}
46
47impl StructuralGraph {
48 fn build_symbol_index(graph: &DiGraph<Node, Edge>) -> HashMap<String, NodeIndex> {
49 let mut index = HashMap::new();
50 for idx in graph.node_indices() {
51 if let Node::Symbol { name, .. } = &graph[idx] {
52 index.entry(name.clone()).or_insert(idx);
53 }
54 }
55 index
56 }
57
58 pub fn from_graph(graph: DiGraph<Node, Edge>) -> Self {
59 let symbol_index = Self::build_symbol_index(&graph);
60 StructuralGraph {
61 graph,
62 symbol_index,
63 }
64 }
65
66 pub(crate) fn rebuild_symbol_index(&mut self) {
67 self.symbol_index = Self::build_symbol_index(&self.graph);
68 }
69
70 pub fn build_from_analysis(entries: &[FileAnalysisOutput]) -> Self {
71 let mut graph = DiGraph::new();
72 let mut seen: HashSet<(NodeIndex, NodeIndex)> = HashSet::new();
73 let mut symbol_index: HashMap<String, NodeIndex> = HashMap::new();
74
75 for entry in entries {
77 let fp = entry.formatted.lines().next().unwrap_or("");
78 let file = graph.add_node(Node::File {
79 path: fp.to_string(),
80 });
81
82 for f in &entry.semantic.functions {
83 let n = graph.add_node(Node::Symbol {
84 name: f.name.clone(),
85 kind: SymbolKind::Function,
86 file_path: fp.to_string(),
87 });
88 if seen.insert((file, n)) {
89 graph.add_edge(file, n, Edge::Contains);
90 }
91 symbol_index.entry(f.name.clone()).or_insert(n);
92 }
93 for c in &entry.semantic.classes {
94 let n = graph.add_node(Node::Symbol {
95 name: c.name.clone(),
96 kind: SymbolKind::Class,
97 file_path: fp.to_string(),
98 });
99 if seen.insert((file, n)) {
100 graph.add_edge(file, n, Edge::Contains);
101 }
102 symbol_index.entry(c.name.clone()).or_insert(n);
103 }
104 for im in &entry.semantic.imports {
105 if !im.module.is_empty() {
106 let n = graph.add_node(Node::Module {
107 path: im.module.clone(),
108 });
109 if seen.insert((file, n)) {
110 graph.add_edge(file, n, Edge::Imports);
111 }
112 }
113 }
114 }
115
116 for entry in entries {
118 for cl in &entry.semantic.calls {
119 let caller = symbol_index.get(&cl.caller).copied();
121 let callee = symbol_index.get(&cl.callee).copied();
122 if let (Some(c), Some(e)) = (caller, callee)
123 && seen.insert((c, e))
124 {
125 graph.add_edge(c, e, Edge::Calls);
126 }
127 }
128 }
129
130 StructuralGraph {
131 graph,
132 symbol_index,
133 }
134 }
135
136 fn bfs_frontier(&self, start: NodeIndex, depth: usize) -> (HashSet<NodeIndex>, Vec<NodeIndex>) {
142 let mut visited = HashSet::new();
143 let mut result = Vec::new();
144 let mut frontier = vec![start];
145 visited.insert(start);
146 for _ in 0..depth {
147 if frontier.is_empty() {
148 break;
149 }
150 let mut next = Vec::new();
151 for node in frontier {
152 for nb in self.graph.neighbors(node) {
153 if visited.insert(nb) {
154 result.push(nb);
155 next.push(nb);
156 }
157 }
158 }
159 frontier = next;
160 }
161 (visited, result)
162 }
163
164 pub fn bfs_blast_radius(&self, symbol: &str, depth: usize) -> Vec<NodeIndex> {
165 let Some(start) = self.symbol_index.get(symbol).copied() else {
166 return vec![];
167 };
168 self.bfs_frontier(start, depth).1
169 }
170
171 pub fn blast_radius_subgraph(
181 &self,
182 symbol: &str,
183 depth: usize,
184 ) -> (Vec<NodeIndex>, Vec<(NodeIndex, NodeIndex, Edge)>) {
185 let Some(start) = self.symbol_index.get(symbol).copied() else {
186 return (vec![], vec![]);
187 };
188
189 let (visited, tail) = self.bfs_frontier(start, depth);
190
191 let mut nodes = vec![start];
193 nodes.extend(tail);
194
195 let edges: Vec<(NodeIndex, NodeIndex, Edge)> = self
197 .graph
198 .edge_references()
199 .filter(|e| visited.contains(&e.source()) && visited.contains(&e.target()))
200 .map(|e| (e.source(), e.target(), e.weight().clone()))
201 .collect();
202
203 (nodes, edges)
204 }
205}
206
207#[cfg(test)]
208mod tests {
209 use super::*;
210 use crate::types::{CallInfo, ClassInfo, FunctionInfo, ImportInfo, SemanticAnalysis};
211
212 fn make_output(
213 path: &str,
214 funcs: Vec<&str>,
215 classes: Vec<&str>,
216 imports: Vec<&str>,
217 calls: Vec<(&str, &str)>,
218 ) -> FileAnalysisOutput {
219 FileAnalysisOutput::new(
220 format!("{}:1:1:1", path),
221 SemanticAnalysis {
222 functions: funcs
223 .into_iter()
224 .map(|n| FunctionInfo {
225 name: n.to_string(),
226 line: 1,
227 end_line: 6,
228 parameters: vec![],
229 return_type: None,
230 })
231 .collect(),
232 classes: classes
233 .into_iter()
234 .map(|n| ClassInfo {
235 name: n.to_string(),
236 line: 1,
237 end_line: 10,
238 methods: vec![],
239 fields: vec![],
240 inherits: vec![],
241 })
242 .collect(),
243 imports: imports
244 .into_iter()
245 .map(|m| ImportInfo {
246 module: m.to_string(),
247 items: vec![],
248 line: 1,
249 })
250 .collect(),
251 references: vec![],
252 call_frequency: Default::default(),
253 calls: calls
254 .into_iter()
255 .map(|(c, e)| CallInfo {
256 caller: c.to_string(),
257 callee: e.to_string(),
258 line: 1,
259 column: 0,
260 arg_count: None,
261 })
262 .collect(),
263 impl_traits: vec![],
264 def_use_sites: vec![],
265 },
266 10,
267 None,
268 )
269 }
270
271 #[test]
272 fn test_build_happy_path() {
273 let e = make_output(
274 "src/main.rs",
275 vec!["main", "helper"],
276 vec!["Config"],
277 vec!["std::collections"],
278 vec![("main", "helper")],
279 );
280 let g = StructuralGraph::build_from_analysis(&[e]);
281 assert!(g.graph.node_count() >= 4, "nodes={}", g.graph.node_count());
282 assert!(g.graph.edge_count() >= 5, "edges={}", g.graph.edge_count());
283 assert!(g.graph.edge_indices().any(|i| g.graph[i] == Edge::Calls));
284 }
285
286 #[test]
287 fn test_build_empty_input() {
288 let e = make_output("src/e.rs", vec![], vec![], vec![], vec![]);
289 let g = StructuralGraph::build_from_analysis(&[e]);
290 assert_eq!(g.graph.node_count(), 1);
291 assert_eq!(g.graph.edge_count(), 0);
292 }
293
294 #[test]
295 fn test_build_dedup_edges() {
299 let e1 = make_output(
300 "src/a.rs",
301 vec!["main", "helper"],
302 vec![],
303 vec![],
304 vec![("main", "helper")],
305 );
306 let e2 = make_output(
307 "src/b.rs",
308 vec!["main", "helper"],
309 vec![],
310 vec![],
311 vec![("main", "helper")],
312 );
313 let g = StructuralGraph::build_from_analysis(&[e1, e2]);
314 let n = g
315 .graph
316 .edge_indices()
317 .filter(|i| g.graph[*i] == Edge::Calls)
318 .count();
319 assert_eq!(
320 n, 1,
321 "expected 1 Calls edge (first-definition-wins), got {}",
322 n
323 );
324 }
325
326 #[test]
327 fn test_bfs_diamond() {
328 let mut g = DiGraph::new();
329 let mut sym = |n: &str| {
330 g.add_node(Node::Symbol {
331 name: n.into(),
332 kind: SymbolKind::Function,
333 file_path: "t.rs".into(),
334 })
335 };
336 let a = sym("A");
337 let b = sym("B");
338 let c = sym("C");
339 let d = sym("D");
340 g.add_edge(a, b, Edge::Calls);
341 g.add_edge(a, c, Edge::Calls);
342 g.add_edge(b, d, Edge::Calls);
343 g.add_edge(c, d, Edge::Calls);
344 let graph = StructuralGraph::from_graph(g);
345 let r = graph.bfs_blast_radius("A", 2);
346 assert_eq!(r.len(), 3, "expected 3 nodes, got {:?}", r);
347 }
348
349 #[test]
350 fn test_bfs_symbol_not_found() {
351 let graph = StructuralGraph::from_graph(DiGraph::new());
352 assert!(graph.bfs_blast_radius("x", 3).is_empty());
353 }
354}