Skip to main content

arbor_graph/
ranking.rs

1//! Centrality ranking for code nodes.
2//!
3//! We use a simplified PageRank variant to score nodes by their
4//! architectural significance. Nodes that are called by many
5//! others rank higher.
6
7use crate::graph::{ArborGraph, NodeId};
8use std::collections::HashMap;
9
10/// Stores centrality scores after computation.
11#[derive(Debug, Default)]
12pub struct CentralityScores {
13    scores: HashMap<NodeId, f64>,
14}
15
16impl CentralityScores {
17    /// Gets the score for a node.
18    pub fn get(&self, id: NodeId) -> f64 {
19        self.scores.get(&id).copied().unwrap_or(0.0)
20    }
21
22    /// Converts to a HashMap for storage in the graph.
23    pub fn into_map(self) -> HashMap<NodeId, f64> {
24        self.scores
25    }
26}
27
28/// Computes centrality scores for all nodes in the graph.
29///
30/// This is a simplified PageRank that:
31/// 1. Initializes all nodes with equal score
32/// 2. Iteratively distributes scores along edges
33/// 3. Applies damping to prevent score concentration
34///
35/// # Arguments
36///
37/// * `graph` - The graph to analyze
38/// * `iterations` - Number of iterations (10-20 is usually enough)
39/// * `damping` - Damping factor (0.85 is standard)
40pub fn compute_centrality(graph: &ArborGraph, iterations: usize, damping: f64) -> CentralityScores {
41    let node_count = graph.node_count();
42    if node_count == 0 {
43        return CentralityScores::default();
44    }
45
46    // Initialize scores
47    let initial_score = 1.0 / node_count as f64;
48    let mut scores: HashMap<NodeId, f64> = graph
49        .node_indexes()
50        .map(|idx| (idx, initial_score))
51        .collect();
52
53    // Count outgoing edges for each node
54    let mut out_degree: HashMap<NodeId, usize> = HashMap::new();
55    for idx in graph.node_indexes() {
56        let callees = graph.get_callees(idx);
57        out_degree.insert(idx, callees.len().max(1)); // Avoid division by zero
58    }
59
60    // Iterate
61    for _ in 0..iterations {
62        let mut new_scores: HashMap<NodeId, f64> = HashMap::new();
63
64        for idx in graph.node_indexes() {
65            // Base score (random jump)
66            let base = (1.0 - damping) / node_count as f64;
67
68            // Score from callers
69            let callers = graph.get_callers(idx);
70            let incoming: f64 = callers
71                .iter()
72                .filter_map(|caller| {
73                    let caller_idx = graph.get_index(&caller.id)?;
74                    let caller_score = scores.get(&caller_idx)?;
75                    let caller_out = *out_degree.get(&caller_idx)? as f64;
76                    Some(caller_score / caller_out)
77                })
78                .sum();
79
80            new_scores.insert(idx, base + damping * incoming);
81        }
82
83        scores = new_scores;
84    }
85
86    // Normalize to [0, 1] range
87    let max_score = scores.values().cloned().fold(0.0f64, f64::max);
88    if max_score > 0.0 {
89        for score in scores.values_mut() {
90            *score /= max_score;
91        }
92    }
93
94    CentralityScores { scores }
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100    use crate::edge::{Edge, EdgeKind};
101    use arbor_core::{CodeNode, NodeKind};
102
103    #[test]
104    fn test_centrality_empty_graph() {
105        let graph = ArborGraph::new();
106        let scores = compute_centrality(&graph, 10, 0.85);
107        assert!(scores.scores.is_empty());
108    }
109
110    #[test]
111    fn test_centrality_single_node() {
112        let mut graph = ArborGraph::new();
113        let node = CodeNode::new("foo", "foo", NodeKind::Function, "test.rs");
114        graph.add_node(node);
115
116        let scores = compute_centrality(&graph, 10, 0.85);
117        assert_eq!(scores.scores.len(), 1);
118    }
119
120    #[test]
121    fn test_centrality_popular_node_ranks_higher() {
122        let mut graph = ArborGraph::new();
123
124        // Create a "popular" function called by many others
125        let popular = CodeNode::new("popular", "popular", NodeKind::Function, "test.rs");
126        let popular_idx = graph.add_node(popular);
127
128        // Create callers
129        for i in 0..5 {
130            let caller = CodeNode::new(
131                format!("caller{}", i),
132                format!("caller{}", i),
133                NodeKind::Function,
134                "test.rs",
135            );
136            let caller_idx = graph.add_node(caller);
137            graph.add_edge(caller_idx, popular_idx, Edge::new(EdgeKind::Calls));
138        }
139
140        let scores = compute_centrality(&graph, 20, 0.85);
141
142        // The popular node should have the highest score
143        let popular_score = scores.get(popular_idx);
144        assert!(popular_score > 0.5, "Popular node should rank high");
145    }
146}