use std::collections::HashMap;
use rayon::prelude::*;
use crate::db::Database;
use crate::types::*;
pub struct ReferenceResolver<'a> {
#[allow(dead_code)]
db: &'a Database,
name_cache: HashMap<String, Vec<Node>>,
qualified_name_cache: HashMap<String, Vec<Node>>,
suffix_cache: HashMap<String, Vec<String>>,
}
impl<'a> ReferenceResolver<'a> {
pub async fn new(db: &'a Database) -> Self {
let all_nodes = db.get_all_nodes().await.unwrap_or_default();
Self::from_nodes(db, &all_nodes)
}
pub fn from_nodes(db: &'a Database, all_nodes: &[Node]) -> Self {
let mut name_cache: HashMap<String, Vec<Node>> = HashMap::new();
let mut qualified_name_cache: HashMap<String, Vec<Node>> = HashMap::new();
let mut suffix_cache: HashMap<String, Vec<String>> = HashMap::new();
for node in all_nodes {
name_cache
.entry(node.name.clone())
.or_default()
.push(node.clone());
let qn = &node.qualified_name;
qualified_name_cache
.entry(qn.clone())
.or_default()
.push(node.clone());
let mut pos = 0;
while let Some(idx) = qn[pos..].find("::") {
let suffix = &qn[pos + idx + 2..];
if !suffix.is_empty() {
suffix_cache
.entry(suffix.to_string())
.or_default()
.push(qn.clone());
}
pos += idx + 2;
}
}
for entries in suffix_cache.values_mut() {
entries.sort_unstable();
entries.dedup();
}
Self {
db,
name_cache,
qualified_name_cache,
suffix_cache,
}
}
pub fn resolve_one(&self, uref: &UnresolvedRef) -> Option<ResolvedRef> {
if uref.reference_name.contains("::") {
if let Some(resolved) = self.try_qualified_match(uref) {
return Some(resolved);
}
}
self.try_exact_name_match(uref)
}
pub fn resolve_all(&self, refs: &[UnresolvedRef]) -> ResolutionResult {
let total = refs.len();
let results: Vec<_> = refs
.par_iter()
.map(|uref| (uref, self.resolve_one(uref)))
.collect();
let mut resolved = Vec::new();
let mut unresolved = Vec::new();
for (uref, res) in results {
match res {
Some(r) => resolved.push(r),
None => unresolved.push(uref.clone()),
}
}
let resolved_count = resolved.len();
ResolutionResult {
resolved,
unresolved,
total,
resolved_count,
}
}
pub fn create_edges(&self, resolved: &[ResolvedRef]) -> Vec<Edge> {
resolved
.iter()
.map(|r| Edge {
source: r.original.from_node_id.clone(),
target: r.target_node_id.clone(),
kind: r.original.reference_kind.clone(),
line: Some(r.original.line),
})
.collect()
}
fn try_qualified_match(&self, uref: &UnresolvedRef) -> Option<ResolvedRef> {
if let Some(candidates) = self.qualified_name_cache.get(&uref.reference_name) {
if let Some(node) = candidates.first() {
return Some(ResolvedRef {
original: uref.clone(),
target_node_id: node.id.clone(),
confidence: 0.95,
resolved_by: "qualified-match".to_string(),
});
}
}
if let Some(full_names) = self.suffix_cache.get(&uref.reference_name) {
for full_name in full_names {
if let Some(candidates) = self.qualified_name_cache.get(full_name) {
if let Some(node) = candidates.first() {
return Some(ResolvedRef {
original: uref.clone(),
target_node_id: node.id.clone(),
confidence: 0.95,
resolved_by: "qualified-match".to_string(),
});
}
}
}
}
None
}
fn try_exact_name_match(&self, uref: &UnresolvedRef) -> Option<ResolvedRef> {
let candidates = self.name_cache.get(&uref.reference_name)?;
if candidates.len() == 1 {
return Some(ResolvedRef {
original: uref.clone(),
target_node_id: candidates[0].id.clone(),
confidence: 0.9,
resolved_by: "exact-match".to_string(),
});
}
let best = self.find_best_match(uref, candidates)?;
Some(ResolvedRef {
original: uref.clone(),
target_node_id: best.id.clone(),
confidence: 0.7,
resolved_by: "exact-match".to_string(),
})
}
fn find_best_match(&self, uref: &UnresolvedRef, candidates: &[Node]) -> Option<Node> {
if candidates.is_empty() {
return None;
}
let mut best_score = i64::MIN;
let mut best_node: Option<&Node> = None;
for node in candidates {
let mut score: i64 = 0;
if node.file_path == uref.file_path {
score += 100;
let distance = node.start_line.abs_diff(uref.line);
let proximity = 20_i64.saturating_sub(i64::from(distance) / 10);
score += proximity.max(0);
}
if node.visibility == Visibility::Pub {
score += 10;
}
if uref.reference_kind == EdgeKind::Calls
&& matches!(
node.kind,
NodeKind::Function
| NodeKind::Method
| NodeKind::StructMethod
| NodeKind::Constructor
| NodeKind::AbstractMethod
)
{
score += 25;
}
if score > best_score {
best_score = score;
best_node = Some(node);
}
}
best_node.cloned()
}
}