#![warn(missing_docs)]
use crate::graph::pdg::{Edge, EdgeMetadata, EdgeType, Node, NodeType, ProgramDependenceGraph};
use crate::parse::prelude::{FlowChannel, FlowFact, ImportInfo, SignatureInfo};
use regex::Regex;
use std::collections::{HashMap, HashSet};
use std::path::Path;
use std::sync::Arc;
type LocalNodeIds = HashMap<String, Vec<crate::graph::pdg::NodeId>>;
pub fn extract_pdg_from_signatures(
signatures: Vec<SignatureInfo>,
source_code: &[u8],
file_path: &str,
language: &str,
) -> ProgramDependenceGraph {
let mut pdg = ProgramDependenceGraph::new();
let mut node_ids: HashMap<String, crate::graph::pdg::NodeId> = HashMap::new();
let mut local_node_ids = LocalNodeIds::new();
let mut seen_qnames = HashSet::new();
let duplicate_qnames: HashSet<&str> = signatures
.iter()
.filter_map(|sig| {
(!seen_qnames.insert(sig.qualified_name.as_str()))
.then_some(sig.qualified_name.as_str())
})
.collect();
for sig in &signatures {
let mut node = signature_to_node(sig, file_path, language);
if duplicate_qnames.contains(sig.qualified_name.as_str()) {
node.id = format!(
"{}:{}@{}..{}",
file_path, sig.qualified_name, sig.byte_range.0, sig.byte_range.1
);
}
let nid = pdg.add_node(node);
local_node_ids
.entry(sig.qualified_name.clone())
.or_default()
.push(nid);
node_ids.entry(sig.qualified_name.clone()).or_insert(nid);
}
let containment = infer_class_nodes_and_containment(
&signatures,
&mut pdg,
&mut node_ids,
&local_node_ids,
file_path,
language,
);
pdg.add_containment_edges(containment);
let data_edges = extract_data_flow_edges_for_nodes(&signatures, &local_node_ids);
pdg.add_data_flow_edges(data_edges);
let inheritance = extract_inheritance_edges(&signatures, &node_ids);
pdg.add_inheritance_edges(inheritance);
let call_edges = extract_call_edges_for_nodes(&signatures, &local_node_ids);
pdg.add_call_edges(call_edges);
extract_flow_edges(&signatures, &node_ids, &mut pdg);
let import_edges = extract_import_edges(
&signatures,
&node_ids,
&mut pdg,
file_path,
language,
source_code,
);
pdg.add_import_edges(import_edges);
pdg
}
fn infer_class_nodes_and_containment(
signatures: &[SignatureInfo],
pdg: &mut ProgramDependenceGraph,
node_ids: &mut HashMap<String, crate::graph::pdg::NodeId>,
local_node_ids: &LocalNodeIds,
file_path: &str,
language: &str,
) -> Vec<(crate::graph::pdg::NodeId, crate::graph::pdg::NodeId)> {
let mut class_methods: HashMap<String, HashSet<crate::graph::pdg::NodeId>> = HashMap::new();
for sig in signatures {
if !sig.is_method {
continue;
}
let normalized = normalize_symbol(&sig.qualified_name);
if let Some(dot_pos) = normalized.rfind('.') {
let class_prefix = normalized[..dot_pos].to_string();
if let Some(method_ids) = local_node_ids.get(&sig.qualified_name) {
class_methods
.entry(class_prefix)
.or_default()
.extend(method_ids);
}
}
}
let mut containment = Vec::new();
let mut class_names: Vec<&String> = class_methods.keys().collect();
class_names.sort();
for class_name in class_names {
let method_nids = &class_methods[class_name];
let already_exists = node_ids.contains_key(class_name)
|| node_ids.keys().any(|k| normalize_symbol(k) == *class_name);
if already_exists {
if let Some(&class_nid) = node_ids.get(class_name).or_else(|| {
node_ids
.iter()
.find(|(k, _)| normalize_symbol(k) == *class_name)
.map(|(_, v)| v)
}) {
for &mnid in method_nids {
containment.push((class_nid, mnid));
}
}
continue;
}
let (min_start, max_end) = method_nids.iter().fold((usize::MAX, 0), |(mn, mx), &mnid| {
pdg.get_node(mnid)
.map(|n| (mn.min(n.byte_range.0), mx.max(n.byte_range.1)))
.unwrap_or((mn, mx))
});
let short_name = class_name
.rsplit('.')
.next()
.unwrap_or(class_name)
.to_string();
let class_complexity: u32 = method_nids
.iter()
.filter_map(|&mnid| pdg.get_node(mnid))
.map(|node| node.complexity)
.sum();
let class_node = Node {
id: format!("{}:{}", file_path, class_name),
node_type: NodeType::Class,
name: short_name,
file_path: Arc::from(file_path),
byte_range: (
if min_start == usize::MAX {
0
} else {
min_start
},
max_end,
),
complexity: if class_complexity > 0 {
class_complexity
} else {
method_nids.len() as u32
},
language: language.to_string(),
};
let class_nid = pdg.add_node(class_node);
node_ids.insert(class_name.clone(), class_nid);
for &mnid in method_nids {
containment.push((class_nid, mnid));
}
}
containment
}
const EXCLUDED_TYPES: &[&str] = &[
"str",
"string",
"String",
"&str",
"int",
"i32",
"i64",
"u32",
"u64",
"usize",
"f32",
"f64",
"bool",
"void",
"None",
"null",
"undefined",
"any",
"Any",
"object",
"Object",
"self",
"Self",
"cls",
"this",
"bytes",
"Bytes",
"Vec",
"List",
"list",
"dict",
"Dict",
"HashMap",
"Option",
"Result",
"Error",
"Exception",
"T",
"U",
"K",
"V",
];
fn is_excluded_type(t: &str) -> bool {
let base = t.split('<').next().unwrap_or(t).trim();
EXCLUDED_TYPES.contains(&base)
}
type DataFlowEdge = (
crate::graph::pdg::NodeId,
crate::graph::pdg::NodeId,
String,
f32,
);
type DataFlowSeen = HashSet<(crate::graph::pdg::NodeId, crate::graph::pdg::NodeId)>;
struct DataFlowIndexes<'a> {
producers: HashMap<String, Vec<&'a SignatureInfo>>,
consumers: HashMap<String, Vec<&'a SignatureInfo>>,
call_set: HashMap<String, HashSet<String>>,
by_normalized_name: HashMap<String, Vec<&'a SignatureInfo>>,
}
pub fn extract_data_flow_edges(
signatures: &[SignatureInfo],
node_ids: &HashMap<String, crate::graph::pdg::NodeId>,
) -> Vec<(
crate::graph::pdg::NodeId,
crate::graph::pdg::NodeId,
String,
f32,
)> {
let local_node_ids = node_ids
.iter()
.map(|(name, &id)| (name.clone(), vec![id]))
.collect();
extract_data_flow_edges_for_nodes(signatures, &local_node_ids)
}
fn extract_data_flow_edges_for_nodes(
signatures: &[SignatureInfo],
node_ids: &LocalNodeIds,
) -> Vec<DataFlowEdge> {
let indexes = build_data_flow_indexes(signatures);
let mut edges = Vec::new();
let mut seen = HashSet::new();
add_return_to_parameter_edges(&indexes, node_ids, &mut edges, &mut seen);
add_shared_return_call_edges(&indexes, node_ids, &mut edges, &mut seen);
add_shared_parameter_call_edges(signatures, &indexes, node_ids, &mut edges, &mut seen);
edges
}
fn build_data_flow_indexes(signatures: &[SignatureInfo]) -> DataFlowIndexes<'_> {
let mut indexes = DataFlowIndexes {
producers: HashMap::new(),
consumers: HashMap::new(),
call_set: HashMap::new(),
by_normalized_name: HashMap::new(),
};
for sig in signatures {
if let Some(ret) = &sig.return_type {
let norm = normalize_type_name(ret);
if !norm.is_empty() && !is_excluded_type(&norm) {
indexes.producers.entry(norm).or_default().push(sig);
}
}
for param in &sig.parameters {
if let Some(t) = ¶m.type_annotation {
let norm = normalize_type_name(t);
if !norm.is_empty() && !is_excluded_type(&norm) {
indexes.consumers.entry(norm).or_default().push(sig);
}
}
}
let calls = sig.calls.iter().map(|c| normalize_symbol(c)).collect();
indexes
.call_set
.insert(normalize_symbol(&sig.qualified_name), calls);
indexes
.by_normalized_name
.entry(normalize_symbol(&sig.qualified_name))
.or_default()
.push(sig);
}
indexes
}
fn add_return_to_parameter_edges(
indexes: &DataFlowIndexes<'_>,
node_ids: &LocalNodeIds,
edges: &mut Vec<DataFlowEdge>,
seen: &mut DataFlowSeen,
) {
for (type_name, producer_sigs) in &indexes.producers {
if let Some(consumer_sigs) = indexes.consumers.get(type_name) {
for prod in producer_sigs {
for cons in consumer_sigs {
if prod.qualified_name == cons.qualified_name {
continue;
}
let (Some(from_ids), Some(to_ids)) = (
node_ids.get(&prod.qualified_name),
node_ids.get(&cons.qualified_name),
) else {
continue;
};
for &from in from_ids {
for &to in to_ids {
if seen.insert((from, to)) {
edges.push((from, to, type_name.clone(), 0.85));
}
}
}
}
}
}
}
}
fn add_shared_return_call_edges(
indexes: &DataFlowIndexes<'_>,
node_ids: &LocalNodeIds,
edges: &mut Vec<DataFlowEdge>,
seen: &mut DataFlowSeen,
) {
for (type_name, ret_sigs) in &indexes.producers {
if ret_sigs.len() < 2 {
continue;
}
for i in 0..ret_sigs.len() {
for j in 0..ret_sigs.len() {
if i == j {
continue;
}
let a = ret_sigs[i];
let b = ret_sigs[j];
let a_norm = normalize_symbol(&a.qualified_name);
let b_norm = normalize_symbol(&b.qualified_name);
let a_calls_b = indexes
.call_set
.get(&a_norm)
.map(|s| s.contains(&b_norm))
.unwrap_or(false);
if a_calls_b {
let (Some(from_ids), Some(to_ids)) = (
node_ids.get(&a.qualified_name),
node_ids.get(&b.qualified_name),
) else {
continue;
};
for &from in from_ids {
for &to in to_ids {
if seen.insert((from, to)) {
edges.push((from, to, format!("ret:{}", type_name), 0.65));
}
}
}
}
}
}
}
}
fn add_shared_parameter_call_edges(
signatures: &[SignatureInfo],
indexes: &DataFlowIndexes<'_>,
node_ids: &LocalNodeIds,
edges: &mut Vec<DataFlowEdge>,
seen: &mut DataFlowSeen,
) {
for sig_a in signatures {
let a_norm = normalize_symbol(&sig_a.qualified_name);
let Some(a_calls) = indexes.call_set.get(&a_norm) else {
continue;
};
for called_norm in a_calls {
let Some(callee_sigs) = indexes.by_normalized_name.get(called_norm) else {
continue;
};
for sig_b in callee_sigs {
let a_types: HashSet<String> = sig_a
.parameters
.iter()
.filter_map(|p| p.type_annotation.as_ref())
.map(|t| normalize_type_name(t))
.filter(|t| !t.is_empty() && !is_excluded_type(t))
.collect();
let b_types: HashSet<String> = sig_b
.parameters
.iter()
.filter_map(|p| p.type_annotation.as_ref())
.map(|t| normalize_type_name(t))
.filter(|t| !t.is_empty() && !is_excluded_type(t))
.collect();
let mut shared: Vec<&String> = a_types.intersection(&b_types).collect();
shared.sort();
if shared.is_empty() {
continue;
}
let (Some(from_ids), Some(to_ids)) = (
node_ids.get(&sig_a.qualified_name),
node_ids.get(&sig_b.qualified_name),
) else {
continue;
};
for &from in from_ids {
for &to in to_ids {
if seen.insert((from, to)) {
edges.push((from, to, format!("param:{}", shared[0]), 0.45));
}
}
}
}
}
}
}
fn normalize_type_name(raw: &str) -> String {
let stripped = raw
.trim()
.trim_start_matches('&')
.trim_start_matches("mut ")
.trim();
if let Some(inner_start) = stripped.find('<') {
let inner = &stripped[inner_start + 1..];
let inner_end = inner.rfind('>').unwrap_or(inner.len());
let inner_type = inner[..inner_end].split(',').next().unwrap_or("").trim();
if !inner_type.is_empty() && !is_excluded_type(inner_type) {
return inner_type.to_string();
}
}
stripped.to_string()
}
const MIN_INHERITANCE_CONFIDENCE: f32 = 0.45;
const COMMON_METHOD_NAMES: &[&str] = &[
"new",
"init",
"__init__",
"constructor",
"create",
"build",
"toString",
"to_string",
"__str__",
"__repr__",
"equals",
"__eq__",
"hashCode",
"__hash__",
"clone",
"__clone__",
"copy",
"dispose",
"close",
"__del__",
"finalize",
"update",
"get",
"set",
"run",
"start",
"stop",
"execute",
];
fn is_common_method(name: &str) -> bool {
COMMON_METHOD_NAMES.contains(&name)
}
const ABSTRACT_BASE_PREFIXES: &[&str] = &["Abstract", "Base", "I"];
const ABSTRACT_BASE_SUFFIXES: &[&str] = &[
"Base",
"Mixin",
"Interface",
"Protocol",
"Trait",
"ABC",
"Abstract",
];
fn looks_like_abstract_base(class_name: &str) -> bool {
ABSTRACT_BASE_PREFIXES
.iter()
.any(|p| class_name.starts_with(p) && class_name.len() > p.len())
|| ABSTRACT_BASE_SUFFIXES
.iter()
.any(|s| class_name.ends_with(s) && class_name.len() > s.len())
}
#[derive(Debug, Default)]
struct InheritanceEvidence {
super_call_confidence: f32,
override_confidence: f32,
naming_confidence: f32,
nesting_confidence: f32,
}
impl InheritanceEvidence {
fn max_confidence(&self) -> f32 {
self.super_call_confidence
.max(self.override_confidence)
.max(self.naming_confidence)
.max(self.nesting_confidence)
}
}
fn group_methods_by_class(signatures: &[SignatureInfo]) -> HashMap<String, Vec<&SignatureInfo>> {
let mut class_methods: HashMap<String, Vec<&SignatureInfo>> = HashMap::new();
for sig in signatures {
if !sig.is_method {
continue;
}
let normalized = normalize_symbol(&sig.qualified_name);
let Some(dot_pos) = normalized.rfind('.') else {
continue;
};
class_methods
.entry(normalized[..dot_pos].to_string())
.or_default()
.push(sig);
}
class_methods
}
fn class_method_names<'class, 'sig>(
class_methods: &'class HashMap<String, Vec<&'sig SignatureInfo>>,
) -> HashMap<&'class str, HashSet<&'sig str>> {
class_methods
.iter()
.map(|(class_name, methods)| {
(
class_name.as_str(),
methods.iter().map(|sig| sig.name.as_str()).collect(),
)
})
.collect()
}
fn method_calls_parent(method: &SignatureInfo, parent_class: &str) -> bool {
let method_name = &method.name;
let super_patterns = [
format!("super.{}", method_name),
format!("super::{}", method_name),
format!("parent.{}", method_name),
format!("Base.{}", method_name),
format!("{}.{}", parent_class, method_name),
format!("{}::{}", parent_class, method_name),
];
method.calls.iter().any(|call| {
let norm_call = normalize_symbol(call);
super_patterns
.iter()
.any(|pat| norm_call.ends_with(&normalize_symbol(pat)))
|| norm_call.starts_with("super.")
|| norm_call.starts_with("super::")
|| norm_call.starts_with("parent.")
})
}
fn class_calls_parent(methods: &[&SignatureInfo], parent_class: &str) -> bool {
methods
.iter()
.any(|method| method_calls_parent(method, parent_class))
}
fn super_call_confidence(
methods_a: &[&SignatureInfo],
methods_b: &[&SignatureInfo],
cls_a: &str,
cls_b: &str,
) -> f32 {
if class_calls_parent(methods_a, cls_b) || class_calls_parent(methods_b, cls_a) {
0.90
} else {
0.0
}
}
fn override_confidence(shared_count: usize) -> f32 {
match shared_count {
0 | 1 => 0.0,
2 => 0.45,
3 => 0.60,
_ => 0.75,
}
}
fn naming_confidence(cls_a: &str, cls_b: &str, shared_count: usize) -> f32 {
if (looks_like_abstract_base(cls_a) || looks_like_abstract_base(cls_b)) && shared_count >= 1 {
0.50
} else {
0.0
}
}
fn is_qualified_class_prefix(prefix: &str, class_name: &str) -> bool {
class_name.starts_with(prefix)
&& class_name
.chars()
.nth(prefix.len())
.map(|character| character == '.')
.unwrap_or(false)
}
fn nesting_confidence(cls_a: &str, cls_b: &str) -> f32 {
if is_qualified_class_prefix(cls_a, cls_b) || is_qualified_class_prefix(cls_b, cls_a) {
0.70
} else {
0.0
}
}
fn inheritance_evidence(
cls_a: &str,
cls_b: &str,
methods_a: &[&SignatureInfo],
methods_b: &[&SignatureInfo],
class_method_names: &HashMap<&str, HashSet<&str>>,
) -> InheritanceEvidence {
let names_a = class_method_names.get(cls_a).cloned().unwrap_or_default();
let names_b = class_method_names.get(cls_b).cloned().unwrap_or_default();
let shared_count = names_a
.intersection(&names_b)
.filter(|&&name| !is_common_method(name))
.count();
let short_a = cls_a.rsplit('.').next().unwrap_or(cls_a);
let short_b = cls_b.rsplit('.').next().unwrap_or(cls_b);
InheritanceEvidence {
super_call_confidence: super_call_confidence(methods_a, methods_b, cls_a, cls_b),
override_confidence: override_confidence(shared_count),
naming_confidence: naming_confidence(short_a, short_b, shared_count),
nesting_confidence: nesting_confidence(cls_a, cls_b),
}
}
fn representative_class_node(
class_name: &str,
class_methods: &HashMap<String, Vec<&SignatureInfo>>,
node_ids: &HashMap<String, crate::graph::pdg::NodeId>,
) -> Option<crate::graph::pdg::NodeId> {
node_ids
.get(class_name)
.or_else(|| {
class_methods
.get(class_name)
.and_then(|methods| methods.first())
.and_then(|sig| node_ids.get(&sig.qualified_name))
})
.copied()
}
pub fn extract_inheritance_edges(
signatures: &[SignatureInfo],
node_ids: &HashMap<String, crate::graph::pdg::NodeId>,
) -> Vec<(crate::graph::pdg::NodeId, crate::graph::pdg::NodeId, f32)> {
let mut edges = Vec::new();
let class_methods = group_methods_by_class(signatures);
let class_names: Vec<&str> = class_methods.keys().map(String::as_str).collect();
if class_names.len() < 2 {
return edges;
}
let method_names_by_class = class_method_names(&class_methods);
for (i, &cls_a) in class_names.iter().enumerate() {
for &cls_b in &class_names[i + 1..] {
let methods_a = &class_methods[cls_a];
let methods_b = &class_methods[cls_b];
let evidence =
inheritance_evidence(cls_a, cls_b, methods_a, methods_b, &method_names_by_class);
let confidence = evidence.max_confidence();
if confidence < MIN_INHERITANCE_CONFIDENCE {
continue;
}
let short_a = cls_a.rsplit('.').next().unwrap_or(cls_a);
let short_b = cls_b.rsplit('.').next().unwrap_or(cls_b);
let (child_cls, parent_cls) = determine_inheritance_direction(
cls_a, cls_b, methods_a, methods_b, &evidence, short_a, short_b,
);
let child_nid = representative_class_node(child_cls, &class_methods, node_ids);
let parent_nid = representative_class_node(parent_cls, &class_methods, node_ids);
if let (Some(child_id), Some(parent_id)) = (child_nid, parent_nid) {
edges.push((child_id, parent_id, confidence));
}
}
}
edges
}
fn determine_inheritance_direction<'a>(
cls_a: &'a str,
cls_b: &'a str,
methods_a: &[&SignatureInfo],
_methods_b: &[&SignatureInfo],
evidence: &InheritanceEvidence,
short_a: &str,
short_b: &str,
) -> (&'a str, &'a str) {
if evidence.super_call_confidence > 0.0 {
let a_calls_super = methods_a.iter().any(|sig| {
sig.calls.iter().any(|c| {
let norm = normalize_symbol(c);
norm.starts_with("super.")
|| norm.starts_with("super::")
|| norm.starts_with("parent.")
|| norm.contains(cls_b)
})
});
if a_calls_super {
return (cls_a, cls_b);
}
return (cls_b, cls_a);
}
if looks_like_abstract_base(short_a) {
return (cls_b, cls_a);
}
if looks_like_abstract_base(short_b) {
return (cls_a, cls_b);
}
if cls_b.starts_with(cls_a) {
return (cls_b, cls_a);
}
if cls_a.starts_with(cls_b) {
return (cls_a, cls_b);
}
if cls_a.len() <= cls_b.len() {
(cls_b, cls_a)
} else {
(cls_a, cls_b)
}
}
fn caller_namespace(qualified_name: &str) -> Option<String> {
let normalized = normalize_symbol(qualified_name);
let segments: Vec<&str> = normalized.split('.').collect();
(segments.len() > 1).then(|| segments[..segments.len() - 1].join("."))
}
fn ordered_resolution_candidates(
call_target: &str,
alias_map: &HashMap<String, String>,
caller_ns: Option<&str>,
) -> Vec<String> {
let mut candidates = vec![call_target.to_string()];
let normalized = normalize_symbol(call_target);
let call_segments: Vec<&str> = normalized.split('.').filter(|s| !s.is_empty()).collect();
if let Some(first) = call_segments.first() {
if let Some(import_path) = alias_map.get(*first) {
let alias_target = if call_segments.len() == 1 {
import_path.clone()
} else {
format!("{}.{}", import_path, call_segments[1..].join("."))
};
candidates.push(alias_target);
}
if let Some(namespace) = caller_ns {
if matches!(
*first,
"self" | "this" | "super" | "Self" | "crate" | "base"
) {
let rest = call_segments[1..].join(".");
if !rest.is_empty() {
candidates.push(format!("{}.{}", namespace, rest));
}
} else if call_segments.len() == 1 {
candidates.push(format!("{}.{}", namespace, first));
}
}
}
candidates
}
fn local_call_targets(
candidates: &[String],
exact_map: &HashMap<String, Vec<crate::graph::pdg::NodeId>>,
last_map: &HashMap<String, Vec<crate::graph::pdg::NodeId>>,
suffix_map: &HashMap<String, Vec<crate::graph::pdg::NodeId>>,
) -> Vec<crate::graph::pdg::NodeId> {
let mut targets = Vec::new();
for candidate in candidates {
let normalized = normalize_symbol(candidate);
let segments: Vec<&str> = normalized.split('.').filter(|s| !s.is_empty()).collect();
if let Some(ids) = exact_map.get(&normalized) {
targets.extend(ids);
}
if let Some(last) = segments.last() {
if let Some(ids) = last_map.get(*last) {
targets.extend(ids);
}
}
for len in 2..=3_usize.min(segments.len()) {
let start = segments.len() - len;
let suffix = segments[start..].join(".");
if let Some(ids) = suffix_map.get(&suffix) {
targets.extend(ids);
}
}
}
targets.sort_unstable();
targets.dedup();
targets
}
fn type_node_target(
call_target: &str,
node_ids: &LocalNodeIds,
last_map: &HashMap<String, Vec<crate::graph::pdg::NodeId>>,
) -> Option<crate::graph::pdg::NodeId> {
let callee_name = normalize_symbol(call_target);
let (scoped_prefix, _) = callee_name.rsplit_once('.')?;
let bare_type = scoped_prefix.rsplit('.').next().unwrap_or(scoped_prefix);
if !bare_type.chars().next().is_some_and(|c| c.is_uppercase()) {
return None;
}
node_ids
.get(scoped_prefix)
.and_then(|ids| ids.first())
.or_else(|| node_ids.get(bare_type).and_then(|ids| ids.first()))
.or_else(|| last_map.get(bare_type).and_then(|ids| ids.first()))
.copied()
}
pub fn extract_call_edges(
signatures: &[SignatureInfo],
node_ids: &HashMap<String, crate::graph::pdg::NodeId>,
) -> Vec<(crate::graph::pdg::NodeId, crate::graph::pdg::NodeId)> {
let local_node_ids = node_ids
.iter()
.map(|(name, &id)| (name.clone(), vec![id]))
.collect();
extract_call_edges_for_nodes(signatures, &local_node_ids)
}
fn extract_call_edges_for_nodes(
signatures: &[SignatureInfo],
node_ids: &LocalNodeIds,
) -> Vec<(crate::graph::pdg::NodeId, crate::graph::pdg::NodeId)> {
let mut edges = Vec::new();
let mut seen = HashSet::new();
let mut exact_map: HashMap<String, Vec<crate::graph::pdg::NodeId>> = HashMap::new();
let mut last_map: HashMap<String, Vec<crate::graph::pdg::NodeId>> = HashMap::new();
let mut suffix_map: HashMap<String, Vec<crate::graph::pdg::NodeId>> = HashMap::new();
for signature in signatures {
if let Some(ids) = node_ids.get(&signature.qualified_name) {
let normalized = normalize_symbol(&signature.qualified_name);
let segments: Vec<&str> = normalized.split('.').filter(|s| !s.is_empty()).collect();
exact_map.entry(normalized.clone()).or_default().extend(ids);
if let Some(last) = segments.last() {
last_map.entry((*last).to_string()).or_default().extend(ids);
}
for len in 2..=3_usize.min(segments.len()) {
let start = segments.len() - len;
suffix_map
.entry(segments[start..].join("."))
.or_default()
.extend(ids);
}
}
}
for signature in signatures {
let Some(caller_ids) = node_ids.get(&signature.qualified_name) else {
continue;
};
let alias_map = import_alias_map(&signature.imports);
let caller_ns = caller_namespace(&signature.qualified_name);
for &caller_id in caller_ids {
for call_target in &signature.calls {
let candidates =
ordered_resolution_candidates(call_target, &alias_map, caller_ns.as_deref());
for target_id in local_call_targets(&candidates, &exact_map, &last_map, &suffix_map)
{
if caller_id != target_id && seen.insert((caller_id, target_id)) {
edges.push((caller_id, target_id));
}
}
if let Some(target_id) = type_node_target(call_target, node_ids, &last_map) {
if caller_id != target_id && seen.insert((caller_id, target_id)) {
edges.push((caller_id, target_id));
}
}
}
}
}
edges
}
fn add_local_flow_fact_edge(
fact: &FlowFact,
channel: &str,
caller_id: crate::graph::pdg::NodeId,
by_normalized: &HashMap<String, crate::graph::pdg::NodeId>,
by_last: &HashMap<String, Vec<crate::graph::pdg::NodeId>>,
pdg: &mut ProgramDependenceGraph,
) {
let Some(target) = resolve_flow_target(&fact.target, by_normalized, by_last) else {
return;
};
if target == caller_id {
return;
}
let mut metadata = EdgeMetadata::with_variable(fact.source.clone());
metadata.channel = Some(channel.to_string());
metadata.position = fact.position;
pdg.add_edge(
caller_id,
target,
Edge {
edge_type: EdgeType::DataDependency,
metadata,
},
);
}
fn extract_flow_edges(
signatures: &[SignatureInfo],
node_ids: &HashMap<String, crate::graph::pdg::NodeId>,
pdg: &mut ProgramDependenceGraph,
) {
let mut by_normalized: HashMap<String, crate::graph::pdg::NodeId> = HashMap::new();
let mut by_last: HashMap<String, Vec<crate::graph::pdg::NodeId>> = HashMap::new();
for sig in signatures {
if let Some(&id) = node_ids.get(&sig.qualified_name) {
by_normalized.insert(normalize_symbol(&sig.qualified_name), id);
if let Some(last) = normalize_symbol(&sig.qualified_name).rsplit('.').next() {
by_last.entry(last.to_string()).or_default().push(id);
}
}
}
let mut external: HashMap<String, crate::graph::pdg::NodeId> = HashMap::new();
for sig in signatures {
let Some(&caller_id) = node_ids.get(&sig.qualified_name) else {
continue;
};
let command = sig
.flow_facts
.iter()
.find(|fact| fact.channel == FlowChannel::CommandArgument && fact.target == "command")
.map(|fact| fact.source.clone());
for fact in &sig.flow_facts {
let (edge_type, target_label, channel) = match &fact.channel {
FlowChannel::Argument | FlowChannel::ReturnValue => {
add_local_flow_fact_edge(
fact,
&flow_channel_name(&fact.channel),
caller_id,
&by_normalized,
&by_last,
pdg,
);
continue;
}
FlowChannel::StateRead | FlowChannel::StateWrite => {
let target = resolve_flow_target(&fact.target, &by_normalized, &by_last);
let target_label = target
.and_then(|id| pdg.get_node(id).map(|node| node.name.to_string()))
.unwrap_or_else(|| fact.target.clone());
(
EdgeType::StateTransition,
target_label,
flow_channel_name(&fact.channel),
)
}
FlowChannel::CommandArgument => {
if fact.target == "command" {
continue;
}
let label = match fact.target.as_str() {
"argv" => command.clone().unwrap_or_else(|| fact.source.clone()),
"env" => fact.source.clone(),
_ => fact.target.clone(),
};
let edge_type = match fact.target.as_str() {
"env" => EdgeType::Environment,
"stdin" => EdgeType::Stdin,
_ => EdgeType::CommandArgument,
};
(edge_type, label, fact.target.clone())
}
FlowChannel::Environment => (
EdgeType::Environment,
fact.target.clone(),
"env".to_string(),
),
FlowChannel::Stdin => (EdgeType::Stdin, fact.target.clone(), "stdin".to_string()),
};
let target_id =
if let Some(id) = resolve_flow_target(&target_label, &by_normalized, &by_last) {
id
} else {
*external.entry(target_label.clone()).or_insert_with(|| {
pdg.add_node(Node {
id: format!("external:{}", target_label),
node_type: NodeType::External,
name: target_label.clone(),
file_path: Arc::from("<external>"),
byte_range: (0, 0),
complexity: 0,
language: "external".to_string(),
})
})
};
let mut metadata = EdgeMetadata::with_variable(fact.source.clone());
metadata.channel = Some(channel);
metadata.position = fact.position;
pdg.add_edge(
caller_id,
target_id,
Edge {
edge_type,
metadata,
},
);
}
}
}
fn resolve_flow_target(
target: &str,
by_normalized: &HashMap<String, crate::graph::pdg::NodeId>,
by_last: &HashMap<String, Vec<crate::graph::pdg::NodeId>>,
) -> Option<crate::graph::pdg::NodeId> {
let normalized = normalize_symbol(target);
by_normalized.get(&normalized).copied().or_else(|| {
normalized
.rsplit('.')
.next()
.and_then(|last| by_last.get(last)?.first().copied())
})
}
fn flow_channel_name(channel: &FlowChannel) -> String {
match channel {
FlowChannel::StateRead => "state_read",
FlowChannel::StateWrite => "state_write",
FlowChannel::Argument => "argument",
FlowChannel::ReturnValue => "return_value",
FlowChannel::CommandArgument => "argv",
FlowChannel::Environment => "env",
FlowChannel::Stdin => "stdin",
}
.to_string()
}
#[path = "extraction_cross_file.rs"]
mod cross_file;
pub use cross_file::{
resolve_cross_file_call_edges, resolve_cross_file_call_edges_for_files,
resolve_cross_file_flow_edges_for_files,
};
fn qualified_name_from_node(node: &Node) -> Option<&str> {
if let Some(qname) = node
.id
.strip_prefix(node.file_path.as_ref())
.and_then(|rest| rest.strip_prefix(':'))
{
return Some(qname.split_once('@').map_or(qname, |(qname, _)| qname));
}
let expected_file_name = Path::new(node.file_path.as_ref()).file_name()?.to_str()?;
node.id.char_indices().find_map(|(idx, ch)| {
if ch != ':' {
return None;
}
let (id_path, rest) = node.id.split_at(idx);
let rest = rest.strip_prefix(':')?;
if rest.is_empty() {
return None;
}
let id_file_name = Path::new(id_path).file_name()?.to_str()?;
(id_file_name == expected_file_name)
.then(|| rest.split_once('@').map_or(rest, |(qname, _)| qname))
})
}
fn import_alias_map(imports: &[ImportInfo]) -> HashMap<String, String> {
let mut alias_map = HashMap::new();
for import in imports {
let alias = import.alias.clone().or_else(|| {
import
.path
.split(['.', ':', '/', '\\'])
.next_back()
.map(|s| s.to_string())
});
if let Some(alias) = alias {
alias_map
.entry(alias)
.or_insert_with(|| import.path.clone());
}
}
alias_map
}
pub fn extract_import_paths_from_source(source_code: &[u8], language: &str) -> HashSet<String> {
let Ok(source) = std::str::from_utf8(source_code) else {
return HashSet::new();
};
let lang = language.to_ascii_lowercase();
let source = strip_block_comments(&lang, source);
match lang.as_str() {
"python" | "py" => extract_python_imports(&source),
"javascript" | "js" | "typescript" | "ts" | "jsx" | "tsx" => extract_js_ts_imports(&source),
"rust" | "rs" => extract_rust_imports(&source),
"go" | "golang" => extract_go_imports(&source),
"java" => extract_java_imports(&source),
"csharp" | "cs" | "c#" => extract_csharp_imports(&source),
"ruby" | "rb" => extract_ruby_imports(&source),
"php" => extract_php_imports(&source),
"lua" => extract_lua_imports(&source),
"scala" => extract_scala_imports(&source),
"c" | "cpp" | "c++" | "cxx" | "cc" | "h" | "hpp" => extract_c_imports(&source),
_ => HashSet::new(),
}
}
fn strip_block_comments(lang: &str, source: &str) -> String {
match lang {
"python" | "py" | "ruby" | "rb" => source.to_string(), _ => {
let mut result = String::with_capacity(source.len());
let mut chars = source.chars().peekable();
while let Some(c) = chars.next() {
if c == '/' && chars.peek() == Some(&'*') {
chars.next(); loop {
match chars.next() {
Some('*') if chars.peek() == Some(&'/') => {
chars.next();
break;
}
None => break,
_ => {}
}
}
result.push(' '); } else {
result.push(c);
}
}
result
}
}
}
fn extract_python_imports(source: &str) -> HashSet<String> {
let mut imports = HashSet::new();
let re_import = Regex::new(r"(?m)^import\s+([\w,\s.]+)").unwrap();
for cap in re_import.captures_iter(source) {
for name in cap[1].split(',') {
let trimmed = name.split_whitespace().next().unwrap_or("").trim();
if !trimmed.is_empty() {
imports.insert(trimmed.to_string());
}
}
}
let re_from =
Regex::new(r"(?s)from\s+([\w.]+)\s+import\s+(?:\(([^)]+)\)|(\w[\w\s,*]*))").unwrap();
for cap in re_from.captures_iter(source) {
let module = cap[1].trim();
imports.insert(module.to_string());
let names_str = cap.get(2).or(cap.get(3)).map(|m| m.as_str()).unwrap_or("");
for name in names_str.split(',') {
let sym = name
.split_whitespace()
.next()
.unwrap_or("")
.trim()
.trim_matches('*');
if !sym.is_empty() && sym != "*" {
imports.insert(format!("{}.{}", module, sym));
}
}
}
imports
}
fn extract_js_ts_imports(source: &str) -> HashSet<String> {
let mut imports = HashSet::new();
let re_named =
Regex::new(r#"(?s)import\s+(?:type\s+)?\{[^}]*\}\s+from\s+['"]([^'"]+)['"]"#).unwrap();
for cap in re_named.captures_iter(source) {
imports.insert(cap[1].trim().to_string());
}
let re_default =
Regex::new(r#"import\s+(?:type\s+)?(?:\*\s+as\s+\w+|\w+)\s+from\s+['"]([^'"]+)['"]"#)
.unwrap();
for cap in re_default.captures_iter(source) {
imports.insert(cap[1].trim().to_string());
}
let re_require = Regex::new(r#"require\s*\(\s*['"]([^'"]+)['"]\s*\)"#).unwrap();
for cap in re_require.captures_iter(source) {
imports.insert(cap[1].trim().to_string());
}
let re_export = Regex::new(r#"export\s+(?:\*|\{[^}]*\})\s+from\s+['"]([^'"]+)['"]"#).unwrap();
for cap in re_export.captures_iter(source) {
imports.insert(cap[1].trim().to_string());
}
imports
}
fn extract_rust_imports(source: &str) -> HashSet<String> {
let mut imports = HashSet::new();
let collapsed = collapse_multiline(source, "use ", ';');
for stmt in &collapsed {
let use_stmt = stmt.trim_start_matches("use ").trim_end_matches(';').trim();
expand_rust_use(use_stmt, &mut imports);
}
imports
}
fn expand_rust_use(stmt: &str, out: &mut HashSet<String>) {
if let Some(brace_start) = stmt.find('{') {
let base = stmt[..brace_start]
.trim()
.trim_end_matches("::")
.replace("::", ".");
let inner = stmt[brace_start + 1..]
.trim_end_matches('}')
.trim_end_matches(';');
for item in split_respecting_braces(inner) {
let item = item.trim();
if item == "self" {
out.insert(base.clone());
continue;
}
if item.contains('{') {
expand_rust_use(&format!("{}::{}", base.replace('.', "::"), item), out);
} else {
let full = format!("{}.{}", base, item.replace("::", "."));
out.insert(full);
}
}
} else {
out.insert(stmt.replace("::", ".").trim_matches('.').to_string());
}
}
fn split_respecting_braces(s: &str) -> Vec<&str> {
let mut result = Vec::new();
let mut depth = 0i32;
let mut last = 0;
for (i, c) in s.char_indices() {
match c {
'{' => depth += 1,
'}' => depth -= 1,
',' if depth == 0 => {
result.push(s[last..i].trim());
last = i + 1;
}
_ => {}
}
}
let tail = s[last..].trim();
if !tail.is_empty() {
result.push(tail);
}
result
}
fn collapse_multiline(source: &str, prefix: &str, terminator: char) -> Vec<String> {
let mut results = Vec::new();
let mut in_stmt = false;
let mut current = String::new();
for line in source.lines() {
let trimmed = line.trim();
if trimmed.starts_with("//") {
continue;
}
if !in_stmt && trimmed.starts_with(prefix) {
in_stmt = true;
current = trimmed.to_string();
} else if in_stmt {
current.push(' ');
current.push_str(trimmed);
}
if in_stmt {
if let Some(end) = current.find(terminator) {
results.push(current[..=end].to_string());
in_stmt = false;
current = String::new();
}
}
}
results
}
fn extract_go_imports(source: &str) -> HashSet<String> {
let mut imports = HashSet::new();
let re_single = Regex::new(r#"import\s+["']([^"']+)["']"#).unwrap();
for cap in re_single.captures_iter(source) {
imports.insert(cap[1].trim().to_string());
}
let re_block = Regex::new(r#"(?s)import\s*\(([^)]+)\)"#).unwrap();
let re_path = Regex::new(r#"["']([^"']+)["']"#).unwrap();
for cap in re_block.captures_iter(source) {
let inner = &cap[1];
for p in re_path.captures_iter(inner) {
imports.insert(p[1].trim().to_string());
}
}
imports
}
fn extract_java_imports(source: &str) -> HashSet<String> {
let mut imports = HashSet::new();
let re = Regex::new(r"(?m)^import(?:\s+static)?\s+([\w.*]+)\s*;").unwrap();
for cap in re.captures_iter(source) {
imports.insert(cap[1].trim().to_string());
}
imports
}
fn extract_csharp_imports(source: &str) -> HashSet<String> {
let mut imports = HashSet::new();
let re = Regex::new(r"(?m)^using(?:\s+static)?\s+([\w.]+)\s*;").unwrap();
for cap in re.captures_iter(source) {
imports.insert(cap[1].trim().to_string());
}
imports
}
fn extract_ruby_imports(source: &str) -> HashSet<String> {
let mut imports = HashSet::new();
let re = Regex::new(r#"(?:require|require_relative|load)\s*['"]([^'"]+)['"]"#).unwrap();
for cap in re.captures_iter(source) {
imports.insert(cap[1].trim().to_string());
}
imports
}
fn extract_php_imports(source: &str) -> HashSet<String> {
let mut imports = HashSet::new();
let re = Regex::new(r"(?m)^use\s+([\w\\]+)(?:\s+as\s+\w+)?\s*;").unwrap();
for cap in re.captures_iter(source) {
let path = cap[1].trim().replace('\\', ".");
imports.insert(path);
}
let re_require =
Regex::new(r#"(?:require|include)(?:_once)?\s*\(?['"]([^'"]+)['"]\)?"#).unwrap();
for cap in re_require.captures_iter(source) {
imports.insert(cap[1].trim().to_string());
}
imports
}
fn extract_lua_imports(source: &str) -> HashSet<String> {
let mut imports = HashSet::new();
let re = Regex::new(r#"require\s*\(?['"]([^'"]+)['"]\)?"#).unwrap();
for cap in re.captures_iter(source) {
imports.insert(cap[1].replace('.', "/").trim().to_string());
}
imports
}
fn extract_scala_imports(source: &str) -> HashSet<String> {
let mut imports = HashSet::new();
let re = Regex::new(r"(?m)^\s*import\s+([^\n]+)$").unwrap();
let selector_re = Regex::new(r"^([\w.]+)(?:\.\{([^}]+)\}|\.(\w+|\*))?$").unwrap();
for cap in re.captures_iter(source) {
let stmt = cap[1].trim();
if let Some(sel) = selector_re.captures(stmt) {
let base = &sel[1];
if let Some(names) = sel.get(2) {
for name in names.as_str().split(',') {
let n = name.trim();
if n != "_" && !n.is_empty() {
imports.insert(format!("{}.{}", base, n));
}
}
} else if let Some(single) = sel.get(3) {
imports.insert(format!("{}.{}", base, single.as_str()));
} else {
imports.insert(base.to_string());
}
}
}
imports
}
fn extract_c_imports(source: &str) -> HashSet<String> {
let mut imports = HashSet::new();
let re = Regex::new(r#"#include\s*[<"']([^>"']+)[>"']"#).unwrap();
for cap in re.captures_iter(source) {
imports.insert(cap[1].trim().to_string());
}
imports
}
fn extract_import_edges(
signatures: &[SignatureInfo],
node_ids: &HashMap<String, crate::graph::pdg::NodeId>,
pdg: &mut ProgramDependenceGraph,
file_path: &str,
language: &str,
source_code: &[u8],
) -> Vec<(crate::graph::pdg::NodeId, crate::graph::pdg::NodeId)> {
let mut edges = Vec::new();
let mut seen: HashSet<(crate::graph::pdg::NodeId, crate::graph::pdg::NodeId)> = HashSet::new();
let mut unique_paths: HashSet<String> = signatures
.iter()
.flat_map(|sig| sig.imports.iter().map(|imp| imp.path.clone()))
.collect();
unique_paths.extend(extract_import_paths_from_source(source_code, language));
if unique_paths.is_empty() {
return edges;
}
let module_sym = format!("{}:__module__", file_path);
let importer_nid = pdg.find_by_symbol(&module_sym).unwrap_or_else(|| {
pdg.add_node(Node {
id: module_sym,
node_type: NodeType::Module,
name: "__module__".to_string(),
file_path: Arc::from(file_path),
byte_range: (0, 0),
complexity: 1,
language: language.to_string(),
})
});
let mut symbol_map: HashMap<String, Vec<crate::graph::pdg::NodeId>> = HashMap::new();
for sig in signatures {
if let Some(&nid) = node_ids.get(&sig.qualified_name) {
let norm = normalize_symbol(&sig.qualified_name);
symbol_map.entry(norm.clone()).or_default().push(nid);
if let Some(last) = norm.split('.').next_back() {
symbol_map.entry(last.to_string()).or_default().push(nid);
}
}
}
let mut external_nodes: HashMap<String, crate::graph::pdg::NodeId> = HashMap::new();
for path in unique_paths {
let targets = resolve_import_targets(&path, &symbol_map);
let targets = if targets.is_empty() {
let eid = *external_nodes.entry(path.clone()).or_insert_with(|| {
pdg.add_node(Node {
id: format!("{}:__external__:{}", file_path, path),
node_type: NodeType::External,
name: path.clone(),
file_path: Arc::from(file_path),
byte_range: (0, 0),
complexity: 1,
language: "external".to_string(),
})
});
vec![eid]
} else {
targets
};
for target in targets {
if target == importer_nid {
continue;
}
if seen.insert((importer_nid, target)) {
edges.push((importer_nid, target));
}
}
}
edges
}
fn resolve_import_targets(
import_path: &str,
symbol_map: &HashMap<String, Vec<crate::graph::pdg::NodeId>>,
) -> Vec<crate::graph::pdg::NodeId> {
let normalized = normalize_symbol(import_path);
let mut targets: Vec<crate::graph::pdg::NodeId> = Vec::new();
if let Some(ids) = symbol_map.get(&normalized) {
targets.extend(ids);
}
let parts: Vec<&str> = normalized.split('.').collect();
for len in 2..=3_usize.min(parts.len()) {
let start = parts.len() - len;
let key = parts[start..].join(".");
if let Some(ids) = symbol_map.get(&key) {
targets.extend(ids);
}
}
if targets.is_empty() {
if let Some(last) = normalized.split('.').next_back() {
if let Some(ids) = symbol_map.get(last) {
targets.extend(ids);
}
}
}
targets.sort_by_key(|id| id.index());
targets.dedup();
targets
}
pub fn normalize_symbol(raw: &str) -> String {
let trimmed = raw.split('(').next().unwrap_or(raw).trim();
trimmed
.replace("?.", ".")
.replace("::", ".")
.replace("->", ".")
.replace(['\\', '/', ':'], ".")
.replace("..", ".")
.trim_matches('.')
.to_string()
}
fn signature_to_node(sig: &SignatureInfo, file_path: &str, language: &str) -> Node {
let node_type = match sig.return_type.as_deref() {
Some("module") => NodeType::Module,
Some("enum_variant") => NodeType::Variable,
Some("enum") | Some("trait") => NodeType::Class,
Some(value) if value.starts_with("struct") => NodeType::Class,
_ if sig.is_method => NodeType::Method,
_ => NodeType::Function,
};
let complexity = if sig.cyclomatic_complexity > 0 {
sig.cyclomatic_complexity
} else {
1u32 + sig.parameters.len() as u32
};
Node {
id: format!("{}:{}", file_path, sig.qualified_name),
node_type,
name: sig.name.clone(),
file_path: Arc::from(file_path),
byte_range: sig.byte_range,
complexity,
language: language.to_string(),
}
}
#[cfg(test)]
#[path = "extraction_test.rs"]
mod tests;