arbor_graph/
builder.rs

1//! Graph builder for constructing the code graph from parsed nodes.
2//!
3//! The builder takes CodeNodes and resolves their references into
4//! actual graph edges.
5
6use 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
13/// Builds an ArborGraph from parsed code nodes.
14///
15/// The builder handles the two-pass process:
16/// 1. Add all nodes to the graph
17/// 2. Resolve references into edges (including cross-file)
18pub struct GraphBuilder {
19    graph: ArborGraph,
20    /// Maps qualified names to node IDs for edge resolution.
21    symbol_table: SymbolTable,
22    /// Legacy map for simple name resolution (within same file)
23    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    /// Creates a new builder.
34    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    /// Adds nodes from a file to the graph.
43    ///
44    /// Call this for each parsed file, then call `resolve_edges`
45    /// when all files are added.
46    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            // Populate Symbol Table
56            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    /// Resolves references into actual graph edges.
67    ///
68    /// This is the second pass after all nodes are added. It looks up
69    /// reference names and creates edges where targets exist.
70    pub fn resolve_edges(&mut self) {
71        // Collect all the edge additions first to avoid borrow issues
72        let mut edges_to_add = Vec::new();
73
74        // Collect indices to avoid borrowing self.graph during iteration
75        let node_indices: Vec<NodeId> = self.graph.node_indexes().collect();
76
77        for from_idx in node_indices {
78            // Get references by cloning to release borrow on graph
79            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                // 1. Try resolving via Symbol Table (FQN)
88                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                // 2. Fallback to legacy ID map
100                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        // Now add the edges
111        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    /// Finishes building and returns the graph.
118    pub fn build(mut self) -> ArborGraph {
119        self.resolve_edges();
120        self.graph
121    }
122
123    /// Builds without resolving edges (for incremental updates).
124    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        // File A: Calls "pkg.Utils.helper"
167        let caller = CodeNode::new("main", "main", NodeKind::Function, "main.rs")
168            .with_references(vec!["pkg.Utils.helper".to_string()]);
169
170        // File B: Defines "pkg.Utils.helper"
171        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}