use std::collections::HashMap;
use std::path::PathBuf;
use petgraph::graph::{DiGraph, EdgeIndex, NodeIndex};
use petgraph::visit::EdgeRef;
use crate::graph::CodeGraph;
use crate::graph::edge::{EdgeData, EdgeKind};
#[cfg(feature = "dataflow")]
use crate::graph::node::DataGraphNode;
use crate::graph::node::{ExternalNode, FileNode, NodeData, SymbolNode};
use crate::language::LangId;
#[cfg(feature = "dataflow")]
use crate::model::DataNodeId;
use crate::model::{FileId, IdGenerator, SnapshotId, Symbol, SymbolId};
#[derive(Debug)]
pub struct GraphBuilder {
graph: DiGraph<NodeData, EdgeData>,
file_to_index: HashMap<FileId, NodeIndex>,
symbol_to_index: HashMap<SymbolId, NodeIndex>,
path_to_file: HashMap<PathBuf, FileId>,
file_id_gen: IdGenerator<FileId>,
snapshot_id: SnapshotId,
edge_index: HashMap<(NodeIndex, NodeIndex, EdgeKind), EdgeIndex>,
external_index: HashMap<String, NodeIndex>,
#[cfg(feature = "dataflow")]
data_to_index: HashMap<DataNodeId, NodeIndex>,
}
impl GraphBuilder {
pub fn new(snapshot_id: SnapshotId) -> Self {
Self {
graph: DiGraph::new(),
file_to_index: HashMap::new(),
symbol_to_index: HashMap::new(),
path_to_file: HashMap::new(),
file_id_gen: IdGenerator::new(),
snapshot_id,
edge_index: HashMap::new(),
external_index: HashMap::new(),
#[cfg(feature = "dataflow")]
data_to_index: HashMap::new(),
}
}
pub fn add_file(&mut self, path: PathBuf, language: LangId) -> FileId {
if let Some(&existing_id) = self.path_to_file.get(&path) {
return existing_id;
}
let id = self.file_id_gen.next();
let node = FileNode {
id,
path: path.clone(),
language,
snapshot_id: self.snapshot_id,
};
let idx = self.graph.add_node(NodeData::File(node));
self.file_to_index.insert(id, idx);
self.path_to_file.insert(path, id);
id
}
pub fn add_symbol(&mut self, symbol: &Symbol) -> Result<NodeIndex, crate::Error> {
if let Some(&existing) = self.symbol_to_index.get(&symbol.id) {
return Ok(existing);
}
let file_id = *self
.path_to_file
.get(&symbol.file_path)
.ok_or_else(|| crate::Error::Graph("file must be added before its symbols".into()))?;
let file_idx = *self
.file_to_index
.get(&file_id)
.ok_or_else(|| crate::Error::Graph("file index must exist".into()))?;
let node = SymbolNode {
id: symbol.id,
name: symbol.name.clone(),
kind: symbol.kind,
file_id,
visibility: symbol.visibility,
source_range: symbol.source_range.clone(),
};
let sym_idx = self.graph.add_node(NodeData::Symbol(node));
self.symbol_to_index.insert(symbol.id, sym_idx);
self.add_edge_internal(file_idx, sym_idx, EdgeKind::Ownership, 1.0);
Ok(sym_idx)
}
#[cfg(feature = "dataflow")]
pub fn add_data_node(&mut self, data_node: &crate::model::DataNode) -> NodeIndex {
if let Some(&existing) = self.data_to_index.get(&data_node.id) {
return existing;
}
let node = DataGraphNode {
id: data_node.id,
symbol_id: data_node.symbol_id,
name: data_node.name.clone(),
scope: data_node.scope,
type_hint: data_node.type_hint.clone(),
source_range: data_node.source_range.clone(),
};
let idx = self.graph.add_node(NodeData::Data(node));
self.data_to_index.insert(data_node.id, idx);
if let Some(symbol_id) = data_node.symbol_id
&& let Some(&sym_idx) = self.symbol_to_index.get(&symbol_id)
{
self.add_edge_internal(sym_idx, idx, EdgeKind::Ownership, 1.0);
}
idx
}
#[cfg(feature = "dataflow")]
pub fn add_flow_edge(
&mut self,
source: crate::model::DataNodeId,
target: crate::model::DataNodeId,
kind: crate::model::FlowKind,
confidence: f32,
) {
let Some(&src_idx) = self.data_to_index.get(&source) else {
return;
};
let Some(&tgt_idx) = self.data_to_index.get(&target) else {
return;
};
self.add_edge_internal_with_flow(src_idx, tgt_idx, EdgeKind::Flow, confidence, Some(kind));
}
pub fn add_import(&mut self, from: FileId, to: PathBuf) {
let Some(&from_idx) = self.file_to_index.get(&from) else {
return; };
if let Some(&to_id) = self.path_to_file.get(&to) {
if let Some(&to_idx) = self.file_to_index.get(&to_id) {
self.add_edge_internal(from_idx, to_idx, EdgeKind::Import, 1.0);
}
return;
}
let raw_path = to.to_string_lossy().to_string();
let to_idx = if let Some(&idx) = self.external_index.get(&raw_path) {
idx
} else {
let language = self
.path_to_file
.iter()
.find(|(_, id)| **id == from)
.map(|(p, _)| {
crate::input::detect_language(p).unwrap_or(crate::language::LangId::Python)
})
.unwrap_or(crate::language::LangId::Python);
let node = ExternalNode {
raw_path: raw_path.clone(),
language,
classification: None,
};
let idx = self.graph.add_node(NodeData::External(node));
self.external_index.insert(raw_path, idx);
idx
};
self.add_edge_internal(from_idx, to_idx, EdgeKind::Import, 1.0);
}
#[cfg(feature = "dataflow")]
fn add_edge_internal_with_flow(
&mut self,
source: NodeIndex,
target: NodeIndex,
kind: EdgeKind,
confidence: f32,
flow_kind: Option<crate::model::FlowKind>,
) {
let confidence = confidence.clamp(0.0, 1.0);
let key = (source, target, kind);
if let Some(&edge_idx) = self.edge_index.get(&key) {
let edge = &mut self.graph[edge_idx];
edge.confidence = edge.confidence.max(confidence);
if edge.flow_kind.is_none() {
edge.flow_kind = flow_kind;
}
return;
}
let edge_idx = self.graph.add_edge(
source,
target,
EdgeData {
kind,
confidence,
flow_kind,
},
);
self.edge_index.insert(key, edge_idx);
}
pub fn add_reference(&mut self, from: SymbolId, to: SymbolId, confidence: f32) {
let Some(&from_idx) = self.symbol_to_index.get(&from) else {
return;
};
let Some(&to_idx) = self.symbol_to_index.get(&to) else {
return;
};
self.add_edge_internal(from_idx, to_idx, EdgeKind::Reference, confidence);
}
fn add_edge_internal(
&mut self,
source: NodeIndex,
target: NodeIndex,
kind: EdgeKind,
confidence: f32,
) {
let confidence = confidence.clamp(0.0, 1.0);
let key = (source, target, kind);
if let Some(&edge_idx) = self.edge_index.get(&key) {
let existing = &mut self.graph[edge_idx];
existing.confidence = existing.confidence.max(confidence);
return;
}
let edge_data = EdgeData::with_confidence(kind, confidence);
let edge_idx = self.graph.add_edge(source, target, edge_data);
self.edge_index.insert(key, edge_idx);
}
pub fn file_id_for_path(&self, path: &std::path::PathBuf) -> Option<FileId> {
self.path_to_file.get(path).copied()
}
pub fn import_adjacency(&self) -> HashMap<FileId, Vec<FileId>> {
let mut adjacency: HashMap<FileId, Vec<FileId>> = HashMap::new();
let mut index_to_file: HashMap<NodeIndex, FileId> = HashMap::new();
for (&file_id, &idx) in &self.file_to_index {
index_to_file.insert(idx, file_id);
}
for edge in self.graph.edge_references() {
if edge.weight().kind == EdgeKind::Import
&& let (Some(&from_id), Some(&to_id)) = (
index_to_file.get(&edge.source()),
index_to_file.get(&edge.target()),
)
{
adjacency.entry(from_id).or_default().push(to_id);
}
}
adjacency
}
pub fn build(self) -> CodeGraph {
CodeGraph::from_parts(
self.graph,
self.edge_index,
self.file_to_index,
self.symbol_to_index,
self.external_index,
self.snapshot_id,
)
}
pub fn node_count(&self) -> usize {
self.graph.node_count()
}
pub fn edge_count(&self) -> usize {
self.graph.edge_count()
}
pub fn external_count(&self) -> usize {
self.external_index.len()
}
pub fn from_extractions<F>(
extractions: &[F],
root: &std::path::Path,
snapshot_id: crate::model::SnapshotId,
diagnostics: &mut Vec<crate::error::Diagnostic>,
) -> (CodeGraph, crate::graph::SccAnalysis)
where
F: std::borrow::Borrow<crate::model::FileExtraction> + Sync,
{
let (graph, scc, _) =
Self::from_extractions_with_scope(extractions, root, snapshot_id, diagnostics);
(graph, scc)
}
pub fn from_extractions_with_scope<F>(
extractions: &[F],
root: &std::path::Path,
snapshot_id: crate::model::SnapshotId,
diagnostics: &mut Vec<crate::error::Diagnostic>,
) -> (
CodeGraph,
crate::graph::SccAnalysis,
crate::graph::resolver::FlattenedScopeCache,
)
where
F: std::borrow::Borrow<crate::model::FileExtraction> + Sync,
{
let mut builder = Self::new(snapshot_id);
for file in extractions {
let file = file.borrow();
builder.add_file(file.path.clone(), file.lang);
}
for file in extractions {
let file = file.borrow();
for symbol in &file.symbols {
if let Err(e) = builder.add_symbol(symbol) {
diagnostics.push(crate::error::Diagnostic {
path: file.path.clone(),
severity: crate::error::Severity::Warning,
message: format!("failed to add symbol to graph: {e}"),
source_range: None,
});
}
}
}
#[cfg(feature = "dataflow")]
{
for file in extractions {
let file = file.borrow();
for data_node in &file.data_nodes {
builder.add_data_node(data_node);
}
}
for file in extractions {
let file = file.borrow();
for flow_edge in &file.flow_edges {
builder.add_flow_edge(
flow_edge.source,
flow_edge.target,
flow_edge.kind,
flow_edge.confidence,
);
}
}
}
let path_to_file_id: HashMap<std::path::PathBuf, crate::model::FileId> = extractions
.iter()
.filter_map(|f| {
let f = f.borrow();
builder
.file_id_for_path(&f.path)
.map(|fid| (f.path.clone(), fid))
})
.collect();
let mut resolvers = HashMap::new();
for lang in crate::language::LangId::all() {
resolvers.insert(lang, crate::language::import_resolver::make_resolver(lang));
}
for file in extractions {
let file = file.borrow();
let Some(&source_fid) = path_to_file_id.get(&file.path) else {
continue;
};
let source_dir = file.path.parent().unwrap_or(std::path::Path::new("."));
if let Some(resolver) = resolvers.get(&file.lang) {
for import in &file.imports {
if let Some(target) =
resolver.resolve(&import.import_specifier, source_dir, root)
{
builder.add_import(source_fid, target);
}
}
}
}
let import_adjacency = builder.import_adjacency();
let ctx = crate::graph::resolver::ResolutionContext::from_extractions(
extractions,
&path_to_file_id,
import_adjacency,
);
let scope_cache = crate::graph::resolver::FlattenedScopeCache::build(&ctx, diagnostics);
let ref_edges = crate::graph::resolver::resolve_all_references(
extractions,
&path_to_file_id,
&scope_cache,
diagnostics,
);
for (from, to, confidence) in ref_edges {
builder.add_reference(from, to, confidence);
}
let graph = builder.build();
#[cfg(feature = "metacall-deploy")]
let mut graph = graph;
#[cfg(feature = "metacall-deploy")]
{
let call_sites: Vec<crate::deploy::scanner::CallSite> = extractions
.iter()
.flat_map(|file| file.borrow().call_sites.iter().cloned())
.collect();
if !call_sites.is_empty() {
let (call_edges, call_diagnostics) =
crate::deploy::client_call::resolve_client_call_edges(
&graph,
extractions,
&call_sites,
root,
);
diagnostics.extend(call_diagnostics);
for (from, to, confidence) in call_edges {
if let (Some(from_idx), Some(to_idx)) =
(graph.symbol_node_index(from), graph.symbol_node_index(to))
{
graph.add_edge_normalized(
from_idx,
to_idx,
EdgeKind::Reference,
confidence,
);
}
}
}
}
let scc = crate::graph::SccAnalysis::analyze(graph.graph());
(graph, scc, scope_cache)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::model::{LineColumn, SourceRange, SymbolKind, Visibility};
fn test_source_range() -> SourceRange {
SourceRange {
byte_start: 0,
byte_end: 10,
start: LineColumn { line: 0, column: 0 },
end: LineColumn {
line: 0,
column: 10,
},
}
}
fn test_symbol(id: u32, name: &str, _file_id: u32) -> Symbol {
Symbol {
id: SymbolId::new(id).unwrap(),
name: name.to_string(),
kind: SymbolKind::Function,
language: LangId::Python,
file_path: PathBuf::from("test.py"),
source_range: test_source_range(),
visibility: Some(Visibility::Public),
signature: None,
docstring: None,
is_async: false,
}
}
#[test]
fn builder_creates_file_nodes() {
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
let path = PathBuf::from("src/main.py");
let id1 = builder.add_file(path.clone(), LangId::Python);
let id2 = builder.add_file(path, LangId::Python);
assert_eq!(id1, id2, "same path should return same FileId");
assert_eq!(builder.node_count(), 1);
}
#[test]
fn builder_creates_symbol_with_ownership() {
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
let path = PathBuf::from("test.py");
let file_id = builder.add_file(path.clone(), LangId::Python);
let symbol = test_symbol(1, "hello", file_id.to_raw());
let _sym_idx = builder.add_symbol(&symbol).unwrap();
assert_eq!(builder.node_count(), 2); assert_eq!(builder.edge_count(), 1); }
#[test]
fn builder_deduplicates_edges() {
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
let path1 = PathBuf::from("a.py");
let path2 = PathBuf::from("b.py");
let file1 = builder.add_file(path1, LangId::Python);
let _file2 = builder.add_file(path2, LangId::Python);
builder.add_import(file1, PathBuf::from("b.py"));
builder.add_import(file1, PathBuf::from("b.py"));
assert_eq!(
builder.edge_count(),
1,
"duplicate edges should be deduplicated"
);
}
#[test]
fn builder_reference_max_merges_confidence() {
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
builder.add_file(PathBuf::from("test.py"), LangId::Python);
let low = test_symbol(10, "low_fn", 0);
let high = test_symbol(11, "high_fn", 0);
builder.add_symbol(&low).unwrap();
builder.add_symbol(&high).unwrap();
builder.add_reference(low.id, high.id, 0.5);
builder.add_reference(low.id, high.id, 0.9);
assert_eq!(builder.edge_count(), 3);
let graph = builder.build();
assert_eq!(graph.edge_count(), 3);
let refs: Vec<_> = graph.edges_of_kind(EdgeKind::Reference).collect();
assert_eq!(refs.len(), 1);
let (src, dst) = refs[0];
let edge = graph.graph().edges_connecting(src, dst).next().unwrap();
assert_eq!(edge.weight().confidence, 0.9);
}
#[test]
fn builder_creates_external_node_for_unknown_import() {
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
let path = PathBuf::from("main.py");
let file_id = builder.add_file(path, LangId::Python);
builder.add_import(file_id, PathBuf::from("external_module.py"));
assert_eq!(
builder.edge_count(),
1,
"external import should create an edge"
);
assert_eq!(builder.external_count(), 1, "should have one external node");
}
#[test]
fn builder_tracks_node_mappings() {
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
let path = PathBuf::from("test.py");
let file_id = builder.add_file(path, LangId::Python);
let symbol = test_symbol(42, "func", file_id.to_raw());
builder.add_symbol(&symbol).unwrap();
assert!(builder.file_to_index.contains_key(&file_id));
assert!(
builder
.symbol_to_index
.contains_key(&SymbolId::new(42).unwrap())
);
}
#[test]
fn from_extractions_builds_graph_with_correct_node_count() {
use crate::model::{FileExtraction, LineColumn, SourceRange, Symbol, SymbolId, SymbolKind};
use std::path::PathBuf;
let sym = Symbol {
id: SymbolId::new(1).unwrap(),
name: "foo".into(),
kind: SymbolKind::Function,
language: LangId::Python,
file_path: PathBuf::from("/proj/a.py"),
source_range: SourceRange {
byte_start: 0,
byte_end: 10,
start: LineColumn { line: 0, column: 0 },
end: LineColumn {
line: 0,
column: 10,
},
},
visibility: None,
signature: None,
docstring: None,
is_async: false,
};
let mut base = FileExtraction::empty(PathBuf::from("/proj/a.py"), LangId::Python);
base.symbols = vec![sym];
let extractions = vec![base];
let mut diags = Vec::new();
let (graph, _scc) = GraphBuilder::from_extractions(
&extractions,
std::path::Path::new("/proj"),
SnapshotId::new(1).unwrap(),
&mut diags,
);
assert_eq!(graph.file_count(), 1);
assert_eq!(graph.symbol_count(), 1);
assert!(diags.is_empty());
}
#[test]
fn from_extractions_populates_scc_analysis() {
use crate::model::FileExtraction;
let extractions: Vec<FileExtraction> = vec![];
let mut diags = Vec::new();
let (graph, scc) = GraphBuilder::from_extractions(
&extractions,
std::path::Path::new("/proj"),
SnapshotId::new(1).unwrap(),
&mut diags,
);
assert_eq!(graph.node_count(), 0);
assert!(!scc.components.iter().any(|c| c.is_cyclic));
}
#[test]
fn from_extractions_accumulates_diagnostics_on_symbol_error() {
use crate::model::{FileExtraction, LineColumn, SourceRange, Symbol, SymbolId, SymbolKind};
use std::path::PathBuf;
let sym = Symbol {
id: SymbolId::new(99).unwrap(),
name: "orphan".into(),
kind: SymbolKind::Function,
language: LangId::Python,
file_path: PathBuf::from("/proj/missing.py"),
source_range: SourceRange {
byte_start: 0,
byte_end: 5,
start: LineColumn { line: 0, column: 0 },
end: LineColumn { line: 0, column: 5 },
},
visibility: None,
signature: None,
docstring: None,
is_async: false,
};
let mut base = FileExtraction::empty(PathBuf::from("/proj/a.py"), LangId::Python);
base.symbols = vec![sym];
let extractions = vec![base];
let mut diags = Vec::new();
let (_graph, _scc) = GraphBuilder::from_extractions(
&extractions,
std::path::Path::new("/proj"),
SnapshotId::new(1).unwrap(),
&mut diags,
);
assert!(!diags.is_empty(), "expected diagnostic for orphan symbol");
}
#[test]
fn from_extractions_resolves_cross_file_imports() {
use crate::model::{FileExtraction, LineColumn, SourceRange, UnresolvedImport};
use std::path::PathBuf;
let mut first = FileExtraction::empty(PathBuf::from("/proj/a.py"), LangId::Python);
first.imports = vec![UnresolvedImport {
import_specifier: "b".into(),
alias: None,
symbol: None,
star: false,
range: SourceRange {
byte_start: 0,
byte_end: 1,
start: LineColumn { line: 0, column: 0 },
end: LineColumn { line: 0, column: 1 },
},
}];
let second = FileExtraction::empty(PathBuf::from("/proj/b.py"), LangId::Python);
let extractions = vec![first, second];
let mut diags = Vec::new();
let (graph, _scc) = GraphBuilder::from_extractions(
&extractions,
std::path::Path::new("/proj"),
SnapshotId::new(1).unwrap(),
&mut diags,
);
assert_eq!(graph.file_count(), 2);
let import_edges: Vec<_> = graph
.edges_of_kind(crate::graph::EdgeKind::Import)
.collect();
assert_eq!(
import_edges.len(),
1,
"expected import edge from a.py to b.py"
);
}
#[cfg(feature = "dataflow")]
#[test]
fn add_data_node_creates_ownership_edge_to_symbol() {
use crate::model::{DataNode, DataNodeId, DataScope};
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
let file_path = PathBuf::from("test.rs");
let _file_id = builder.add_file(file_path.clone(), LangId::Rust);
let sym = Symbol {
id: SymbolId::new(1).unwrap(),
name: "my_fn".into(),
kind: SymbolKind::Function,
language: LangId::Rust,
file_path,
source_range: test_source_range(),
visibility: Some(Visibility::Public),
signature: None,
docstring: None,
is_async: false,
};
let sym_idx = builder.add_symbol(&sym).unwrap();
let data = DataNode {
id: DataNodeId::new(1).unwrap(),
symbol_id: Some(SymbolId::new(1).unwrap()),
name: Some("x".into()),
scope: DataScope::Local,
type_hint: Some("i32".into()),
source_range: test_source_range(),
};
builder.add_data_node(&data);
let graph = builder.build();
let ownership_edges: Vec<_> = graph.edges_of_kind(EdgeKind::Ownership).collect();
assert_eq!(
ownership_edges.len(),
2,
"expected file→symbol and symbol→data ownership edges"
);
let data_ownership_edges: Vec<_> = ownership_edges
.into_iter()
.filter(|(src, _)| *src == sym_idx)
.collect();
assert!(
!data_ownership_edges.is_empty(),
"expected symbol→data ownership edge"
);
}
}