Skip to main content

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;
12use tracing::warn;
13
14/// Builds an ArborGraph from parsed code nodes.
15///
16/// The builder handles the two-pass process:
17/// 1. Add all nodes to the graph
18/// 2. Resolve references into edges (including cross-file)
19pub struct GraphBuilder {
20    graph: ArborGraph,
21    /// Maps qualified names to node IDs for edge resolution.
22    symbol_table: SymbolTable,
23    /// Legacy map for simple name resolution (within same file)
24    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    /// Creates a new builder.
35    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    /// Adds nodes from a file to the graph.
44    ///
45    /// Call this for each parsed file, then call `resolve_edges`
46    /// when all files are added.
47    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            // Populate Symbol Table
57            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    /// Resolves references into actual graph edges.
68    ///
69    /// This is the second pass after all nodes are added. It looks up
70    /// reference names and creates edges where targets exist.
71    pub fn resolve_edges(&mut self) {
72        // Collect all the edge additions first to avoid borrow issues
73        let mut edges_to_add = Vec::new();
74
75        // Collect indices to avoid borrowing self.graph during iteration
76        let node_indices: Vec<NodeId> = self.graph.node_indexes().collect();
77
78        for from_idx in node_indices {
79            // Get references and file by cloning to release borrow on graph
80            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                // 1. Try exact FQN match
87                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                // 2. Try context-aware resolution (suffix match with locality)
95                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                // 3. Unresolved - warn
106                warn!(
107                    "Unresolved reference '{}' in {}",
108                    reference,
109                    from_file.display()
110                );
111            }
112        }
113
114        // Now add the edges
115        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    /// Finishes building and returns the graph.
122    pub fn build(mut self) -> ArborGraph {
123        self.resolve_edges();
124        self.graph
125    }
126
127    /// Builds without resolving edges (for incremental updates).
128    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        // File A: Calls "pkg.Utils.helper"
171        let caller = CodeNode::new("main", "main", NodeKind::Function, "main.rs")
172            .with_references(vec!["pkg.Utils.helper".to_string()]);
173
174        // File B: Defines "pkg.Utils.helper"
175        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}