1use crate::edge::{Edge, EdgeKind};
7use crate::graph::ArborGraph;
8use arbor_core::CodeNode;
9use std::collections::HashMap;
10
11pub struct GraphBuilder {
17 graph: ArborGraph,
18 name_to_id: HashMap<String, String>,
20}
21
22impl Default for GraphBuilder {
23 fn default() -> Self {
24 Self::new()
25 }
26}
27
28impl GraphBuilder {
29 pub fn new() -> Self {
31 Self {
32 graph: ArborGraph::new(),
33 name_to_id: HashMap::new(),
34 }
35 }
36
37 pub fn add_nodes(&mut self, nodes: Vec<CodeNode>) {
42 for node in nodes {
43 let id = node.id.clone();
44 let name = node.name.clone();
45 let qualified = node.qualified_name.clone();
46
47 self.graph.add_node(node);
48
49 self.name_to_id.insert(name.clone(), id.clone());
51 self.name_to_id.insert(qualified, id);
52 }
53 }
54
55 pub fn resolve_edges(&mut self) {
60 let mut edges_to_add = Vec::new();
62
63 for node in self.graph.nodes() {
64 let from_id = &node.id;
65
66 for reference in &node.references {
67 if let Some(to_id) = self.name_to_id.get(reference) {
69 if from_id != to_id {
70 edges_to_add.push((from_id.clone(), to_id.clone(), reference.clone()));
71 }
72 }
73 }
74 }
75
76 for (from_id, to_id, _ref_name) in edges_to_add {
78 if let (Some(from_idx), Some(to_idx)) =
79 (self.graph.get_index(&from_id), self.graph.get_index(&to_id))
80 {
81 self.graph
82 .add_edge(from_idx, to_idx, Edge::new(EdgeKind::Calls));
83 }
84 }
85 }
86
87 pub fn build(mut self) -> ArborGraph {
89 self.resolve_edges();
90 self.graph
91 }
92
93 pub fn build_without_resolve(self) -> ArborGraph {
95 self.graph
96 }
97}
98
99#[cfg(test)]
100mod tests {
101 use super::*;
102 use arbor_core::NodeKind;
103
104 #[test]
105 fn test_builder_adds_nodes() {
106 let mut builder = GraphBuilder::new();
107
108 let node1 = CodeNode::new("foo", "foo", NodeKind::Function, "test.rs");
109 let node2 = CodeNode::new("bar", "bar", NodeKind::Function, "test.rs");
110
111 builder.add_nodes(vec![node1, node2]);
112 let graph = builder.build();
113
114 assert_eq!(graph.node_count(), 2);
115 }
116
117 #[test]
118 fn test_builder_resolves_edges() {
119 let mut builder = GraphBuilder::new();
120
121 let caller = CodeNode::new("caller", "caller", NodeKind::Function, "test.rs")
122 .with_references(vec!["callee".to_string()]);
123 let callee = CodeNode::new("callee", "callee", NodeKind::Function, "test.rs");
124
125 builder.add_nodes(vec![caller, callee]);
126 let graph = builder.build();
127
128 assert_eq!(graph.node_count(), 2);
129 assert_eq!(graph.edge_count(), 1);
130 }
131}