use std::collections::HashMap;
use std::path::PathBuf;
use petgraph::graph::{DiGraph, NodeIndex};
use petgraph::visit::EdgeRef;
use crate::graph::CodeGraph;
use crate::graph::edge::{EdgeData, EdgeKind};
use crate::graph::node::{ExternalNode, FileNode, NodeData, SymbolNode};
use crate::language::LangId;
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_normalizer: EdgeNormalizer,
external_index: HashMap<String, NodeIndex>,
}
#[derive(Debug, Default)]
struct EdgeNormalizer {
seen: std::collections::HashSet<(NodeIndex, NodeIndex, EdgeKind)>,
}
impl EdgeNormalizer {
fn is_new(&mut self, source: NodeIndex, target: NodeIndex, kind: EdgeKind) -> bool {
self.seen.insert((source, target, kind))
}
}
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_normalizer: EdgeNormalizer::default(),
external_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)
}
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);
}
pub fn add_metacall_load(
&mut self,
from_path: PathBuf,
target_lang: LangId,
scripts: &[String],
confidence: f64,
root: &std::path::Path,
) {
let Some(&from_fid) = self.path_to_file.get(&from_path) else {
return;
};
let Some(&from_idx) = self.file_to_index.get(&from_fid) else {
return;
};
for script in scripts {
let target_path = root.join(script);
if let Some(&to_fid) = self.path_to_file.get(&target_path) {
if let Some(&to_idx) = self.file_to_index.get(&to_fid) {
let edge_data = EdgeData {
kind: EdgeKind::Import,
confidence: confidence as f32,
};
self.graph.add_edge(from_idx, to_idx, edge_data);
}
} else {
let raw_path = script.clone();
if let Some(&to_idx) = self.external_index.get(&raw_path) {
let edge_data = EdgeData {
kind: EdgeKind::Import,
confidence: confidence as f32,
};
self.graph.add_edge(from_idx, to_idx, edge_data);
} else {
let node = ExternalNode {
raw_path: raw_path.clone(),
language: target_lang,
classification: None,
};
let idx = self.graph.add_node(NodeData::External(node));
self.external_index.insert(raw_path, idx);
let edge_data = EdgeData {
kind: EdgeKind::Import,
confidence: confidence as f32,
};
self.graph.add_edge(from_idx, idx, edge_data);
}
}
}
}
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,
) {
if !self.edge_normalizer.is_new(source, target, kind) {
if let Some(edge_idx) = self.graph.find_edge(source, target) {
let existing = &mut self.graph[edge_idx];
if existing.kind == kind {
existing.confidence = existing.confidence.max(confidence);
}
}
return;
}
let edge_data = EdgeData { kind, confidence };
self.graph.add_edge(source, target, edge_data);
}
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 {
graph: self.graph,
file_to_index: self.file_to_index,
symbol_to_index: self.symbol_to_index,
external_index: self.external_index,
snapshot_id: 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(
extractions: &[crate::model::FileExtraction],
root: &std::path::Path,
snapshot_id: crate::model::SnapshotId,
diagnostics: &mut Vec<crate::error::Diagnostic>,
) -> (CodeGraph, crate::graph::SccAnalysis) {
let mut builder = Self::new(snapshot_id);
for file in extractions {
builder.add_file(file.path.clone(), file.lang);
}
for file in extractions {
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,
});
}
}
}
let path_to_file_id: HashMap<std::path::PathBuf, crate::model::FileId> = extractions
.iter()
.filter_map(|f| {
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 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();
let scc = crate::graph::SccAnalysis::analyze(&graph.graph);
(graph, scc)
}
}
#[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(id),
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(1));
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(1));
let path = PathBuf::from("test.py");
let file_id = builder.add_file(path.clone(), LangId::Python);
let symbol = test_symbol(1, "hello", file_id.0);
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(1));
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_creates_external_node_for_unknown_import() {
let mut builder = GraphBuilder::new(SnapshotId(1));
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(1));
let path = PathBuf::from("test.py");
let file_id = builder.add_file(path, LangId::Python);
let symbol = test_symbol(42, "func", file_id.0);
builder.add_symbol(&symbol).unwrap();
assert!(builder.file_to_index.contains_key(&file_id));
assert!(builder.symbol_to_index.contains_key(&SymbolId(42)));
}
#[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(1),
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 extractions = vec![FileExtraction {
path: PathBuf::from("/proj/a.py"),
lang: LangId::Python,
symbols: vec![sym],
imports: vec![],
references: vec![],
diagnostics: vec![],
ast_node_count: 0,
}];
let mut diags = Vec::new();
let (graph, _scc) = GraphBuilder::from_extractions(
&extractions,
std::path::Path::new("/proj"),
SnapshotId(1),
&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(1),
&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(99),
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 extractions = vec![FileExtraction {
path: PathBuf::from("/proj/a.py"),
lang: LangId::Python,
symbols: vec![sym],
imports: vec![],
references: vec![],
diagnostics: vec![],
ast_node_count: 0,
}];
let mut diags = Vec::new();
let (_graph, _scc) = GraphBuilder::from_extractions(
&extractions,
std::path::Path::new("/proj"),
SnapshotId(1),
&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 extractions = vec![
FileExtraction {
path: PathBuf::from("/proj/a.py"),
lang: LangId::Python,
symbols: vec![],
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 },
},
}],
references: vec![],
diagnostics: vec![],
ast_node_count: 0,
},
FileExtraction {
path: PathBuf::from("/proj/b.py"),
lang: LangId::Python,
symbols: vec![],
imports: vec![],
references: vec![],
diagnostics: vec![],
ast_node_count: 0,
},
];
let mut diags = Vec::new();
let (graph, _scc) = GraphBuilder::from_extractions(
&extractions,
std::path::Path::new("/proj"),
SnapshotId(1),
&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"
);
}
}