use std::collections::{HashMap, HashSet};
pub const DEFAULT_DAMPING: f32 = 0.85;
pub const DEFAULT_MAX_ITERATIONS: usize = 50;
pub const DEFAULT_TOLERANCE: f32 = 1e-6;
#[derive(Debug, Default, Clone)]
pub struct CodeGraph {
out_edges: HashMap<u64, Vec<u64>>,
nodes: HashSet<u64>,
}
impl CodeGraph {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn add_node(&mut self, node: u64) {
self.nodes.insert(node);
}
pub fn add_edge(&mut self, from: u64, to: u64) {
self.nodes.insert(from);
self.nodes.insert(to);
if from == to {
return;
}
self.out_edges.entry(from).or_default().push(to);
}
pub fn add_undirected_edge(&mut self, a: u64, b: u64) {
self.add_edge(a, b);
self.add_edge(b, a);
}
#[must_use]
pub fn node_count(&self) -> usize {
self.nodes.len()
}
#[must_use]
pub fn edge_count(&self) -> usize {
self.out_edges.values().map(Vec::len).sum()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.nodes.is_empty()
}
pub fn nodes(&self) -> impl Iterator<Item = u64> + '_ {
self.nodes.iter().copied()
}
}
#[derive(Debug, Clone, Copy)]
pub struct PageRankConfig {
pub damping: f32,
pub max_iterations: usize,
pub tolerance: f32,
}
impl Default for PageRankConfig {
fn default() -> Self {
Self {
damping: DEFAULT_DAMPING,
max_iterations: DEFAULT_MAX_ITERATIONS,
tolerance: DEFAULT_TOLERANCE,
}
}
}
#[must_use]
pub fn compute_pagerank(graph: &CodeGraph) -> HashMap<u64, f32> {
compute_pagerank_with(graph, PageRankConfig::default())
}
#[must_use]
pub fn compute_pagerank_with(graph: &CodeGraph, config: PageRankConfig) -> HashMap<u64, f32> {
if graph.is_empty() {
return HashMap::new();
}
let n_nodes = graph.node_count();
#[allow(clippy::cast_precision_loss)]
let n_f = n_nodes as f32;
let base = (1.0 - config.damping) / n_f;
let mut node_list: Vec<u64> = graph.nodes().collect();
node_list.sort_unstable();
let node_index: HashMap<u64, usize> =
node_list.iter().enumerate().map(|(i, n)| (*n, i)).collect();
let mut in_neighbors: Vec<Vec<(usize, f32)>> = vec![Vec::new(); n_nodes];
let mut out_degree: Vec<f32> = vec![0.0; n_nodes];
for (from, tos) in &graph.out_edges {
let Some(&u_idx) = node_index.get(from) else {
continue;
};
#[allow(clippy::cast_precision_loss)]
let deg = tos.len() as f32;
out_degree[u_idx] = deg;
for to in tos {
if let Some(&v_idx) = node_index.get(to) {
in_neighbors[v_idx].push((u_idx, deg));
}
}
}
let mut rank: Vec<f32> = vec![1.0 / n_f; n_nodes];
let mut next: Vec<f32> = vec![0.0; n_nodes];
for _iter in 0..config.max_iterations {
let mut dangling: f32 = 0.0;
for i in 0..n_nodes {
if out_degree[i] == 0.0 {
dangling += rank[i];
}
}
let dangling_share = config.damping * dangling / n_f;
next.fill(base + dangling_share);
for v_idx in 0..n_nodes {
for &(u_idx, deg) in &in_neighbors[v_idx] {
next[v_idx] += config.damping * rank[u_idx] / deg;
}
}
let mut delta: f32 = 0.0;
for i in 0..n_nodes {
delta += (next[i] - rank[i]).abs();
}
std::mem::swap(&mut rank, &mut next);
if delta < config.tolerance {
break;
}
}
node_list.into_iter().zip(rank).collect()
}
#[must_use]
pub fn build_code_graph(
call_edges: &[(u64, u64)],
type_ref_edges: &[(u64, u64)],
hierarchy_edges: &[(u64, u64)],
all_chunk_ids: &[u64],
) -> CodeGraph {
let mut g = CodeGraph::new();
for &id in all_chunk_ids {
g.add_node(id);
}
for &(caller, callee) in call_edges {
g.add_edge(caller, callee);
}
for &(def, usage) in type_ref_edges {
g.add_undirected_edge(def, usage);
}
for &(child, parent) in hierarchy_edges {
g.add_undirected_edge(child, parent);
}
g
}
#[cfg(test)]
mod tests {
use super::*;
fn approx_eq(a: f32, b: f32, eps: f32) -> bool {
(a - b).abs() < eps
}
#[test]
fn empty_graph_returns_empty_ranks() {
let g = CodeGraph::new();
let ranks = compute_pagerank(&g);
assert!(ranks.is_empty());
}
#[test]
fn single_node_has_unit_rank() {
let mut g = CodeGraph::new();
g.add_node(1);
let ranks = compute_pagerank(&g);
assert_eq!(ranks.len(), 1);
assert!(
approx_eq(ranks[&1], 1.0, 1e-4),
"single node rank should be ~1.0, got {}",
ranks[&1]
);
}
#[test]
fn three_node_chain_central_node_wins() {
let mut g = CodeGraph::new();
g.add_edge(1, 2);
g.add_edge(2, 3);
g.add_edge(3, 2);
let ranks = compute_pagerank(&g);
assert_eq!(ranks.len(), 3);
assert!(
ranks[&2] > ranks[&1],
"node 2 ({}) should outrank node 1 ({})",
ranks[&2],
ranks[&1]
);
assert!(
ranks[&2] > ranks[&3],
"node 2 ({}) should outrank node 3 ({})",
ranks[&2],
ranks[&3]
);
let sum: f32 = ranks.values().sum();
assert!(
approx_eq(sum, 1.0, 1e-3),
"ranks should sum to ~1.0, got {sum}"
);
}
#[test]
fn star_graph_central_hub_dominates() {
let mut g = CodeGraph::new();
for leaf in 1..=5u64 {
g.add_edge(leaf, 0);
}
let ranks = compute_pagerank(&g);
let hub = ranks[&0];
for leaf in 1..=5u64 {
assert!(
hub > ranks[&leaf],
"hub rank ({hub}) should exceed leaf {} rank ({})",
leaf,
ranks[&leaf]
);
}
}
#[test]
fn isolated_nodes_get_base_rank() {
let mut g = CodeGraph::new();
g.add_node(10);
g.add_node(11);
g.add_edge(20, 21);
let ranks = compute_pagerank(&g);
assert_eq!(ranks.len(), 4);
assert!(approx_eq(ranks[&10], ranks[&11], 1e-4));
assert!(ranks[&10] > 0.0);
}
#[test]
fn self_loops_are_ignored() {
let mut g = CodeGraph::new();
g.add_edge(1, 1);
g.add_edge(2, 3);
let ranks = compute_pagerank(&g);
assert!(ranks.contains_key(&1));
assert!(ranks[&1] > 0.0);
}
#[test]
fn ranks_sum_to_unity() {
let mut g = CodeGraph::new();
let edges = [(1, 2), (2, 3), (3, 1), (4, 1), (5, 4), (6, 7), (7, 8)];
for (a, b) in edges {
g.add_edge(a, b);
}
let ranks = compute_pagerank(&g);
let sum: f32 = ranks.values().sum();
assert!(
approx_eq(sum, 1.0, 5e-3),
"ranks should sum to ~1.0, got {sum}"
);
}
#[test]
fn build_code_graph_merges_all_edge_types() {
let calls = vec![(1u64, 2u64)];
let type_refs = vec![(3u64, 4u64)];
let hierarchy = vec![(5u64, 6u64)];
let all_ids = vec![1u64, 2, 3, 4, 5, 6, 7]; let g = build_code_graph(&calls, &type_refs, &hierarchy, &all_ids);
assert_eq!(g.node_count(), 7);
assert_eq!(g.edge_count(), 5);
let ranks = compute_pagerank(&g);
assert_eq!(ranks.len(), 7);
assert!(ranks.contains_key(&7));
}
#[test]
fn convergence_respects_max_iterations() {
let mut g = CodeGraph::new();
for i in 0..20u64 {
for j in 0..20u64 {
if i != j {
g.add_edge(i, j);
}
}
}
let cfg = PageRankConfig {
damping: 0.85,
max_iterations: 3,
tolerance: 1e-12, };
let ranks = compute_pagerank_with(&g, cfg);
let sum: f32 = ranks.values().sum();
assert!(
approx_eq(sum, 1.0, 5e-3),
"ranks should remain a distribution even with iter cap: sum={sum}"
);
let expected = 1.0 / 20.0;
for v in ranks.values() {
assert!(
approx_eq(*v, expected, 0.02),
"complete-graph rank should be ~1/N, got {v}"
);
}
}
}