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