1use crate::edge::{Edge, EdgeKind};
7use crate::graph::{ArborGraph, NodeId};
8use crate::symbol_table::SymbolTable;
9use arbor_core::CodeNode;
10use std::collections::HashMap;
11use std::path::PathBuf;
12use tracing::warn;
13
14pub struct GraphBuilder {
20 graph: ArborGraph,
21 symbol_table: SymbolTable,
23 name_to_id: HashMap<String, String>,
25}
26
27impl Default for GraphBuilder {
28 fn default() -> Self {
29 Self::new()
30 }
31}
32
33impl GraphBuilder {
34 pub fn new() -> Self {
36 Self {
37 graph: ArborGraph::new(),
38 symbol_table: SymbolTable::new(),
39 name_to_id: HashMap::new(),
40 }
41 }
42
43 pub fn add_nodes(&mut self, nodes: Vec<CodeNode>) {
48 for node in nodes {
49 let id_str = node.id.clone();
50 let name = node.name.clone();
51 let qualified = node.qualified_name.clone();
52 let file = PathBuf::from(&node.file);
53
54 let node_idx = self.graph.add_node(node);
55
56 if !qualified.is_empty() {
58 self.symbol_table
59 .insert(qualified.clone(), node_idx, file.clone());
60 }
61
62 self.name_to_id.insert(name.clone(), id_str.clone());
63 self.name_to_id.insert(qualified, id_str);
64 }
65 }
66
67 pub fn resolve_edges(&mut self) {
72 let mut edges_to_add = Vec::new();
74
75 let node_indices: Vec<NodeId> = self.graph.node_indexes().collect();
77
78 for from_idx in node_indices {
79 let (references, from_file) = {
81 let node = self.graph.get(from_idx).unwrap();
82 (node.references.clone(), PathBuf::from(&node.file))
83 };
84
85 for reference in references {
86 if let Some(to_idx) = self.symbol_table.resolve(&reference) {
88 if from_idx != to_idx {
89 edges_to_add.push((from_idx, to_idx, reference.clone()));
90 }
91 continue;
92 }
93
94 if let Some(to_idx) = self
96 .symbol_table
97 .resolve_with_context(&reference, &from_file)
98 {
99 if from_idx != to_idx {
100 edges_to_add.push((from_idx, to_idx, reference.clone()));
101 }
102 continue;
103 }
104
105 warn!(
107 "Unresolved reference '{}' in {}",
108 reference,
109 from_file.display()
110 );
111 }
112 }
113
114 for (from_id, to_id, _ref_name) in edges_to_add {
116 self.graph
117 .add_edge(from_id, to_id, Edge::new(EdgeKind::Calls));
118 }
119 }
120
121 pub fn build(mut self) -> ArborGraph {
123 self.resolve_edges();
124 self.graph
125 }
126
127 pub fn build_without_resolve(self) -> ArborGraph {
129 self.graph
130 }
131}
132
133#[cfg(test)]
134mod tests {
135 use super::*;
136 use arbor_core::NodeKind;
137
138 #[test]
139 fn test_builder_adds_nodes() {
140 let mut builder = GraphBuilder::new();
141
142 let node1 = CodeNode::new("foo", "foo", NodeKind::Function, "test.rs");
143 let node2 = CodeNode::new("bar", "bar", NodeKind::Function, "test.rs");
144
145 builder.add_nodes(vec![node1, node2]);
146 let graph = builder.build();
147
148 assert_eq!(graph.node_count(), 2);
149 }
150
151 #[test]
152 fn test_builder_resolves_edges() {
153 let mut builder = GraphBuilder::new();
154
155 let caller = CodeNode::new("caller", "caller", NodeKind::Function, "test.rs")
156 .with_references(vec!["callee".to_string()]);
157 let callee = CodeNode::new("callee", "callee", NodeKind::Function, "test.rs");
158
159 builder.add_nodes(vec![caller, callee]);
160 let graph = builder.build();
161
162 assert_eq!(graph.node_count(), 2);
163 assert_eq!(graph.edge_count(), 1);
164 }
165
166 #[test]
167 fn test_cross_file_resolution() {
168 let mut builder = GraphBuilder::new();
169
170 let caller = CodeNode::new("main", "main", NodeKind::Function, "main.rs")
172 .with_references(vec!["pkg.Utils.helper".to_string()]);
173
174 let mut callee = CodeNode::new("helper", "helper", NodeKind::Method, "utils.rs");
176 callee.qualified_name = "pkg.Utils.helper".to_string();
177
178 builder.add_nodes(vec![caller]);
179 builder.add_nodes(vec![callee]);
180
181 let graph = builder.build();
182
183 assert_eq!(graph.node_count(), 2);
184 assert_eq!(
185 graph.edge_count(),
186 1,
187 "Should resolve cross-file edge via FQN"
188 );
189 }
190}