use std::collections::{HashSet, VecDeque};
use std::path::{Path, PathBuf};
use anyhow::Result;
use petgraph::Direction;
use petgraph::stable_graph::{NodeIndex, StableGraph};
use petgraph::visit::{EdgeRef, IntoEdgeReferences};
use serde::{Deserialize, Serialize};
use tree_sitter::{Node, Parser};
use crate::symgraph::registry::SymbolRegistry;
use crate::symgraph::resolve::{NameIndex, bare_name, rank_matches, resolve_callee};
use crate::symgraph::symbol::{SymbolKind, detect_language, extract_symbols};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SymbolNode {
pub file: PathBuf,
pub name: String,
pub kind: SymbolKind,
pub start_line: usize,
}
pub use crate::symgraph::contracts::EdgeKind;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SymbolEdge {
pub from: String,
pub to: String,
pub kind: EdgeKind,
}
pub type Edge = SymbolEdge;
struct RawEdge {
caller: NodeIndex,
callee: String,
kind: EdgeKind,
}
#[derive(Debug)]
pub enum SymbolMatch<'a> {
NotFound,
Unique(&'a SymbolNode),
Ambiguous {
chosen: &'a SymbolNode,
alternatives: Vec<&'a SymbolNode>,
},
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SymbolGraph {
#[serde(rename = "graph")]
inner: StableGraph<SymbolNode, EdgeKind>,
#[serde(skip, default)]
names: NameIndex,
}
impl Default for SymbolGraph {
fn default() -> Self {
Self {
inner: StableGraph::new(),
names: NameIndex::default(),
}
}
}
impl SymbolGraph {
pub fn new() -> Self {
Self::default()
}
pub fn node_count(&self) -> usize {
self.inner.node_count()
}
pub fn edge_count(&self) -> usize {
self.inner.edge_count()
}
pub fn inner(&self) -> &StableGraph<SymbolNode, EdgeKind> {
&self.inner
}
pub fn nodes(&self) -> Vec<&SymbolNode> {
self.inner.node_indices().map(|i| &self.inner[i]).collect()
}
pub fn edges(&self) -> Vec<SymbolEdge> {
self.inner
.edge_references()
.map(|er| {
let from = self.inner[er.source()].name.clone();
let to = self.inner[er.target()].name.clone();
SymbolEdge {
from,
to,
kind: er.weight().clone(),
}
})
.collect()
}
fn add_node(&mut self, node: SymbolNode) -> NodeIndex {
let name = node.name.clone();
self.add_node_as(node, &name)
}
fn add_node_as(&mut self, node: SymbolNode, symbol: &str) -> NodeIndex {
let file = node.file.display().to_string();
let callable = matches!(node.kind, SymbolKind::Function | SymbolKind::Method);
let idx = self.inner.add_node(node);
self.names.insert(&file, symbol, idx, callable);
idx
}
fn add_edge_resolved(&mut self, caller: NodeIndex, callee: &str, kind: EdgeKind) {
let caller_file = self.inner[caller].file.display().to_string();
let require_callable = kind == EdgeKind::Calls;
if let Some((target, _grounds)) =
resolve_callee(&self.names, &caller_file, callee, require_callable)
{
self.inner.add_edge(caller, target, kind);
}
}
pub fn rebuild_name_index(&mut self) {
self.names = NameIndex::default();
for idx in self.inner.node_indices() {
let node = &self.inner[idx];
let file = node.file.display().to_string();
let name = node.name.clone();
let callable = matches!(node.kind, SymbolKind::Function | SymbolKind::Method);
self.names.insert(&file, &name, idx, callable);
}
}
pub fn build_from_file(file: &Path) -> Result<SymbolGraph> {
let source = std::fs::read_to_string(file)?;
let Some((lang, lang_tag)) = detect_language(file) else {
return Ok(SymbolGraph::default());
};
let symbols = extract_symbols(&source, lang.clone(), file);
let mut sorted: Vec<_> = symbols.iter().collect();
sorted.sort_by_key(|a| a.start_line);
let mut graph = SymbolGraph::default();
let placed: Vec<(&&crate::symgraph::symbol::Symbol, NodeIndex)> = sorted
.iter()
.map(|s| {
let idx = graph.add_node(SymbolNode {
file: s.file.clone(),
name: s.name.clone(),
kind: s.kind,
start_line: s.start_line,
});
(s, idx)
})
.collect();
let mut raw_edges: Vec<RawEdge> = Vec::new();
let mut parser = Parser::new();
if parser.set_language(&lang).is_ok()
&& let Some(tree) = parser.parse(&source, None)
{
let bytes = source.as_bytes();
for (sym, caller) in &placed {
if !matches!(sym.kind, SymbolKind::Function | SymbolKind::Method) {
continue;
}
if let Some(node) =
node_for_byte_range(tree.root_node(), sym.start_byte, sym.end_byte)
{
collect_calls(node, bytes, lang_tag, *caller, &mut raw_edges);
}
}
let file_stem = file
.file_stem()
.and_then(|s| s.to_str())
.unwrap_or("")
.to_string();
let mut stem_idx: Option<NodeIndex> = placed
.iter()
.find(|(s, _)| s.name == file_stem)
.map(|(_, i)| *i);
for sym in &symbols {
if !matches!(sym.kind, SymbolKind::Import) {
continue;
}
if stem_idx.is_none() && !file_stem.is_empty() {
stem_idx = Some(graph.add_node(SymbolNode {
file: file.to_path_buf(),
name: file_stem.clone(),
kind: SymbolKind::Unknown,
start_line: 0,
}));
}
if let Some(caller) = stem_idx {
raw_edges.push(RawEdge {
caller,
callee: sym.name.clone(),
kind: EdgeKind::Imports,
});
}
}
}
for e in raw_edges {
graph.add_edge_resolved(e.caller, &e.callee, e.kind);
}
Ok(graph)
}
pub fn build_from_registry(registry: &SymbolRegistry) -> Self {
let mut graph = SymbolGraph::default();
let entries: Vec<_> = registry.iter().collect();
let placed: Vec<NodeIndex> = entries
.iter()
.map(|(id, entry)| {
let node = SymbolNode {
file: entry
.assigned_file
.clone()
.unwrap_or_else(|| PathBuf::from("")),
name: bare_name(id.as_str()).to_string(),
kind: registry_kind_to_symbol_kind(&entry.kind),
start_line: 0,
};
graph.add_node_as(node, id.as_str())
})
.collect();
for ((_, entry), &caller) in entries.iter().zip(placed.iter()) {
for dep in &entry.dependencies {
graph.add_edge_resolved(caller, dep.as_str(), EdgeKind::Calls);
}
}
graph
}
pub fn resolve_symbol(&self, name: &str) -> SymbolMatch<'_> {
let hits = self.ranked_indices(name);
match hits.len() {
0 => SymbolMatch::NotFound,
1 => SymbolMatch::Unique(&self.inner[hits[0]]),
_ => SymbolMatch::Ambiguous {
chosen: &self.inner[hits[0]],
alternatives: hits[1..].iter().map(|&i| &self.inner[i]).collect(),
},
}
}
fn ranked_indices(&self, name: &str) -> Vec<NodeIndex> {
rank_matches(&self.names, name, |i| {
self.inner.edges_directed(i, Direction::Outgoing).count()
+ self.inner.edges_directed(i, Direction::Incoming).count()
})
}
fn idx_of(&self, name: &str) -> Option<NodeIndex> {
self.ranked_indices(name).first().copied()
}
pub fn callers_of(&self, name: &str) -> Vec<&SymbolNode> {
let Some(target) = self.idx_of(name) else {
return Vec::new();
};
let mut seen: HashSet<NodeIndex> = HashSet::new();
let mut out = Vec::new();
for er in self.inner.edges_directed(target, Direction::Incoming) {
if *er.weight() != EdgeKind::Calls {
continue;
}
let src = er.source();
if seen.insert(src) {
out.push(&self.inner[src]);
}
}
out
}
pub fn callees_of(&self, name: &str) -> Vec<&SymbolNode> {
let Some(source) = self.idx_of(name) else {
return Vec::new();
};
let mut seen: HashSet<NodeIndex> = HashSet::new();
let mut out = Vec::new();
for er in self.inner.edges_directed(source, Direction::Outgoing) {
if *er.weight() != EdgeKind::Calls {
continue;
}
let dst = er.target();
if seen.insert(dst) {
out.push(&self.inner[dst]);
}
}
out
}
pub fn context_for(&self, name: &str, depth: usize) -> Vec<&SymbolNode> {
if depth == 0 {
return Vec::new();
}
let Some(start) = self.idx_of(name) else {
return Vec::new();
};
let mut visited: HashSet<NodeIndex> = HashSet::new();
visited.insert(start);
let mut queue: VecDeque<(NodeIndex, usize)> = VecDeque::new();
queue.push_back((start, 0));
let mut out_idx: Vec<NodeIndex> = Vec::new();
while let Some((cur, d)) = queue.pop_front() {
if d >= depth {
continue;
}
for er in self.inner.edges_directed(cur, Direction::Outgoing) {
if *er.weight() != EdgeKind::Calls {
continue;
}
let next = er.target();
if visited.insert(next) {
out_idx.push(next);
queue.push_back((next, d + 1));
}
}
for er in self.inner.edges_directed(cur, Direction::Incoming) {
if *er.weight() != EdgeKind::Calls {
continue;
}
let next = er.source();
if visited.insert(next) {
out_idx.push(next);
queue.push_back((next, d + 1));
}
}
}
out_idx.into_iter().map(|i| &self.inner[i]).collect()
}
}
fn registry_kind_to_symbol_kind(k: &crate::symgraph::registry::SymbolKind) -> SymbolKind {
use crate::symgraph::registry::SymbolKind as R;
match k {
R::Function | R::Test | R::TestSuite => SymbolKind::Function,
R::Method => SymbolKind::Method,
R::Class => SymbolKind::Class,
R::Struct => SymbolKind::Struct,
R::Trait => SymbolKind::Trait,
R::Impl => SymbolKind::Impl,
R::Import => SymbolKind::Import,
R::TypeAlias => SymbolKind::TypeAlias,
R::Const => SymbolKind::Const,
R::Unknown => SymbolKind::Unknown,
}
}
fn node_for_byte_range<'a>(root: Node<'a>, start: usize, end: usize) -> Option<Node<'a>> {
if root.start_byte() == start && root.end_byte() == end {
return Some(root);
}
let mut cursor = root.walk();
for child in root.children(&mut cursor) {
if child.start_byte() <= start
&& child.end_byte() >= end
&& let Some(found) = node_for_byte_range(child, start, end)
{
return Some(found);
}
}
None
}
fn collect_calls(node: Node, bytes: &[u8], lang: &str, caller: NodeIndex, out: &mut Vec<RawEdge>) {
let kind = node.kind();
let is_call = match lang {
"rust" | "javascript" | "go" => kind == "call_expression",
"python" => kind == "call",
_ => false,
};
if is_call && let Some(callee) = call_target_name(node, bytes, lang) {
out.push(RawEdge {
caller,
callee,
kind: EdgeKind::Calls,
});
}
let mut cursor = node.walk();
for child in node.children(&mut cursor) {
collect_calls(child, bytes, lang, caller, out);
}
}
fn call_target_name(node: Node, bytes: &[u8], lang: &str) -> Option<String> {
let func_node = match lang {
"rust" | "javascript" | "go" => node
.child_by_field_name("function")
.or_else(|| node.child(0)),
"python" => node
.child_by_field_name("function")
.or_else(|| node.child(0)),
_ => None,
}?;
let raw = func_node.utf8_text(bytes).ok()?;
let last = raw.rsplit("::").next().unwrap_or(raw);
let last = last.rsplit('.').next().unwrap_or(last);
Some(last.trim().to_string())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::symgraph::registry::{SymbolEntry, SymbolId, SymbolKind as RKind, SymbolRegistry};
use std::io::Write;
use tempfile::NamedTempFile;
fn entry(id: &str, file: &str, deps: &[&str]) -> SymbolEntry {
let mut e = SymbolEntry::new(
SymbolId(id.to_string()),
RKind::Function,
format!("fn {}() {{}}", bare_name(id)),
"rust",
);
e.assigned_file = Some(PathBuf::from(file));
e.dependencies = deps.iter().map(|d| SymbolId((*d).to_string())).collect();
e
}
fn registry_of(entries: Vec<SymbolEntry>) -> SymbolRegistry {
let mut reg = SymbolRegistry::new(PathBuf::from("/proj"));
for e in entries {
reg.insert(e);
}
reg
}
fn callee_files(g: &SymbolGraph, caller: &str) -> Vec<String> {
g.callees_of(caller)
.iter()
.map(|n| n.file.display().to_string())
.collect()
}
#[test]
fn same_file_callee_wins_over_an_earlier_registered_twin() {
let g = SymbolGraph::build_from_registry(®istry_of(vec![
entry("agents::stamp::write", "crates/agents/src/stamp.rs", &[]),
entry("search::store::write", "crates/search/src/store.rs", &[]),
entry(
"search::store::upsert",
"crates/search/src/store.rs",
&["write"],
),
]));
assert_eq!(
callee_files(&g, "upsert"),
vec!["crates/search/src/store.rs".to_string()],
);
}
#[test]
fn bare_name_collision_across_crates_creates_no_edge() {
let g = SymbolGraph::build_from_registry(®istry_of(vec![
entry("a::alpha::run", "crates/a/src/lib.rs", &[]),
entry("b::beta::run", "crates/b/src/lib.rs", &[]),
entry("c::gamma::start", "crates/c/src/lib.rs", &["run"]),
]));
assert!(
g.callees_of("start").is_empty(),
"ambiguous callee resolved anyway: {:?}",
callee_files(&g, "start"),
);
}
#[test]
fn directory_scope_beats_a_distant_twin() {
let g = SymbolGraph::build_from_registry(®istry_of(vec![
entry("far::helper", "crates/far/src/util.rs", &[]),
entry("near::helper", "crates/near/src/util.rs", &[]),
entry("near::start", "crates/near/src/lib.rs", &["helper"]),
]));
assert_eq!(
callee_files(&g, "start"),
vec!["crates/near/src/util.rs".to_string()],
);
}
#[test]
fn corpus_unique_name_still_resolves_across_files() {
let g = SymbolGraph::build_from_registry(®istry_of(vec![
entry("a::alpha::only_one", "crates/a/src/lib.rs", &[]),
entry("c::gamma::start", "crates/c/src/lib.rs", &["only_one"]),
]));
assert_eq!(
callee_files(&g, "start"),
vec!["crates/a/src/lib.rs".to_string()],
);
}
#[test]
fn cross_language_twin_is_not_an_edge() {
let g = SymbolGraph::build_from_registry(®istry_of(vec![
entry("ui::chat::get", "ui/src/chatStream.ts", &[]),
entry("a::alpha::start", "crates/a/src/lib.rs", &["get"]),
]));
assert!(
g.callees_of("start").is_empty(),
"cross-language callee resolved: {:?}",
callee_files(&g, "start"),
);
}
#[test]
fn path_qualified_name_anchors_on_the_file_it_names() {
let g = SymbolGraph::build_from_registry(®istry_of(vec![
entry("agents::stamp::write", "crates/agents/src/stamp.rs", &[]),
entry("search::store::write", "crates/search/src/store.rs", &[]),
entry(
"search::store::upsert",
"crates/search/src/store.rs",
&["write"],
),
]));
match g.resolve_symbol("crates/agents/src/stamp.rs::write") {
SymbolMatch::Unique(n) => assert_eq!(
n.file.display().to_string(),
"crates/agents/src/stamp.rs",
"anchored on the wrong file",
),
other => panic!("expected Unique on stamp.rs, got {other:?}"),
}
assert!(
g.callers_of("crates/agents/src/stamp.rs::write").is_empty(),
"anchored on the store.rs twin: {:?}",
g.callers_of("crates/agents/src/stamp.rs::write"),
);
let callers = g.callers_of("src/store.rs::write");
assert_eq!(callers.len(), 1, "got {callers:?}");
assert_eq!(callers[0].name, "upsert");
}
#[test]
fn a_call_never_lands_on_a_container_in_the_callers_file() {
let mut container = entry("m::Helper", "crates/a/src/lib.rs", &[]);
container.kind = RKind::Struct;
let g = SymbolGraph::build_from_registry(®istry_of(vec![
container,
entry("m::start", "crates/a/src/lib.rs", &["m::Helper"]),
]));
assert!(
g.callees_of("start").is_empty(),
"call resolved to a container: {:?}",
callee_files(&g, "start"),
);
}
#[test]
fn sibling_extensions_of_one_language_stay_ambiguous() {
let g = SymbolGraph::build_from_registry(®istry_of(vec![
entry("lib::a::get", "ui/lib/a.ts", &[]),
entry("widgets::b::get", "ui/widgets/b.tsx", &[]),
entry("app::main::start", "ui/app/main.ts", &["get"]),
]));
assert!(
g.callees_of("start").is_empty(),
"sibling-extension twin resolved: {:?}",
callee_files(&g, "start"),
);
}
#[test]
fn ambiguous_bare_name_reports_every_candidate() {
let g = SymbolGraph::build_from_registry(®istry_of(vec![
entry("agents::stamp::write", "crates/agents/src/stamp.rs", &[]),
entry("search::store::write", "crates/search/src/store.rs", &[]),
entry(
"search::store::upsert",
"crates/search/src/store.rs",
&["write"],
),
]));
match g.resolve_symbol("write") {
SymbolMatch::Ambiguous {
chosen,
alternatives,
} => {
assert_eq!(
chosen.file.display().to_string(),
"crates/search/src/store.rs"
);
assert_eq!(alternatives.len(), 1);
}
other => panic!("expected Ambiguous, got {other:?}"),
}
assert!(matches!(
g.resolve_symbol("no_such_symbol"),
SymbolMatch::NotFound
));
}
#[test]
fn build_from_registry_smoke() {
use std::collections::BTreeSet;
let tmp = tempfile::TempDir::new().unwrap();
let mut reg = SymbolRegistry::new(tmp.path().to_path_buf());
let mut caller = SymbolEntry::new(
SymbolId::new("m", "caller"),
RKind::Function,
"fn caller() { callee(); }".into(),
"rust",
);
let mut deps = BTreeSet::new();
deps.insert(SymbolId("callee".into()));
caller.dependencies = deps;
reg.insert(caller);
let callee = SymbolEntry::new(
SymbolId::new("m", "callee"),
RKind::Function,
"fn callee() {}".into(),
"rust",
);
reg.insert(callee);
let g = SymbolGraph::build_from_registry(®);
assert_eq!(g.node_count(), 2);
let names: Vec<&str> = g.nodes().iter().map(|n| n.name.as_str()).collect();
assert!(names.contains(&"caller"));
assert!(names.contains(&"callee"));
let edges = g.edges();
assert!(
edges
.iter()
.any(|e| e.from == "caller" && e.to == "callee" && e.kind == EdgeKind::Calls),
"expected caller -> callee Calls edge, got {edges:?}",
);
}
#[test]
fn kg_calls_edge_between_two_functions() {
let src = "fn caller() { callee(); }\n\nfn callee() {}\n";
let mut tmp = NamedTempFile::new().unwrap();
tmp.write_all(src.as_bytes()).unwrap();
let p = tmp.path().with_extension("rs");
std::fs::copy(tmp.path(), &p).unwrap();
let g = SymbolGraph::build_from_file(&p).unwrap();
let _ = std::fs::remove_file(&p);
let edges = g.edges();
let calls: Vec<&SymbolEdge> = edges.iter().filter(|e| e.kind == EdgeKind::Calls).collect();
assert!(
calls.iter().any(|e| e.from == "caller" && e.to == "callee"),
"expected caller -> callee Calls edge, got {edges:?}",
);
assert!(!g.callers_of("callee").is_empty());
assert!(!g.callees_of("caller").is_empty());
}
}