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