1use crate::types::{Edge, GraphData, Node};
2use std::collections::{HashMap, HashSet, VecDeque};
3
4pub const DEFAULT_AFFECTED_RELATIONS: &[&str] = &[
5 "calls",
6 "references",
7 "imports",
8 "imports_from",
9 "re_exports",
10 "inherits",
11 "extends",
12 "implements",
13 "uses",
14 "mixes_in",
15 "embeds",
16];
17
18#[derive(Debug, Clone, PartialEq)]
19pub struct AffectedHit {
20 pub node_id: String,
21 pub depth: usize,
22 pub via_relation: String,
23}
24
25pub fn resolve_seed(graph: &GraphData, query: &str) -> Option<String> {
26 let node_ids: HashSet<&str> = graph.nodes.iter().map(|n| n.id.as_str()).collect();
27 if node_ids.contains(query) {
28 return Some(query.to_string());
29 }
30 let query_lower = query.to_lowercase();
31 let exact_label: Vec<&str> = graph
32 .nodes
33 .iter()
34 .filter(|n| n.label.to_lowercase() == query_lower)
35 .map(|n| n.id.as_str())
36 .collect();
37 if exact_label.len() == 1 {
38 return Some(exact_label[0].to_string());
39 }
40 let exact_source: Vec<&str> = graph
41 .nodes
42 .iter()
43 .filter(|n| n.source_file.to_lowercase() == query_lower)
44 .map(|n| n.id.as_str())
45 .collect();
46 if exact_source.len() == 1 {
47 return Some(exact_source[0].to_string());
48 }
49 let contains: Vec<&str> = graph
50 .nodes
51 .iter()
52 .filter(|n| n.label.to_lowercase().contains(&query_lower))
53 .map(|n| n.id.as_str())
54 .collect();
55 if contains.len() == 1 {
56 return Some(contains[0].to_string());
57 }
58 None
59}
60
61pub fn affected_nodes(
62 graph: &GraphData,
63 seed: &str,
64 relations: &[&str],
65 depth: usize,
66) -> Vec<AffectedHit> {
67 let relation_set: HashSet<&str> = relations.iter().copied().collect();
68
69 let mut reverse_adj: HashMap<&str, Vec<(&str, &str)>> = HashMap::new();
70 for edge in &graph.edges {
71 reverse_adj
72 .entry(edge.target.as_str())
73 .or_default()
74 .push((edge.source.as_str(), edge.relation.as_str()));
75 }
76
77 let mut seen: HashSet<&str> = HashSet::new();
78 seen.insert(seed);
79 let mut queue: VecDeque<(&str, usize)> = VecDeque::new();
80 queue.push_back((seed, 0));
81 let mut hits: Vec<AffectedHit> = Vec::new();
82
83 while let Some((current, current_depth)) = queue.pop_front() {
84 if current_depth >= depth {
85 continue;
86 }
87 if let Some(incoming) = reverse_adj.get(current) {
88 for &(source, relation) in incoming {
89 if !relation_set.contains(relation) {
90 continue;
91 }
92 if seen.contains(source) {
93 continue;
94 }
95 seen.insert(source);
96 hits.push(AffectedHit {
97 node_id: source.to_string(),
98 depth: current_depth + 1,
99 via_relation: relation.to_string(),
100 });
101 queue.push_back((source, current_depth + 1));
102 }
103 }
104 }
105 hits
106}
107
108pub fn format_affected(graph: &GraphData, query: &str, relations: &[&str], depth: usize) -> String {
109 let node_map: HashMap<&str, &Node> = graph.nodes.iter().map(|n| (n.id.as_str(), n)).collect();
110
111 let seed = match resolve_seed(graph, query) {
112 Some(s) => s,
113 None => return format!("No unique node match for {}", query),
114 };
115
116 let hits = affected_nodes(graph, &seed, relations, depth);
117 let seed_label = node_map
118 .get(seed.as_str())
119 .map(|n| n.label.as_str())
120 .unwrap_or(&seed);
121
122 let mut lines = vec![
123 format!("Affected nodes for {}", seed_label),
124 format!("Relations: {}", relations.join(", ")),
125 format!("Depth: {}", depth),
126 ];
127
128 if hits.is_empty() {
129 lines.push("No affected nodes found.".to_string());
130 return lines.join("\n");
131 }
132
133 for hit in &hits {
134 let node = node_map.get(hit.node_id.as_str());
135 let label = node.map(|n| n.label.as_str()).unwrap_or(&hit.node_id);
136 let location = node
137 .map(|n| {
138 if let Some(ref loc) = n.source_location {
139 format!("{}:{}", n.source_file, loc)
140 } else {
141 n.source_file.clone()
142 }
143 })
144 .unwrap_or_else(|| "-".to_string());
145 lines.push(format!("- {} [{}] {}", label, hit.via_relation, location));
146 }
147 lines.join("\n")
148}
149
150pub fn load_graph_json(path: &std::path::Path) -> crate::error::Result<GraphData> {
151 let text = std::fs::read_to_string(path)?;
152 let v: serde_json::Value =
153 serde_json::from_str(&text).map_err(crate::error::CodeSynapseError::Serialization)?;
154
155 let nodes: Vec<Node> = v
156 .get("nodes")
157 .and_then(|n| n.as_array())
158 .map(|arr| {
159 arr.iter()
160 .filter_map(|item| serde_json::from_value(item.clone()).ok())
161 .collect()
162 })
163 .unwrap_or_default();
164
165 let edge_arr = v
166 .get("edges")
167 .or_else(|| v.get("links"))
168 .and_then(|e| e.as_array());
169
170 let edges: Vec<Edge> = edge_arr
171 .map(|arr| {
172 arr.iter()
173 .filter_map(|item| {
174 let obj = item.as_object()?;
175 let source = obj.get("source")?.as_str()?.to_string();
176 let target = obj.get("target")?.as_str()?.to_string();
177 let relation = obj
178 .get("relation")
179 .and_then(|v| v.as_str())
180 .unwrap_or("")
181 .to_string();
182 let confidence = obj
183 .get("confidence")
184 .and_then(|v| v.as_str())
185 .unwrap_or("EXTRACTED")
186 .to_string();
187 let source_file = obj
188 .get("source_file")
189 .and_then(|v| v.as_str())
190 .map(|s| s.to_string());
191 let weight = obj.get("weight").and_then(|v| v.as_f64()).unwrap_or(1.0);
192 Some(Edge {
193 source,
194 target,
195 relation,
196 confidence,
197 source_file,
198 weight,
199 context: None,
200 })
201 })
202 .collect()
203 })
204 .unwrap_or_default();
205
206 Ok(GraphData {
207 nodes,
208 edges,
209 hyperedges: None,
210 })
211}
212
213#[cfg(test)]
214mod tests {
215 use super::*;
216 use crate::types::{Edge, Node};
217 use std::collections::HashMap;
218
219 fn make_node(id: &str, label: &str, source_file: &str) -> Node {
220 Node {
221 id: id.to_string(),
222 label: label.to_string(),
223 file_type: "code".to_string(),
224 source_file: source_file.to_string(),
225 source_location: None,
226 community: None,
227 rationale: None,
228 docstring: None,
229 metadata: HashMap::new(),
230 }
231 }
232
233 fn make_edge(source: &str, target: &str, relation: &str) -> Edge {
234 Edge {
235 source: source.to_string(),
236 target: target.to_string(),
237 relation: relation.to_string(),
238 confidence: "EXTRACTED".to_string(),
239 source_file: None,
240 weight: 1.0,
241 context: None,
242 }
243 }
244
245 fn test_graph() -> GraphData {
246 GraphData {
247 nodes: vec![
248 make_node("target", "Foo", "pkg/foo.py"),
249 make_node("caller", "X()", "app.py"),
250 make_node("barrel", "__init__.py", "pkg/__init__.py"),
251 make_node("consumer", "app.py", "app.py"),
252 ],
253 edges: vec![
254 make_edge("caller", "target", "calls"),
255 make_edge("barrel", "target", "re_exports"),
256 make_edge("consumer", "target", "imports"),
257 ],
258 hyperedges: None,
259 }
260 }
261
262 #[test]
263 fn test_resolve_seed_by_id() {
264 let g = test_graph();
265 assert_eq!(resolve_seed(&g, "target"), Some("target".to_string()));
266 }
267
268 #[test]
269 fn test_resolve_seed_by_label() {
270 let g = test_graph();
271 assert_eq!(resolve_seed(&g, "Foo"), Some("target".to_string()));
272 }
273
274 #[test]
275 fn test_resolve_seed_case_insensitive() {
276 let g = test_graph();
277 assert_eq!(resolve_seed(&g, "foo"), Some("target".to_string()));
278 }
279
280 #[test]
281 fn test_resolve_seed_not_found() {
282 let g = test_graph();
283 assert_eq!(resolve_seed(&g, "nonexistent"), None);
284 }
285
286 #[test]
287 fn test_resolve_seed_ambiguous_returns_none() {
288 let g = GraphData {
289 nodes: vec![
290 make_node("n1", "Foo", "a.py"),
291 make_node("n2", "Foo", "b.py"),
292 ],
293 edges: vec![],
294 hyperedges: None,
295 };
296 assert_eq!(resolve_seed(&g, "Foo"), None);
297 }
298
299 #[test]
300 fn test_affected_nodes_all_relations() {
301 let g = test_graph();
302 let hits = affected_nodes(&g, "target", DEFAULT_AFFECTED_RELATIONS, 2);
303 let ids: HashSet<&str> = hits.iter().map(|h| h.node_id.as_str()).collect();
304 assert!(ids.contains("caller"));
305 assert!(ids.contains("barrel"));
306 assert!(ids.contains("consumer"));
307 }
308
309 #[test]
310 fn test_affected_nodes_relation_filter() {
311 let g = test_graph();
312 let hits = affected_nodes(&g, "target", &["calls"], 2);
313 let ids: HashSet<&str> = hits.iter().map(|h| h.node_id.as_str()).collect();
314 assert!(ids.contains("caller"));
315 assert!(!ids.contains("barrel"));
316 assert!(!ids.contains("consumer"));
317 }
318
319 #[test]
320 fn test_affected_nodes_depth_zero_returns_empty() {
321 let g = test_graph();
322 let hits = affected_nodes(&g, "target", DEFAULT_AFFECTED_RELATIONS, 0);
323 assert!(hits.is_empty());
324 }
325
326 #[test]
327 fn test_affected_nodes_depth_propagation() {
328 let g = GraphData {
329 nodes: vec![
330 make_node("a", "A", "a.py"),
331 make_node("b", "B", "b.py"),
332 make_node("c", "C", "c.py"),
333 ],
334 edges: vec![make_edge("b", "a", "calls"), make_edge("c", "b", "calls")],
335 hyperedges: None,
336 };
337 let hits_d1 = affected_nodes(&g, "a", &["calls"], 1);
338 let hits_d2 = affected_nodes(&g, "a", &["calls"], 2);
339 assert_eq!(hits_d1.len(), 1);
340 assert_eq!(hits_d2.len(), 2);
341 }
342
343 #[test]
344 fn test_format_affected_contains_expected_output() {
345 let g = test_graph();
346 let out = format_affected(&g, "Foo", DEFAULT_AFFECTED_RELATIONS, 2);
347 assert!(
348 out.contains("Affected nodes for Foo"),
349 "missing header: {out}"
350 );
351 assert!(out.contains("X()"), "missing X(): {out}");
352 assert!(out.contains("calls"), "missing calls: {out}");
353 assert!(out.contains("__init__.py"), "missing __init__: {out}");
354 assert!(out.contains("re_exports"), "missing re_exports: {out}");
355 assert!(out.contains("imports"), "missing imports: {out}");
356 }
357
358 #[test]
359 fn test_format_affected_relation_filter() {
360 let g = test_graph();
361 let out = format_affected(&g, "Foo", &["calls"], 2);
362 assert!(
363 out.contains("Relations: calls"),
364 "missing relations header: {out}"
365 );
366 assert!(out.contains("X()"), "missing X(): {out}");
367 assert!(!out.contains("__init__.py"), "unexpected __init__: {out}");
368 }
369
370 #[test]
371 fn test_format_affected_not_found() {
372 let g = test_graph();
373 let out = format_affected(&g, "nonexistent", DEFAULT_AFFECTED_RELATIONS, 2);
374 assert!(out.contains("No unique node match"));
375 }
376
377 #[test]
378 fn test_format_affected_no_hits() {
379 let g = GraphData {
380 nodes: vec![make_node("isolated", "Isolated", "x.py")],
381 edges: vec![],
382 hyperedges: None,
383 };
384 let out = format_affected(&g, "isolated", DEFAULT_AFFECTED_RELATIONS, 2);
385 assert!(out.contains("No affected nodes found."));
386 }
387}