use std::collections::BTreeMap;
use crate::cache::ScanCache;
use crate::extract::{RefKind, Vis};
#[derive(Debug, Default)]
pub struct Ranking {
pub file_rank: BTreeMap<String, f64>,
pub name_refs: BTreeMap<String, u64>,
pub name_users: BTreeMap<String, BTreeMap<String, u64>>,
}
const DAMPING: f64 = 0.85;
const ITERATIONS: usize = 30;
pub fn rank(cache: &ScanCache) -> Ranking {
let nodes: Vec<&str> = cache
.files
.iter()
.filter_map(|(rel, entry)| entry.lang.map(|_| rel.as_str()))
.collect();
if nodes.is_empty() {
return Ranking::default();
}
let mut definers: BTreeMap<&str, Vec<usize>> = BTreeMap::new();
for (index, rel) in nodes.iter().enumerate() {
let entry = &cache.files[*rel];
let Some(lang) = entry.lang else {
continue;
};
let Some(x) = cache.parses.get(&(lang, entry.hash)) else {
continue;
};
for def in &x.defs {
if def.vis == Vis::Pub {
definers.entry(&def.name).or_default().push(index);
}
}
}
let mut edge_weights: Vec<BTreeMap<usize, f64>> = vec![BTreeMap::new(); nodes.len()];
let mut name_refs: BTreeMap<String, u64> = BTreeMap::new();
let mut name_users: BTreeMap<String, BTreeMap<String, u64>> = BTreeMap::new();
for (source, rel) in nodes.iter().enumerate() {
let entry = &cache.files[*rel];
let Some(lang) = entry.lang else {
continue;
};
let Some(x) = cache.parses.get(&(lang, entry.hash)) else {
continue;
};
for r in &x.refs {
let Some(defs) = definers.get(r.name.as_str()) else {
continue;
};
let targets: Vec<usize> = defs
.iter()
.copied()
.filter(|target| *target != source)
.collect();
if targets.is_empty() {
continue;
}
let weight = match r.kind {
RefKind::Call => 3,
RefKind::Import => 1,
};
*name_refs.entry(r.name.clone()).or_default() += weight;
*name_users
.entry(r.name.clone())
.or_default()
.entry((*rel).to_string())
.or_default() += weight;
let w = 1.0 / targets.len() as f64;
let out = &mut edge_weights[source];
for t in targets {
*out.entry(t).or_default() += w;
}
}
}
let edges: Vec<Vec<(usize, f64)>> = edge_weights
.into_iter()
.map(|out| {
let total: f64 = out.values().sum();
out.into_iter()
.map(|(target, weight)| (target, weight / total))
.collect()
})
.collect();
let n = nodes.len() as f64;
let mut rank = vec![1.0 / n; nodes.len()];
let mut next = vec![0.0; nodes.len()];
for _ in 0..ITERATIONS {
next.fill((1.0 - DAMPING) / n);
let mut dangling = 0.0f64;
for (source, out) in edges.iter().enumerate() {
if out.is_empty() {
dangling += rank[source];
continue;
}
for &(target, weight) in out {
next[target] += DAMPING * rank[source] * weight;
}
}
let spread = DAMPING * dangling / n;
for v in &mut next {
*v += spread;
}
std::mem::swap(&mut rank, &mut next);
}
Ranking {
file_rank: nodes
.into_iter()
.zip(rank)
.map(|(path, score)| (path.to_string(), score))
.collect(),
name_refs,
name_users,
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::cache::FileEntry;
use crate::extract::{Extraction, RefName, SymKind, Symbol};
use crate::lang::Lang;
fn file(cache: &mut ScanCache, rel: &str, defs: &[(&str, Vis)], refs: &[&str]) {
let hash = *blake3::hash(rel.as_bytes()).as_bytes();
cache.files.insert(
rel.into(),
FileEntry {
mtime: (1, 0),
size: 1,
ino: 0,
hash,
lang: Some(Lang::Python),
},
);
cache.parses.insert(
(Lang::Python, hash),
Extraction {
defs: defs
.iter()
.enumerate()
.map(|(i, (name, vis))| Symbol {
line: i as u32 + 1,
end_line: i as u32 + 2,
name: name.to_string(),
kind: SymKind::Fn,
vis: *vis,
sig: format!("def {name}()"),
terms: Vec::new(),
})
.collect(),
refs: refs
.iter()
.enumerate()
.map(|(i, name)| RefName {
line: i as u32 + 10,
name: name.to_string(),
kind: RefKind::Call,
})
.collect(),
},
);
}
#[test]
fn heavily_referenced_file_ranks_highest() {
let mut cache = ScanCache::default();
file(&mut cache, "core.py", &[("gean_core", Vis::Pub)], &[]);
file(&mut cache, "a.py", &[("a_fn", Vis::Pub)], &["gean_core"]);
file(&mut cache, "b.py", &[("b_fn", Vis::Pub)], &["gean_core"]);
file(
&mut cache,
"c.py",
&[("c_fn", Vis::Pub)],
&["gean_core", "a_fn"],
);
let r = rank(&cache);
let core = r.file_rank["core.py"];
assert!(core > r.file_rank["a.py"], "core outranks a");
assert!(core > r.file_rank["b.py"]);
assert_eq!(r.name_refs["gean_core"], 9, "3 callers x call-weight 3");
assert_eq!(r.name_refs["a_fn"], 3, "1 caller x call-weight 3");
let total: f64 = r.file_rank.values().sum();
assert!((total - 1.0).abs() < 1e-6, "rank mass conserved: {total}");
}
#[test]
fn private_defs_attract_no_edges_and_self_refs_ignored() {
let mut cache = ScanCache::default();
file(
&mut cache,
"x.py",
&[("_hidden", Vis::Priv), ("shown", Vis::Pub)],
&["shown"],
);
file(&mut cache, "y.py", &[], &["_hidden"]);
let r = rank(&cache);
assert!(
!r.name_refs.contains_key("_hidden"),
"private never counted"
);
assert!(
!r.name_refs.contains_key("shown"),
"self-ref not cross-file"
);
}
}