pub mod builder;
pub mod edge;
pub mod naming;
pub mod node;
pub mod resolver;
pub mod scc;
use std::collections::HashMap;
pub use builder::{AnalysisParts, GraphBuilder};
pub use edge::{ConfidenceTier, EdgeData, EdgeKind, confidence_tier};
pub use node::{
DataGraphNode, ExternalClassification, ExternalNode, FileNode, NodeData, SymbolNode,
};
pub use scc::{DeployabilityHint, Scc, SccAnalysis};
use crate::language::LangId;
use crate::model::{FileId, SnapshotId, SymbolId};
use petgraph::graph::{DiGraph, EdgeIndex, NodeIndex};
#[derive(Debug, Clone)]
pub struct CodeGraph {
graph: DiGraph<NodeData, EdgeData>,
edge_index: HashMap<(NodeIndex, NodeIndex, EdgeKind), EdgeIndex>,
pub(crate) file_to_index: HashMap<FileId, NodeIndex>,
pub(crate) symbol_to_index: HashMap<SymbolId, NodeIndex>,
pub(crate) external_index: HashMap<String, NodeIndex>,
pub snapshot_id: SnapshotId,
}
impl CodeGraph {
pub fn new(snapshot_id: SnapshotId) -> Self {
Self {
graph: DiGraph::new(),
edge_index: HashMap::new(),
file_to_index: HashMap::new(),
symbol_to_index: HashMap::new(),
external_index: HashMap::new(),
snapshot_id,
}
}
pub(crate) fn from_parts(
graph: DiGraph<NodeData, EdgeData>,
edge_index: HashMap<(NodeIndex, NodeIndex, EdgeKind), EdgeIndex>,
file_to_index: HashMap<FileId, NodeIndex>,
symbol_to_index: HashMap<SymbolId, NodeIndex>,
external_index: HashMap<String, NodeIndex>,
snapshot_id: SnapshotId,
) -> Self {
Self {
graph,
edge_index,
file_to_index,
symbol_to_index,
external_index,
snapshot_id,
}
}
pub fn graph(&self) -> &DiGraph<NodeData, EdgeData> {
&self.graph
}
pub fn add_node(&mut self, node: NodeData) -> NodeIndex {
self.graph.add_node(node)
}
pub fn add_file_node(&mut self, node: FileNode) -> NodeIndex {
let file_id = node.id;
let idx = self.graph.add_node(NodeData::File(node));
self.file_to_index.insert(file_id, idx);
idx
}
pub fn add_symbol_node(&mut self, node: SymbolNode) -> NodeIndex {
let symbol_id = node.id;
let idx = self.graph.add_node(NodeData::Symbol(node));
self.symbol_to_index.insert(symbol_id, idx);
idx
}
pub fn file_node_index(&self, file_id: FileId) -> Option<NodeIndex> {
self.file_to_index.get(&file_id).copied()
}
pub fn symbol_node_index(&self, symbol_id: SymbolId) -> Option<NodeIndex> {
self.symbol_to_index.get(&symbol_id).copied()
}
pub fn file_node(&self, file_id: FileId) -> Option<&FileNode> {
let idx = self.file_node_index(file_id)?;
self.graph.node_weight(idx).and_then(|data| data.as_file())
}
pub fn symbol_node(&self, symbol_id: SymbolId) -> Option<&SymbolNode> {
let idx = self.symbol_node_index(symbol_id)?;
self.graph
.node_weight(idx)
.and_then(|data| data.as_symbol())
}
pub fn file_count(&self) -> usize {
self.file_to_index.len()
}
pub fn symbol_count(&self) -> usize {
self.symbol_to_index.len()
}
pub fn external_count(&self) -> usize {
self.external_index.len()
}
pub fn node_count(&self) -> usize {
self.graph.node_count()
}
pub fn edge_count(&self) -> usize {
self.graph.edge_count()
}
pub fn get_or_create_external_node(&mut self, raw_path: String, language: LangId) -> NodeIndex {
if let Some(&idx) = self.external_index.get(&raw_path) {
return idx;
}
let node = NodeData::External(ExternalNode {
raw_path: raw_path.clone(),
language,
classification: None,
});
let idx = self.graph.add_node(node);
self.external_index.insert(raw_path, idx);
idx
}
pub fn add_edge_normalized(
&mut self,
source: NodeIndex,
target: NodeIndex,
kind: EdgeKind,
confidence: f32,
) {
self.add_edge_normalized_with_flow(source, target, kind, confidence, None);
}
pub fn add_edge_normalized_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) {
self.graph[edge_idx].merge_repeated(confidence, 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 files(&self) -> impl Iterator<Item = (FileId, &FileNode)> + '_ {
let mut files: Vec<(FileId, &FileNode)> = self
.file_to_index
.iter()
.filter_map(|(file_id, &idx)| {
self.graph
.node_weight(idx)
.and_then(|data| data.as_file().map(|f| (*file_id, f)))
})
.collect();
files.sort_by(|a, b| a.1.path.cmp(&b.1.path));
files.into_iter()
}
pub fn symbols(&self) -> impl Iterator<Item = (SymbolId, &SymbolNode)> + '_ {
let mut symbols: Vec<(SymbolId, &SymbolNode)> = self
.symbol_to_index
.iter()
.filter_map(|(symbol_id, &idx)| {
self.graph
.node_weight(idx)
.and_then(|data| data.as_symbol().map(|s| (*symbol_id, s)))
})
.collect();
symbols.sort_by(|a, b| {
let left = self.file_node(a.1.file_id).map(|file| &file.path);
let right = self.file_node(b.1.file_id).map(|file| &file.path);
left.cmp(&right)
.then(a.1.name.cmp(&b.1.name))
.then(a.0.to_raw().cmp(&b.0.to_raw()))
});
symbols.into_iter()
}
pub fn edges_of_kind(
&self,
kind: EdgeKind,
) -> impl Iterator<Item = (NodeIndex, NodeIndex)> + '_ {
self.graph.edge_indices().filter_map(move |edge_idx| {
let (source, target) = self.graph.edge_endpoints(edge_idx)?;
let weight = self.graph.edge_weight(edge_idx)?;
if weight.kind == kind {
Some((source, target))
} else {
None
}
})
}
pub fn reference_edges(&self) -> impl Iterator<Item = (SymbolId, SymbolId, f32)> + '_ {
self.graph.edge_indices().filter_map(move |edge_idx| {
let weight = self.graph.edge_weight(edge_idx)?;
if weight.kind != EdgeKind::Reference {
return None;
}
let (source, target) = self.graph.edge_endpoints(edge_idx)?;
let source_id = self.graph.node_weight(source)?.as_symbol()?.id;
let target_id = self.graph.node_weight(target)?.as_symbol()?.id;
Some((source_id, target_id, weight.confidence))
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::language::LangId;
use crate::model::{LineColumn, SourceRange, Symbol, SymbolKind, Visibility, ids::SnapshotId};
use std::path::PathBuf;
fn test_range() -> SourceRange {
SourceRange {
byte_start: 0,
byte_end: 10,
start: LineColumn { line: 1, column: 0 },
end: LineColumn {
line: 1,
column: 10,
},
}
}
fn test_file_node(id: u32, path: &str) -> NodeData {
NodeData::File(FileNode {
id: FileId::new(id).unwrap(),
path: PathBuf::from(path),
language: LangId::Rust,
snapshot_id: SnapshotId::new(1).unwrap(),
})
}
fn test_symbol(id: u32, name: &str) -> Symbol {
Symbol {
id: SymbolId::new(id).unwrap(),
name: name.to_string(),
kind: SymbolKind::Function,
language: LangId::Rust,
file_path: PathBuf::from("test.rs"),
source_range: test_range(),
name_range: None,
visibility: Some(Visibility::Public),
signature: None,
docstring: None,
is_async: false,
}
}
#[test]
fn code_graph_new_empty() {
let graph = CodeGraph::new(SnapshotId::new(1).unwrap());
assert_eq!(graph.node_count(), 0);
assert_eq!(graph.edge_count(), 0);
assert_eq!(graph.snapshot_id.to_raw(), 1);
}
#[test]
fn code_graph_file_lookup_returns_none_for_missing() {
let graph = CodeGraph::new(SnapshotId::new(1).unwrap());
assert!(graph.file_node(FileId::new(1).unwrap()).is_none());
assert!(graph.file_node_index(FileId::new(1).unwrap()).is_none());
}
#[test]
fn code_graph_symbol_lookup_returns_none_for_missing() {
let graph = CodeGraph::new(SnapshotId::new(1).unwrap());
assert!(graph.symbol_node(SymbolId::new(1).unwrap()).is_none());
assert!(graph.symbol_node_index(SymbolId::new(1).unwrap()).is_none());
}
#[test]
fn builder_produces_valid_code_graph() {
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
let file_id = builder.add_file(PathBuf::from("test.rs"), LangId::Rust);
let symbol = test_symbol(1, "main");
let _sym_idx = builder.add_symbol(&symbol).unwrap();
let graph = builder.build();
assert_eq!(graph.file_count(), 1);
assert_eq!(graph.symbol_count(), 1);
assert_eq!(graph.node_count(), 2);
assert_eq!(graph.edge_count(), 1);
let file_lookup = graph.file_node(file_id);
assert!(file_lookup.is_some());
assert_eq!(file_lookup.unwrap().language, LangId::Rust);
let sym_lookup = graph.symbol_node(SymbolId::new(1).unwrap());
assert!(sym_lookup.is_some());
assert_eq!(sym_lookup.unwrap().name, "main");
}
#[test]
fn code_graph_iteration_over_files() {
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
builder.add_file(PathBuf::from("a.rs"), LangId::Rust);
builder.add_file(PathBuf::from("b.py"), LangId::Python);
let graph = builder.build();
let files: Vec<_> = graph.files().collect();
assert_eq!(files.len(), 2);
}
#[test]
fn file_iteration_is_path_sorted() {
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
for name in ["e.rs", "a.rs", "d.rs", "b.rs", "c.rs"] {
builder.add_file(PathBuf::from(name), LangId::Rust);
}
let graph = builder.build();
let paths: Vec<PathBuf> = graph.files().map(|(_, file)| file.path.clone()).collect();
let mut sorted = paths.clone();
sorted.sort();
assert_eq!(paths, sorted, "file iteration must be path sorted");
let again: Vec<PathBuf> = graph.files().map(|(_, file)| file.path.clone()).collect();
assert_eq!(paths, again, "file iteration must be stable");
}
#[test]
fn symbol_iteration_is_path_then_name_sorted() {
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
let entries = [
("b.rs", 3, "zeta"),
("a.rs", 1, "beta"),
("a.rs", 2, "alpha"),
("c.rs", 4, "gamma"),
];
for (path, id, name) in entries.iter() {
builder.add_file(PathBuf::from(path), LangId::Rust);
let mut symbol = test_symbol(*id, name);
symbol.file_path = PathBuf::from(path);
builder.add_symbol(&symbol).unwrap();
}
let mut sorted_entries = entries;
sorted_entries.sort_by_key(|(path, _, name)| (*path, *name));
let graph = builder.build();
let ordered: Vec<(String, String)> = graph
.symbols()
.map(|(_, symbol)| {
let path = graph
.file_node(symbol.file_id)
.map(|file| file.path.display().to_string())
.unwrap_or_default();
(path, symbol.name.clone())
})
.collect();
let entries = sorted_entries;
let expected: Vec<(String, String)> = entries
.iter()
.map(|(path, _, name)| (path.to_string(), name.to_string()))
.collect();
assert_eq!(ordered, expected, "symbols must be path then name sorted");
}
#[test]
fn code_graph_iteration_over_symbols() {
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
let _file_id = builder.add_file(PathBuf::from("test.rs"), LangId::Rust);
let sym1 = test_symbol(1, "func_a");
let sym2 = test_symbol(2, "func_b");
builder.add_symbol(&sym1).unwrap();
builder.add_symbol(&sym2).unwrap();
let graph = builder.build();
let symbols: Vec<_> = graph.symbols().collect();
assert_eq!(symbols.len(), 2);
}
#[test]
fn code_graph_edges_of_kind_filtering() {
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
let file1 = builder.add_file(PathBuf::from("a.rs"), LangId::Rust);
let _file2 = builder.add_file(PathBuf::from("b.rs"), LangId::Rust);
builder.add_import(file1, PathBuf::from("b.rs"));
let graph = builder.build();
let ownership_edges: Vec<_> = graph.edges_of_kind(EdgeKind::Ownership).collect();
let import_edges: Vec<_> = graph.edges_of_kind(EdgeKind::Import).collect();
assert_eq!(ownership_edges.len(), 0); assert_eq!(import_edges.len(), 1);
}
#[test]
fn build_populated_edge_index_normalizes_post_build_edges() {
let mut builder = GraphBuilder::new(SnapshotId::new(1).unwrap());
let file_a = builder.add_file(PathBuf::from("a.rs"), LangId::Rust);
let file_b = builder.add_file(PathBuf::from("b.rs"), LangId::Rust);
builder.add_import(file_a, PathBuf::from("b.rs"));
let mut graph = builder.build();
assert_eq!(graph.edge_count(), 1);
let a_idx = graph.file_node_index(file_a).unwrap();
let b_idx = graph.file_node_index(file_b).unwrap();
graph.add_edge_normalized(a_idx, b_idx, EdgeKind::Import, 0.5);
graph.add_edge_normalized(a_idx, b_idx, EdgeKind::Import, 0.8);
assert_eq!(graph.edge_count(), 1);
graph.add_edge_normalized(a_idx, b_idx, EdgeKind::Reference, 0.9);
assert_eq!(graph.edge_count(), 2);
}
#[test]
fn add_edge_normalized_handles_multiple_edge_kinds_between_same_nodes() {
let mut graph = CodeGraph::new(SnapshotId::new(1).unwrap());
let n1 = graph.add_node(test_file_node(1, "a.rs"));
let n2 = graph.add_node(test_file_node(2, "b.rs"));
graph.add_edge_normalized(n1, n2, EdgeKind::Reference, 0.7);
graph.add_edge_normalized(n1, n2, EdgeKind::Import, 0.5);
graph.add_edge_normalized(n1, n2, EdgeKind::Reference, 0.9);
let ref_count = graph
.graph()
.edges_connecting(n1, n2)
.filter(|e| e.weight().kind == EdgeKind::Reference)
.count();
assert_eq!(
ref_count, 1,
"Expected 1 Reference edge, but found {ref_count}"
);
assert_eq!(graph.edge_count(), 2);
}
#[test]
fn add_edge_normalized_with_flow_preserves_first_flow_kind() {
use crate::model::{DataNodeId, DataScope, FlowKind};
let mut graph = CodeGraph::new(SnapshotId::new(1).unwrap());
let n1 = graph.add_node(NodeData::Data(DataGraphNode {
id: DataNodeId::new(1).unwrap(),
symbol_id: None,
name: Some("x".into()),
scope: DataScope::Local,
type_hint: None,
source_range: test_range(),
}));
let n2 = graph.add_node(NodeData::Data(DataGraphNode {
id: DataNodeId::new(2).unwrap(),
symbol_id: None,
name: Some("y".into()),
scope: DataScope::Local,
type_hint: None,
source_range: test_range(),
}));
graph.add_edge_normalized_with_flow(n1, n2, EdgeKind::Flow, 0.9, Some(FlowKind::DefUse));
graph.add_edge_normalized_with_flow(n1, n2, EdgeKind::Flow, 0.8, Some(FlowKind::Argument));
assert_eq!(graph.edge_count(), 1);
let edge = graph.graph().edges_connecting(n1, n2).next().unwrap();
assert_eq!(edge.weight().flow_kind, Some(FlowKind::DefUse));
assert_eq!(edge.weight().confidence, 0.9);
}
#[test]
fn add_edge_normalized_duplicate_never_grows_edge_count() {
let mut graph = CodeGraph::new(SnapshotId::new(1).unwrap());
let n1 = graph.add_node(test_file_node(1, "a.rs"));
let n2 = graph.add_node(test_file_node(2, "b.rs"));
graph.add_edge_normalized(n1, n2, EdgeKind::Import, 0.5);
for confidence in [0.6, 0.3, 0.9, 0.4, 0.7] {
graph.add_edge_normalized(n1, n2, EdgeKind::Import, confidence);
}
assert_eq!(graph.edge_count(), 1);
let edge = graph.graph().edges_connecting(n1, n2).next().unwrap();
assert_eq!(edge.weight().confidence, 0.9);
}
#[test]
fn add_edge_normalized_counts_only_distinct_triples() {
let mut graph = CodeGraph::new(SnapshotId::new(1).unwrap());
let n1 = graph.add_node(test_file_node(1, "a.rs"));
let n2 = graph.add_node(test_file_node(2, "b.rs"));
let n3 = graph.add_node(test_file_node(3, "c.rs"));
graph.add_edge_normalized(n1, n2, EdgeKind::Import, 0.5);
graph.add_edge_normalized(n1, n2, EdgeKind::Reference, 0.6);
graph.add_edge_normalized(n1, n3, EdgeKind::Reference, 0.7);
graph.add_edge_normalized(n1, n2, EdgeKind::Import, 0.9);
assert_eq!(graph.edge_count(), 3);
}
}