use crate::code_tree::models::{FileInfo, FunctionInfo, TypeRelationship};
use rayon::prelude::*;
use std::collections::{HashMap, HashSet};
pub struct CallEdge {
pub caller: String,
pub callee: String,
pub call_lines: String,
pub call_count: i64,
}
#[derive(Debug, Clone, Copy, Default, serde::Serialize)]
pub struct CallResolutionStats {
pub total_calls: u64,
pub excluded_noise: u64,
pub no_candidate: u64,
pub ambiguous_dropped: u64,
pub resolved_call_sites: u64,
pub resolved_edges: u64,
pub resolved_via_inheritance: u64,
}
#[derive(Debug, Clone, Copy, Default)]
struct Counts {
total: u64,
excluded: u64,
no_candidate: u64,
ambiguous: u64,
resolved: u64,
inherited: u64,
}
fn short_type_name(name: &str) -> &str {
let mut cut = 0usize;
for sep in ["::", ".", "/"] {
if let Some(i) = name.rfind(sep) {
let after = i + sep.len();
if after > cut {
cut = after;
}
}
}
&name[cut..]
}
fn build_ancestor_map(rels: &[TypeRelationship]) -> HashMap<&str, HashSet<&str>> {
let mut parents: HashMap<&str, Vec<&str>> = HashMap::new();
for tr in rels {
if tr.relationship == "extends" || tr.relationship == "implements" {
if let Some(tgt) = tr.target_type.as_deref() {
parents
.entry(short_type_name(&tr.source_type))
.or_default()
.push(short_type_name(tgt));
}
}
}
let mut out: HashMap<&str, HashSet<&str>> = HashMap::with_capacity(parents.len());
for &child in parents.keys() {
let mut seen: HashSet<&str> = HashSet::new();
let mut stack: Vec<&str> = parents.get(child).cloned().unwrap_or_default();
while let Some(p) = stack.pop() {
if p != child && seen.insert(p) {
if let Some(gp) = parents.get(p) {
stack.extend(gp.iter().copied());
}
}
}
out.insert(child, seen);
}
out
}
type FnMatchResult<'a> = (Vec<(&'a str, &'a str, u32)>, Counts);
fn qname_starts_with_any(qname: &str, scopes: &[String]) -> bool {
for scope in scopes {
if scope.is_empty() {
continue;
}
if qname.len() > scope.len()
&& qname.starts_with(scope.as_str())
&& (qname.as_bytes()[scope.len()] == b'.'
|| (qname.len() > scope.len() + 1 && &qname[scope.len()..scope.len() + 2] == "::"))
{
return true;
}
}
false
}
fn infer_lang_group(qname: &str) -> &'static str {
if qname.contains("::") {
"rust_cpp"
} else if qname.contains('/') {
"go_ts_js"
} else {
"python_java"
}
}
pub fn build_call_edges(
functions: &[FunctionInfo],
files: &[FileInfo],
excluded_names: &std::collections::HashSet<&str>,
max_targets: usize,
type_relationships: &[TypeRelationship],
) -> (Vec<CallEdge>, CallResolutionStats) {
let verbose = std::env::var_os("KGLITE_CODE_TREE_VERBOSE").is_some();
let t0 = std::time::Instant::now();
let mut name_lookup: HashMap<&str, Vec<&str>> = HashMap::new();
for fn_info in functions {
name_lookup
.entry(fn_info.name.as_str())
.or_default()
.push(fn_info.qualified_name.as_str());
}
let mut qname_to_owner: HashMap<&str, &str> = HashMap::new();
let mut qname_to_prefix: HashMap<&str, &str> = HashMap::new();
for fn_info in functions {
let qn = fn_info.qualified_name.as_str();
for sep in ["::", ".", "/"] {
if let Some(idx) = qn.rfind(sep) {
let owner_path = &qn[..idx];
qname_to_prefix.insert(qn, owner_path);
let mut short = owner_path;
for sep2 in ["::", ".", "/"] {
if let Some(i2) = owner_path.rfind(sep2) {
short = &owner_path[i2 + sep2.len()..];
break;
}
}
qname_to_owner.insert(qn, short);
break;
}
}
}
let qname_to_file: HashMap<&str, &str> = functions
.iter()
.map(|f| (f.qualified_name.as_str(), f.file_path.as_str()))
.collect();
let file_imports: HashMap<&str, &Vec<String>> = files
.iter()
.filter(|f| !f.imports.is_empty())
.map(|f| (f.path.as_str(), &f.imports))
.collect();
let ancestors = build_ancestor_map(type_relationships);
if verbose {
eprintln!(
"[calls] lookup build: {:.3}s",
t0.elapsed().as_secs_f64()
);
}
let t_match = std::time::Instant::now();
let per_fn: Vec<FnMatchResult> = functions
.par_iter()
.map(|fn_info| {
let caller_qn = fn_info.qualified_name.as_str();
let caller_lang = infer_lang_group(caller_qn);
let caller_prefix = qname_to_prefix.get(caller_qn).copied();
let caller_owner = qname_to_owner.get(caller_qn).copied();
let caller_file = fn_info.file_path.as_str();
let mut out: Vec<(&str, &str, u32)> = Vec::new();
let mut counts = Counts::default();
for (called_name, line) in &fn_info.calls {
counts.total += 1;
let (explicit_hint, method_name) = match called_name.rfind('.') {
Some(idx) => (Some(&called_name[..idx]), &called_name[idx + 1..]),
None => (None, called_name.as_str()),
};
if excluded_names.contains(method_name) {
counts.excluded += 1;
continue;
}
let Some(candidates) = name_lookup.get(method_name) else {
counts.no_candidate += 1;
continue;
};
if candidates.len() == 1 {
counts.resolved += 1;
let target = candidates[0];
if target != caller_qn {
out.push((caller_qn, target, *line));
}
continue;
}
let mut targets: &[&str] = candidates.as_slice();
let mut filtered: Vec<&str>;
let implicit_hint = if explicit_hint.is_none() {
caller_owner
} else {
None
};
let mut owner_hint_hit = false;
if let Some(hint) = explicit_hint.or(implicit_hint) {
filtered = targets
.iter()
.copied()
.filter(|t| qname_to_owner.get(t).copied() == Some(hint))
.collect();
if !filtered.is_empty() {
targets = &filtered[..];
owner_hint_hit = true;
}
}
if !owner_hint_hit && targets.len() > 1 {
if let Some(owner) = implicit_hint {
if let Some(anc) = ancestors.get(owner) {
let inh: Vec<&str> = candidates
.iter()
.copied()
.filter(|t| {
qname_to_owner
.get(t)
.copied()
.is_some_and(|o| anc.contains(o))
})
.collect();
if inh.len() == 1 {
counts.resolved += 1;
counts.inherited += 1;
if inh[0] != caller_qn {
out.push((caller_qn, inh[0], *line));
}
continue;
} else if !inh.is_empty() {
filtered = inh;
targets = &filtered[..];
}
}
}
}
if targets.len() > 1 {
if let Some(prefix) = caller_prefix {
let narrowed: Vec<&str> = targets
.iter()
.copied()
.filter(|t| qname_to_prefix.get(t).copied() == Some(prefix))
.collect();
if !narrowed.is_empty() {
filtered = narrowed;
targets = &filtered[..];
}
}
}
if targets.len() > 1 {
if let Some(imports) = file_imports.get(caller_file) {
let narrowed: Vec<&str> = targets
.iter()
.copied()
.filter(|t| qname_starts_with_any(t, imports))
.collect();
if !narrowed.is_empty() {
filtered = narrowed;
targets = &filtered[..];
}
}
}
if targets.len() > 1 {
let narrowed: Vec<&str> = targets
.iter()
.copied()
.filter(|t| qname_to_file.get(t).copied() == Some(caller_file))
.collect();
if !narrowed.is_empty() {
filtered = narrowed;
targets = &filtered[..];
}
}
if targets.len() > 1 {
let narrowed: Vec<&str> = targets
.iter()
.copied()
.filter(|t| infer_lang_group(t) == caller_lang)
.collect();
if !narrowed.is_empty() {
filtered = narrowed;
targets = &filtered[..];
}
}
if targets.len() > max_targets {
counts.ambiguous += 1;
continue;
}
counts.resolved += 1;
for &target in targets {
if target != caller_qn {
out.push((caller_qn, target, *line));
}
}
}
(out, counts)
})
.collect();
let mut stats = CallResolutionStats::default();
for (_, c) in &per_fn {
stats.total_calls += c.total;
stats.excluded_noise += c.excluded;
stats.no_candidate += c.no_candidate;
stats.ambiguous_dropped += c.ambiguous;
stats.resolved_call_sites += c.resolved;
stats.resolved_via_inheritance += c.inherited;
}
let total: usize = per_fn.iter().map(|(v, _)| v.len()).sum();
let mut seen: HashMap<(&str, &str), Vec<u32>> = HashMap::with_capacity(total);
for (edges, _) in per_fn {
for (caller, callee, line) in edges {
seen.entry((caller, callee)).or_default().push(line);
}
}
if verbose {
eprintln!(
"[calls] match loop: {:.3}s ({} entries)",
t_match.elapsed().as_secs_f64(),
seen.len()
);
}
let t_out = std::time::Instant::now();
let mut keys: Vec<(&str, &str)> = seen.keys().copied().collect();
keys.sort_unstable();
let result: Vec<CallEdge> = keys
.into_iter()
.map(|(caller, callee)| {
let mut lines = seen.remove(&(caller, callee)).unwrap_or_default();
lines.sort_unstable();
lines.dedup();
let count = lines.len() as i64;
let mut call_lines = String::with_capacity(lines.len() * 4);
for (i, l) in lines.iter().enumerate() {
if i > 0 {
call_lines.push(',');
}
use std::fmt::Write;
let _ = write!(call_lines, "{}", l);
}
CallEdge {
caller: caller.to_string(),
callee: callee.to_string(),
call_lines,
call_count: count,
}
})
.collect();
stats.resolved_edges = result.len() as u64;
if verbose {
eprintln!(
"[calls] output build: {:.3}s ({} edges, {}/{} call sites resolved)",
t_out.elapsed().as_secs_f64(),
stats.resolved_edges,
stats.resolved_call_sites,
stats.total_calls.saturating_sub(stats.excluded_noise),
);
}
(result, stats)
}
#[cfg(test)]
mod stats_tests {
use super::*;
use crate::code_tree::models::{FileInfo, FunctionInfo, TypeRelationship};
fn func(qn: &str, file: &str, calls: &[(&str, u32)]) -> FunctionInfo {
FunctionInfo {
name: qn.rsplit(['.', ':']).next().unwrap_or(qn).to_string(),
qualified_name: qn.to_string(),
file_path: file.to_string(),
calls: calls.iter().map(|(n, l)| (n.to_string(), *l)).collect(),
..Default::default()
}
}
#[test]
fn stats_classify_every_call_site() {
let functions = vec![
func(
"a.foo",
"a.py",
&[("bar", 1), ("bar", 2), ("external_thing", 3), ("noisy", 4)],
),
func("a.bar", "a.py", &[]),
];
let files = vec![FileInfo {
path: "a.py".into(),
..Default::default()
}];
let mut noise = std::collections::HashSet::new();
noise.insert("noisy");
let (edges, stats) = build_call_edges(&functions, &files, &noise, 5, &[]);
assert_eq!(stats.total_calls, 4);
assert_eq!(stats.excluded_noise, 1);
assert_eq!(stats.no_candidate, 1);
assert_eq!(stats.ambiguous_dropped, 0);
assert_eq!(stats.resolved_call_sites, 2); assert_eq!(stats.resolved_edges, 1); assert_eq!(edges.len(), 1);
}
#[test]
fn inheritance_tier_resolves_self_call_to_ancestor_method() {
let files = vec![FileInfo {
path: "a.py".into(),
..Default::default()
}];
let functions = vec![
func("mod.Base.m", "a.py", &[]),
func("mod.Other.m", "a.py", &[]),
func("mod.Derived.caller", "a.py", &[("m", 1)]),
];
let rels = vec![TypeRelationship {
source_type: "mod.Derived".into(),
target_type: Some("mod.Base".into()),
relationship: "extends".into(),
methods: vec![],
}];
let noise = std::collections::HashSet::new();
let (edges, stats) = build_call_edges(&functions, &files, &noise, 5, &rels);
assert_eq!(stats.resolved_via_inheritance, 1);
let pairs: Vec<(&str, &str)> = edges
.iter()
.map(|e| (e.caller.as_str(), e.callee.as_str()))
.collect();
assert!(pairs.contains(&("mod.Derived.caller", "mod.Base.m")));
assert!(!pairs.iter().any(|(_, callee)| *callee == "mod.Other.m"));
}
#[test]
fn no_type_relationships_means_no_inheritance_resolution() {
let files = vec![FileInfo {
path: "a.py".into(),
..Default::default()
}];
let functions = vec![
func("mod.Base.m", "a.py", &[]),
func("mod.Other.m", "a.py", &[]),
func("mod.Derived.caller", "a.py", &[("m", 1)]),
];
let noise = std::collections::HashSet::new();
let (_edges, stats) = build_call_edges(&functions, &files, &noise, 5, &[]);
assert_eq!(stats.resolved_via_inheritance, 0);
}
}