lean_ctx/core/
pagerank.rs

1//! PageRank computation on the Property Graph.
2//!
3//! Provides a reusable `compute` function for importance-weighted ranking
4//! (repo map, dashboard architecture views).
5
6use std::collections::{HashMap, HashSet};
7
8#[cfg(test)]
9use rusqlite::Connection;
10
11pub struct PageRankInput {
12    pub files: HashSet<String>,
13    pub forward: HashMap<String, Vec<String>>,
14}
15
16impl PageRankInput {
17    /// The raw file graph of the store, not checked against current sources
18    /// (test fixture).
19    #[cfg(test)]
20    pub(crate) fn from_connection(conn: &Connection) -> Self {
21        let mut files: HashSet<String> = HashSet::new();
22        let mut forward: HashMap<String, Vec<String>> = HashMap::new();
23
24        if let Ok(mut stmt) = conn.prepare(
25            "SELECT DISTINCT p.path
26             FROM nodes n JOIN paths p ON p.id = n.file_id
27             WHERE n.kind = 'file'",
28        ) && let Ok(rows) = stmt.query_map([], |row| row.get::<_, String>(0))
29        {
30            for f in rows.flatten() {
31                files.insert(f);
32            }
33        }
34
35        let edge_sql = "
36            SELECT DISTINCT p1.path, p2.path
37            FROM edges e
38            JOIN nodes n1 ON e.source_id = n1.id
39            JOIN nodes n2 ON e.target_id = n2.id
40            JOIN paths p1 ON p1.id = n1.file_id
41            JOIN paths p2 ON p2.id = n2.file_id
42            WHERE n1.kind = 'file' AND n2.kind = 'file'
43              AND n1.file_id != n2.file_id
44        ";
45        if let Ok(mut stmt) = conn.prepare(edge_sql)
46            && let Ok(rows) = stmt.query_map([], |row| {
47                Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
48            })
49        {
50            for row in rows.flatten() {
51                let (src, tgt) = row;
52                forward.entry(src).or_default().push(tgt);
53            }
54        }
55
56        for deps in forward.values_mut() {
57            deps.sort();
58            deps.dedup();
59        }
60
61        Self { files, forward }
62    }
63}
64
65pub fn compute(input: &PageRankInput, damping: f64, iterations: usize) -> HashMap<String, f64> {
66    compute_personalized(input, damping, iterations, &[])
67}
68
69/// Personalized PageRank: if `seed_files` is non-empty, teleportation bias goes
70/// to those files instead of uniform distribution. Handles dangling nodes
71/// (nodes with no outgoing edges) by redistributing their rank.
72pub(crate) fn compute_personalized(
73    input: &PageRankInput,
74    damping: f64,
75    iterations: usize,
76    seed_files: &[String],
77) -> HashMap<String, f64> {
78    let n = input.files.len();
79    if n == 0 {
80        return HashMap::new();
81    }
82
83    let personalization: HashMap<String, f64> = if seed_files.is_empty() {
84        let uniform = 1.0 / n as f64;
85        input.files.iter().map(|f| (f.clone(), uniform)).collect()
86    } else {
87        let valid_seeds: Vec<&String> = seed_files
88            .iter()
89            .filter(|f| input.files.contains(*f))
90            .collect();
91        if valid_seeds.is_empty() {
92            let uniform = 1.0 / n as f64;
93            input.files.iter().map(|f| (f.clone(), uniform)).collect()
94        } else {
95            let weight = 1.0 / valid_seeds.len() as f64;
96            let mut p = HashMap::new();
97            for f in &valid_seeds {
98                p.insert((*f).clone(), weight);
99            }
100            p
101        }
102    };
103
104    let dangling: HashSet<&String> = input
105        .files
106        .iter()
107        .filter(|f| !input.forward.contains_key(*f) || input.forward[*f].is_empty())
108        .collect();
109
110    let init = 1.0 / n as f64;
111    let mut rank: HashMap<String, f64> = input.files.iter().map(|f| (f.clone(), init)).collect();
112
113    let eps = 1e-8;
114    for _ in 0..iterations {
115        let dangling_sum: f64 = dangling
116            .iter()
117            .map(|f| rank.get(*f).copied().unwrap_or(0.0))
118            .sum();
119
120        let mut new_rank: HashMap<String, f64> = HashMap::with_capacity(n);
121
122        for f in &input.files {
123            let teleport = personalization.get(f).copied().unwrap_or(0.0);
124            let dangling_contrib = personalization.get(f).copied().unwrap_or(0.0) * dangling_sum;
125            new_rank.insert(
126                f.clone(),
127                (1.0 - damping) * teleport + damping * dangling_contrib,
128            );
129        }
130
131        for (node, neighbors) in &input.forward {
132            if neighbors.is_empty() {
133                continue;
134            }
135            let share = rank.get(node).copied().unwrap_or(0.0) / neighbors.len() as f64;
136            for neighbor in neighbors {
137                if let Some(nr) = new_rank.get_mut(neighbor) {
138                    *nr += damping * share;
139                }
140            }
141        }
142
143        let diff: f64 = input
144            .files
145            .iter()
146            .map(|f| {
147                (rank.get(f).copied().unwrap_or(0.0) - new_rank.get(f).copied().unwrap_or(0.0))
148                    .abs()
149            })
150            .sum();
151        rank = new_rank;
152
153        if diff < eps {
154            break;
155        }
156    }
157
158    rank
159}
160
161#[cfg(test)]
162mod tests {
163    use super::*;
164    use crate::core::property_graph::{CodeGraph, Edge, EdgeKind, Node};
165
166    #[test]
167    fn pagerank_basic() {
168        let g = CodeGraph::open_in_memory().unwrap();
169        let a = g.upsert_node(&Node::file("a.rs")).unwrap();
170        let b = g.upsert_node(&Node::file("b.rs")).unwrap();
171        let c = g.upsert_node(&Node::file("c.rs")).unwrap();
172
173        g.upsert_edge(&Edge::new(a, b, EdgeKind::Imports)).unwrap();
174        g.upsert_edge(&Edge::new(a, c, EdgeKind::Imports)).unwrap();
175        g.upsert_edge(&Edge::new(b, c, EdgeKind::Imports)).unwrap();
176
177        let input = PageRankInput::from_connection(g.connection());
178        let ranks = compute(&input, 0.85, 30);
179
180        assert_eq!(ranks.len(), 3);
181        let rank_c = ranks.get("c.rs").copied().unwrap_or(0.0);
182        let rank_a = ranks.get("a.rs").copied().unwrap_or(0.0);
183        assert!(
184            rank_c > rank_a,
185            "c.rs should rank higher (more incoming): c={rank_c} a={rank_a}"
186        );
187    }
188
189    #[test]
190    fn personalized_pagerank_boosts_seed() {
191        let g = CodeGraph::open_in_memory().unwrap();
192        let a = g.upsert_node(&Node::file("a.rs")).unwrap();
193        let b = g.upsert_node(&Node::file("b.rs")).unwrap();
194        let c = g.upsert_node(&Node::file("c.rs")).unwrap();
195
196        g.upsert_edge(&Edge::new(a, b, EdgeKind::Imports)).unwrap();
197        g.upsert_edge(&Edge::new(b, c, EdgeKind::Imports)).unwrap();
198
199        let input = PageRankInput::from_connection(g.connection());
200
201        let uniform = compute_personalized(&input, 0.85, 50, &[]);
202        let seeded = compute_personalized(&input, 0.85, 50, &["a.rs".to_string()]);
203
204        let a_uniform = uniform.get("a.rs").copied().unwrap_or(0.0);
205        let a_seeded = seeded.get("a.rs").copied().unwrap_or(0.0);
206
207        assert!(
208            a_seeded > a_uniform,
209            "seeded a.rs ({a_seeded}) should rank higher than uniform ({a_uniform})"
210        );
211    }
212
213    #[test]
214    fn early_convergence() {
215        let g = CodeGraph::open_in_memory().unwrap();
216        let a = g.upsert_node(&Node::file("a.rs")).unwrap();
217        let b = g.upsert_node(&Node::file("b.rs")).unwrap();
218        g.upsert_edge(&Edge::new(a, b, EdgeKind::Imports)).unwrap();
219        g.upsert_edge(&Edge::new(b, a, EdgeKind::Imports)).unwrap();
220
221        let input = PageRankInput::from_connection(g.connection());
222        let ranks = compute(&input, 0.85, 1000);
223        assert_eq!(ranks.len(), 2);
224    }
225}
lean_ctx/core/pagerank.rs

lean_ctx/core/
pagerank.rs
