Skip to main content

aptu_coder_core/graph/
structural.rs

1// SPDX-FileCopyrightText: 2026 aptu-coder contributors
2// SPDX-License-Identifier: Apache-2.0
3//! Structural knowledge graph over petgraph DiGraph with BFS blast-radius traversal.
4
5use 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        // Pass 1: Add all File, Symbol, and Module nodes; populate symbol_index.
76        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        // Pass 2: Resolve call edges against the now-complete symbol_index.
117        for entry in entries {
118            for cl in &entry.semantic.calls {
119                // Name-only keying preserves the prior first-definition-wins semantics: the first node inserted for a given name wins.
120                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    /// BFS traversal returning both the visited set (including start) and the tail
137    /// (neighbors discovered, excluding start).
138    ///
139    /// The visited set contains all nodes reached up to the specified depth.
140    /// The tail is the BFS-order sequence of nodes discovered, not including start.
141    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    /// Blast-radius subgraph including both nodes and edges.
172    ///
173    /// Returns a tuple of (nodes, edges) where:
174    /// - nodes: Vec<NodeIndex> with the start symbol first, followed by all discovered nodes in BFS order
175    /// - edges: Vec<(NodeIndex, NodeIndex, Edge)> containing every edge whose source and target
176    ///   are both in the visited set (not just edges walked by the BFS tree), allowing clients
177    ///   to fully reconstruct the subgraph's connectivity
178    ///
179    /// If the symbol is not found, returns (vec![], vec![]).
180    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        // Build node list: start first, then all discovered nodes in BFS order
192        let mut nodes = vec![start];
193        nodes.extend(tail);
194
195        // Collect all edges whose both source and target are in the visited set
196        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    /// Two files with the same call edge produce exactly 1 Calls edge because the
296    /// HashMap index resolves both callers to the same first-definition node, and
297    /// the `seen` HashSet deduplicates the edge.
298    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}