use crate::parser::{ImportInfo, ParseResult};
use crate::schema::{CodeEdge, CodeEdgePredicate, CodeNode, CodeNodeKind};
use crate::scip_calls;
use std::collections::HashMap;
use std::path::Path;
#[derive(Debug, Default)]
pub struct NameResolver {
name_to_id: HashMap<String, String>,
short_name_to_id: HashMap<String, String>,
}
impl NameResolver {
pub fn from_nodes(nodes: &[CodeNode]) -> Self {
let mut resolver = Self::default();
for node in nodes {
resolver.name_to_id.insert(node.id.clone(), node.id.clone());
let qualified = node_to_qualified_name(node);
if !qualified.is_empty() {
resolver.name_to_id.insert(qualified, node.id.clone());
}
if !node.name.is_empty() {
resolver
.short_name_to_id
.entry(node.name.clone())
.or_insert_with(|| node.id.clone());
}
}
resolver
}
pub fn resolve(&self, name: &str) -> Option<&String> {
self.name_to_id
.get(name)
.or_else(|| self.short_name_to_id.get(name))
}
}
pub fn extract_edges(parse_results: &[ParseResult], resolver: &NameResolver) -> Vec<CodeEdge> {
let mut edges = Vec::new();
let all_nodes: Vec<&CodeNode> = parse_results.iter().flat_map(|r| &r.nodes).collect();
for node in &all_nodes {
if let Some(parent_id) = &node.parent_id {
edges.push(CodeEdge {
source_id: parent_id.clone(),
target_id: node.id.clone(),
predicate: CodeEdgePredicate::Contains,
weight: None,
commit_id: None,
});
}
}
let all_imports: Vec<&ImportInfo> = parse_results.iter().flat_map(|r| &r.imports).collect();
for imp in &all_imports {
let target = resolve_import(imp, resolver);
let target_id = target.unwrap_or_else(|| format!("ext:{}", imp.module));
edges.push(CodeEdge {
source_id: imp.file_node_id.clone(),
target_id,
predicate: CodeEdgePredicate::Imports,
weight: None,
commit_id: None,
});
}
for node in &all_nodes {
if node.kind == CodeNodeKind::Class
&& let Some(sig) = &node.signature
{
let bases = extract_bases_from_signature(sig);
for base in bases {
let target_id = resolver
.resolve(&base)
.cloned()
.unwrap_or_else(|| format!("ext:{base}"));
edges.push(CodeEdge {
source_id: node.id.clone(),
target_id,
predicate: CodeEdgePredicate::InheritsFrom,
weight: None,
commit_id: None,
});
}
}
}
edges
}
pub fn extract_call_edges(
parse_results: &[ParseResult],
_resolver: &NameResolver,
source_texts: &HashMap<String, String>,
crate_dir: Option<&Path>,
) -> Vec<CodeEdge> {
let all_nodes: Vec<CodeNode> = parse_results.iter().flat_map(|r| r.nodes.clone()).collect();
if let Some(dir) = crate_dir {
let result = scip_calls::extract_scip_call_edges(dir, &all_nodes);
if !result.edges.is_empty() {
tracing::info!(
"SCIP call edges: {} resolved, {} unresolved ({} docs)",
result.symbols_resolved,
result.unresolved_references,
result.documents_processed
);
return result.edges;
} else if !result.warnings.is_empty() {
tracing::debug!(
"SCIP unavailable or failed, falling back to text scan: {}",
result.warnings.join("; ")
);
}
}
text_scan_call_edges(&all_nodes, source_texts)
}
fn text_scan_call_edges(
all_nodes: &[CodeNode],
source_texts: &HashMap<String, String>,
) -> Vec<CodeEdge> {
let mut edges = Vec::new();
let callable_names: HashMap<&str, &str> = all_nodes
.iter()
.filter(|n| {
matches!(
n.kind,
CodeNodeKind::Function | CodeNodeKind::Method | CodeNodeKind::Test
)
})
.map(|n| (n.name.as_str(), n.id.as_str()))
.collect();
for node in all_nodes {
if !matches!(
node.kind,
CodeNodeKind::Function | CodeNodeKind::Method | CodeNodeKind::Test
) {
continue;
}
let file_path = node_id_to_path(&node.id);
let Some(source) = source_texts.get(&file_path) else {
continue;
};
for (name, target_id) in &callable_names {
if *target_id == node.id {
continue; }
let call_pattern = format!("{name}(");
let count = source.matches(&call_pattern).count();
if count > 0 {
let valid = source.contains(&format!(".{call_pattern}"))
|| source.contains(&format!(" {call_pattern}"))
|| source.contains(&format!("={call_pattern}"))
|| source.contains(&format!("({call_pattern}"))
|| source.starts_with(&call_pattern);
if valid {
edges.push(CodeEdge {
source_id: node.id.clone(),
target_id: target_id.to_string(),
predicate: CodeEdgePredicate::Calls,
weight: Some(count as f32),
commit_id: None,
});
}
}
}
}
edges
}
fn node_to_qualified_name(node: &CodeNode) -> String {
let id = &node.id;
let after_colon = id.split_once(':').map(|(_, rest)| rest).unwrap_or(id);
let (path, name) = after_colon.split_once("::").unwrap_or((after_colon, ""));
let mod_path = path.strip_suffix(".py").unwrap_or(path).replace('/', ".");
if name.is_empty() {
mod_path
} else if name.contains("::") {
format!("{mod_path}.{}", name.replace("::", "."))
} else {
format!("{mod_path}.{name}")
}
}
fn node_id_to_path(id: &str) -> String {
let after_colon = id.split_once(':').map(|(_, rest)| rest).unwrap_or(id);
after_colon
.split_once("::")
.map(|(path, _)| path.to_string())
.unwrap_or_else(|| after_colon.to_string())
}
fn resolve_import(imp: &ImportInfo, resolver: &NameResolver) -> Option<String> {
let mod_id = format!("mod:{}.py", imp.module.replace('.', "/"));
if let Some(id) = resolver.resolve(&mod_id) {
return Some(id.clone());
}
let file_id = format!("file:{}.py", imp.module.replace('.', "/"));
if let Some(id) = resolver.resolve(&file_id) {
return Some(id.clone());
}
resolver.resolve(&imp.module).cloned()
}
fn extract_bases_from_signature(sig: &str) -> Vec<String> {
let Some(start) = sig.find('(') else {
return Vec::new();
};
let Some(end) = sig.rfind(')') else {
return Vec::new();
};
if start >= end {
return Vec::new();
}
sig[start + 1..end]
.split(',')
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect()
}
pub fn extract_cross_file_edges(
parse_results: &[ParseResult],
resolver: &crate::module_resolver::RustModuleResolver,
) -> Vec<CodeEdge> {
let mut edges = Vec::new();
let all_nodes: Vec<&CodeNode> = parse_results.iter().flat_map(|r| &r.nodes).collect();
let all_imports: Vec<&ImportInfo> = parse_results.iter().flat_map(|r| &r.imports).collect();
for imp in &all_imports {
let resolved = resolver.resolve_use(&imp.module, &imp.names);
for (name, target_id) in &resolved {
edges.push(CodeEdge {
source_id: imp.file_node_id.clone(),
target_id: target_id.clone(),
predicate: CodeEdgePredicate::Imports,
weight: None,
commit_id: None,
});
let use_node = all_nodes.iter().find(|n| {
n.kind == CodeNodeKind::RustUse
&& n.name == *name
&& n.parent_id.as_deref() == Some(&imp.file_node_id)
});
if let Some(use_node) = use_node
&& use_node
.signature
.as_ref()
.is_some_and(|s| s.starts_with("pub"))
{
edges.push(CodeEdge {
source_id: imp.file_node_id.clone(),
target_id: target_id.clone(),
predicate: CodeEdgePredicate::ReExports,
weight: None,
commit_id: None,
});
}
}
}
for node in &all_nodes {
if node.kind != CodeNodeKind::RustImpl {
continue;
}
if let Some(sig) = &node.signature
&& let Some(trait_name) = extract_trait_from_impl(sig)
{
let target_id = resolver
.resolve_name(&trait_name)
.unwrap_or_else(|| format!("ext:{trait_name}"));
edges.push(CodeEdge {
source_id: node.id.clone(),
target_id,
predicate: CodeEdgePredicate::ImplementsTrait,
weight: None,
commit_id: None,
});
}
}
for node in &all_nodes {
if node.kind != CodeNodeKind::RustTest && node.kind != CodeNodeKind::Test {
continue;
}
if let Some(target_name) = node.name.strip_prefix("test_")
&& let Some(target_id) = resolver.resolve_name(target_name)
{
edges.push(CodeEdge {
source_id: node.id.clone(),
target_id,
predicate: CodeEdgePredicate::TestTargets,
weight: None,
commit_id: None,
});
}
}
edges
}
fn extract_trait_from_impl(sig: &str) -> Option<String> {
let sig = sig.strip_prefix("impl").map(|s| s.trim())?;
if let Some(for_idx) = sig.find(" for ") {
let trait_part = sig[..for_idx].trim();
let name = trait_part
.split('<')
.next()
.unwrap_or(trait_part)
.trim()
.to_string();
if name.is_empty() { None } else { Some(name) }
} else {
None }
}
#[cfg(test)]
mod tests {
use super::*;
use crate::parser::parse_python_file;
use std::path::PathBuf;
const SAMPLE: &str = r#"
import os
from pathlib import Path
class Parent:
"""Base class."""
def greet(self):
return "hello"
class Child(Parent):
"""Child class."""
def greet(self):
return "hi"
def helper():
"""A helper."""
return 42
def caller():
"""Calls helper."""
x = helper()
return x
"#;
#[test]
fn test_containment_edges() {
let path = PathBuf::from("test.py");
let result = parse_python_file(&path, SAMPLE).unwrap();
let resolver = NameResolver::from_nodes(&result.nodes);
let edges = extract_edges(&[result], &resolver);
let containment: Vec<_> = edges
.iter()
.filter(|e| e.predicate == CodeEdgePredicate::Contains)
.collect();
assert!(
containment.len() >= 5,
"Expected >= 5 containment edges, got {}",
containment.len()
);
}
#[test]
fn test_import_edges() {
let path = PathBuf::from("test.py");
let result = parse_python_file(&path, SAMPLE).unwrap();
let resolver = NameResolver::from_nodes(&result.nodes);
let edges = extract_edges(&[result], &resolver);
let imports: Vec<_> = edges
.iter()
.filter(|e| e.predicate == CodeEdgePredicate::Imports)
.collect();
assert!(
imports.len() >= 2,
"Expected >= 2 import edges, got {}",
imports.len()
);
let ext_imports: Vec<_> = imports
.iter()
.filter(|e| e.target_id.starts_with("ext:"))
.collect();
assert!(
!ext_imports.is_empty(),
"External imports should have ext: prefix"
);
}
#[test]
fn test_inheritance_edges() {
let path = PathBuf::from("test.py");
let result = parse_python_file(&path, SAMPLE).unwrap();
let resolver = NameResolver::from_nodes(&result.nodes);
let edges = extract_edges(&[result], &resolver);
let inheritance: Vec<_> = edges
.iter()
.filter(|e| e.predicate == CodeEdgePredicate::InheritsFrom)
.collect();
assert_eq!(
inheritance.len(),
1,
"Should have 1 inheritance edge (Child→Parent)"
);
assert!(inheritance[0].source_id.contains("Child"));
}
#[test]
fn test_extract_bases_from_signature() {
assert_eq!(
extract_bases_from_signature("class Child(Parent, Mixin)"),
vec!["Parent", "Mixin"]
);
assert_eq!(
extract_bases_from_signature("class Foo"),
Vec::<String>::new()
);
assert_eq!(
extract_bases_from_signature("class Bar()"),
Vec::<String>::new()
);
}
#[test]
fn test_name_resolver() {
let nodes = vec![CodeNode {
id: "func:brain/utils.py::helper".to_string(),
kind: CodeNodeKind::Function,
parent_id: None,
name: "helper".to_string(),
signature: None,
docstring: None,
body_hash: None,
body: None,
loc: None,
cyclomatic_complexity: None,
coverage_pct: None,
last_modified: None,
..Default::default()
}];
let resolver = NameResolver::from_nodes(&nodes);
assert_eq!(
resolver.resolve("helper"),
Some(&"func:brain/utils.py::helper".to_string())
);
assert_eq!(
resolver.resolve("func:brain/utils.py::helper"),
Some(&"func:brain/utils.py::helper".to_string())
);
}
}