ll-core 0.2.0

Hybrid semantic search engine: BM25 + vector + PageRank + cross-encoder reranking over SQLite, with on-device ONNX inference
//! Personalized PageRank over a wikilink graph.
//!
//! The graph is represented as an adjacency map where each key is a note path
//! and each value is the list of paths that note links to (outgoing edges).

use std::collections::{HashMap, HashSet};

/// Directed adjacency map for the note wikilink graph.
///
/// Keys are note paths (e.g. `"brain/0-inbox/my-note.md"`). Values are lists
/// of paths that the key note links to (outgoing edges). Paths not present as
/// keys are sink nodes with no outgoing edges.
pub type GraphEdges = HashMap<String, Vec<String>>;

/// Stop iterating once the whole score vector moves less than this in one
/// step. Two orders of magnitude below the `1e-6` floor the results are
/// filtered by, so a step too small to trip this cannot change which nodes
/// survive that filter, nor their order beyond the same margin.
const CONVERGENCE_EPSILON: f64 = 1e-8;

/// Total absolute change across every node between two iterations. Nodes
/// present in one map and not the other count their full value: an entry
/// appearing for the first time has moved from zero.
fn l1_delta(prev: &HashMap<String, f64>, next: &HashMap<String, f64>) -> f64 {
    let mut delta = 0.0;
    for (node, score) in next {
        delta += (score - prev.get(node).copied().unwrap_or(0.0)).abs();
    }
    for (node, score) in prev {
        if !next.contains_key(node) {
            delta += score.abs();
        }
    }
    delta
}

/// Rank notes reachable from `seeds` using personalized PageRank.
///
/// Seeds are the starting set (e.g. vector + FTS top results). The algorithm
/// runs power iteration with damping factor `damping`, stopping when the score
/// vector settles or after `iterations` steps, whichever comes first.
/// Seed nodes are excluded from the output -- only related notes are returned.
///
/// Returns at most [`crate::TOP_K`] results as `(path, score)` pairs sorted
/// by descending score. Returns an empty vec if `seeds` or `graph` is empty.
pub fn personalized_pagerank(
    graph: &GraphEdges,
    seeds: &[String],
    damping: f32,
    iterations: usize,
) -> Vec<(String, f64)> {
    personalized_pagerank_holdout(graph, seeds, damping, iterations, None)
}

/// [`personalized_pagerank`] with a held-out node (eval-only).
///
/// When `holdout` is `Some(path)`, that node's adjacency list is skipped
/// during the walk — equivalent to removing the node's edges from the graph
/// without cloning it — and the node is filtered from the output. Inbound
/// edges may still feed score into the held-out node, but it acts as a sink:
/// it never propagates and is never returned.
pub fn personalized_pagerank_holdout(
    graph: &GraphEdges,
    seeds: &[String],
    damping: f32,
    iterations: usize,
    holdout: Option<&str>,
) -> Vec<(String, f64)> {
    if seeds.is_empty() || graph.is_empty() {
        return Vec::new();
    }

    let masked = |node: &str| holdout == Some(node);

    let mut inlink_counts: HashMap<&str, usize> = HashMap::new();
    for (node, targets) in graph {
        if masked(node) {
            continue;
        }
        for t in targets {
            *inlink_counts.entry(t.as_str()).or_default() += 1;
        }
    }

    let mut seed_weights: Vec<(&str, f64)> = seeds
        .iter()
        .filter(|s| graph.contains_key(s.as_str()) && !masked(s.as_str()))
        .map(|s| {
            let inlinks = inlink_counts.get(s.as_str()).copied().unwrap_or(0);
            (s.as_str(), 1.0 / (inlinks as f64 + 1.0))
        })
        .collect();

    let total_weight: f64 = seed_weights.iter().map(|(_, w)| w).sum();
    if total_weight == 0.0 {
        return Vec::new();
    }
    for (_, w) in &mut seed_weights {
        *w /= total_weight;
    }

    let d = damping as f64;
    let mut scores: HashMap<String, f64> = HashMap::new();
    let seed_set: HashSet<&str> = seeds.iter().map(|s| s.as_str()).collect();

    for &(s, w) in &seed_weights {
        scores.insert(s.to_string(), w);
    }

    for _ in 0..iterations {
        let mut new_scores: HashMap<String, f64> = HashMap::new();

        for &(s, w) in &seed_weights {
            *new_scores.entry(s.to_string()).or_default() += (1.0 - d) * w;
        }

        for (node, score) in &scores {
            if masked(node) {
                continue;
            }
            if let Some(neighbors) = graph.get(node) {
                let share = d * score / neighbors.len() as f64;
                for neighbor in neighbors {
                    *new_scores.entry(neighbor.clone()).or_default() += share;
                }
            }
        }

        let converged = l1_delta(&scores, &new_scores) < CONVERGENCE_EPSILON;
        scores = new_scores;
        if converged {
            break;
        }
    }

    let mut results: Vec<(String, f64)> = scores
        .into_iter()
        .filter(|(path, score)| *score > 1e-6 && !seed_set.contains(path.as_str()) && !masked(path))
        .collect();
    results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    results.truncate(crate::config::TOP_K);
    results
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_ppr_empty_graph() {
        let graph: GraphEdges = HashMap::new();
        let results = personalized_pagerank(&graph, &["a.md".to_string()], 0.5, 20);
        assert!(results.is_empty());
    }

    #[test]
    fn test_ppr_empty_seeds() {
        let mut graph: GraphEdges = HashMap::new();
        graph.insert("a".into(), vec!["b".into()]);
        graph.insert("b".into(), vec!["a".into()]);
        let results = personalized_pagerank(&graph, &[], 0.5, 20);
        assert!(results.is_empty());
    }

    #[test]
    fn test_ppr_holdout_masks_edges_and_filters_node() {
        // Undirected: a <-> s, s <-> g, a <-> x. Holding out s must (1) stop
        // the walk reaching g through s's edges, (2) keep s itself out of the
        // output even though the a -> s edge still feeds it score, and
        // (3) leave second-order reachability (a -> x) intact.
        let mut graph: GraphEdges = HashMap::new();
        graph.insert("a".into(), vec!["s".into(), "x".into()]);
        graph.insert("s".into(), vec!["a".into(), "g".into()]);
        graph.insert("g".into(), vec!["s".into()]);
        graph.insert("x".into(), vec!["a".into()]);

        let seeds = vec!["a".to_string()];
        let unmasked = personalized_pagerank(&graph, &seeds, 0.5, 20);
        let unmasked_paths: Vec<&str> = unmasked.iter().map(|r| r.0.as_str()).collect();
        assert!(unmasked_paths.contains(&"s") && unmasked_paths.contains(&"g"));

        let held = personalized_pagerank_holdout(&graph, &seeds, 0.5, 20, Some("s"));
        let held_paths: Vec<&str> = held.iter().map(|r| r.0.as_str()).collect();
        assert!(
            !held_paths.contains(&"s"),
            "held-out node must not be returned: {held_paths:?}"
        );
        assert!(
            !held_paths.contains(&"g"),
            "held-out node's edges must not propagate: {held_paths:?}"
        );
        assert!(
            held_paths.contains(&"x"),
            "unrelated structure must survive the holdout: {held_paths:?}"
        );
    }

    // Convergence: the early stop must not change what the walk produces, only
    // how long it takes to get there. Comparing against a high iteration cap is
    // the check that matters -- if stopping early changed the ranking, the
    // whole change would be a silent quality regression.
    //
    // Scores, not positions, are what these compare. A clique gives every
    // non-seed node the same score, and `sort_by` on equal keys leaves them in
    // whatever order the HashMap yielded, which Rust randomises per process. An
    // earlier version of this test asserted position equality and failed in CI
    // on exactly that: n1 against n11, two nodes with identical scores. That
    // tie order is nondeterministic on main too, so asserting it pins noise.
    fn dense_graph(n: usize) -> GraphEdges {
        let mut graph: GraphEdges = HashMap::new();
        for i in 0..n {
            let targets: Vec<String> = (0..n)
                .filter(|j| *j != i)
                .map(|j| format!("n{j}"))
                .collect();
            graph.insert(format!("n{i}"), targets);
        }
        graph
    }

    /// Path -> score, so a comparison is about values rather than tie order.
    fn score_map(results: &[(String, f64)]) -> HashMap<&str, f64> {
        results.iter().map(|(p, s)| (p.as_str(), *s)).collect()
    }

    fn assert_same_scores(capped: &[(String, f64)], uncapped: &[(String, f64)], case: &str) {
        let a = score_map(capped);
        let b = score_map(uncapped);
        assert_eq!(
            a.len(),
            b.len(),
            "{case}: early stop changed how many nodes survive the score floor"
        );
        for (path, score) in &a {
            let other = b
                .get(path)
                .unwrap_or_else(|| panic!("{case}: {path} missing from the uncapped run"));
            assert!(
                (score - other).abs() < 1e-6,
                "{case}: early stop changed {path} score: {score} vs {other}"
            );
        }
    }

    #[test]
    fn test_ppr_early_stop_matches_full_iteration() {
        // Same call twice, once with room to converge and once with far more
        // than it needs. Identical scores mean the stop fires only after the
        // vector has settled.
        let mut chain: GraphEdges = HashMap::new();
        for i in 0..30 {
            chain.insert(format!("n{i}"), vec![format!("n{}", (i + 1) % 30)]);
        }
        for (case, graph) in [("dense", dense_graph(12)), ("cycle", chain)] {
            let seeds = vec!["n0".to_string()];
            let capped = personalized_pagerank(&graph, &seeds, 0.5, 20);
            let uncapped = personalized_pagerank(&graph, &seeds, 0.5, 400);
            assert_same_scores(&capped, &uncapped, case);
        }
    }

    #[test]
    fn test_ppr_early_stop_preserves_a_real_ranking() {
        // The tie-free case: a chain gives each hop a distinct, decreasing
        // score, so position equality IS meaningful here and pins that the
        // early stop preserves genuine rank order.
        let mut graph: GraphEdges = HashMap::new();
        graph.insert("a".into(), vec!["b".into()]);
        graph.insert("b".into(), vec!["c".into()]);
        graph.insert("c".into(), vec!["d".into()]);
        graph.insert("d".into(), vec!["e".into()]);

        let seeds = vec!["a".to_string()];
        let capped = personalized_pagerank(&graph, &seeds, 0.5, 20);
        let uncapped = personalized_pagerank(&graph, &seeds, 0.5, 400);

        let order: Vec<&str> = capped.iter().map(|r| r.0.as_str()).collect();
        assert_eq!(
            order,
            vec!["b", "c", "d", "e"],
            "hop order is score order on a chain"
        );
        assert_eq!(
            order,
            uncapped.iter().map(|r| r.0.as_str()).collect::<Vec<_>>(),
            "early stop changed rank order on a graph with no ties"
        );
        assert_same_scores(&capped, &uncapped, "chain");
    }

    #[test]
    fn test_ppr_holdout_early_stop_matches_full_iteration() {
        // The holdout path shares the loop, so it shares the stop.
        let graph = dense_graph(10);
        let seeds = vec!["n0".to_string()];
        let capped = personalized_pagerank_holdout(&graph, &seeds, 0.5, 20, Some("n3"));
        let uncapped = personalized_pagerank_holdout(&graph, &seeds, 0.5, 400, Some("n3"));
        assert!(
            !score_map(&capped).contains_key("n3"),
            "held-out node must stay out of the result"
        );
        assert_same_scores(&capped, &uncapped, "holdout");
    }

    #[test]
    fn test_ppr_iteration_cap_is_still_honoured() {
        // A single step cannot have converged on this graph, so the cap must
        // still bound the walk: one iteration reaches only the seed's direct
        // neighbours, not the node two hops out.
        let mut graph: GraphEdges = HashMap::new();
        graph.insert("a".into(), vec!["b".into()]);
        graph.insert("b".into(), vec!["c".into()]);
        graph.insert("c".into(), vec!["d".into()]);

        let one = personalized_pagerank(&graph, &["a".to_string()], 0.5, 1);
        let paths: Vec<&str> = one.iter().map(|r| r.0.as_str()).collect();
        assert!(
            paths.contains(&"b"),
            "one step must reach the direct neighbour: {paths:?}"
        );
        assert!(
            !paths.contains(&"c"),
            "one step must not reach two hops out: {paths:?}"
        );
    }

    #[test]
    fn test_l1_delta_counts_appearing_and_vanishing_nodes() {
        // A node present on one side only has moved by its whole value. Missing
        // that would let the walk stop while score was still arriving.
        let mut prev: HashMap<String, f64> = HashMap::new();
        prev.insert("a".into(), 0.5);
        prev.insert("gone".into(), 0.25);
        let mut next: HashMap<String, f64> = HashMap::new();
        next.insert("a".into(), 0.5);
        next.insert("new".into(), 0.125);

        assert!((l1_delta(&prev, &next) - 0.375).abs() < 1e-12);
        assert_eq!(l1_delta(&prev, &prev), 0.0);
    }

    #[test]
    fn test_ppr_chain() {
        let mut graph: GraphEdges = HashMap::new();
        graph.insert("a".into(), vec!["b".into()]);
        graph.insert("b".into(), vec!["a".into(), "c".into()]);
        graph.insert("c".into(), vec!["b".into(), "d".into()]);
        graph.insert("d".into(), vec!["c".into()]);

        let results = personalized_pagerank(&graph, &["a".into()], 0.5, 20);
        assert!(!results.is_empty());
        let paths: Vec<&str> = results.iter().map(|r| r.0.as_str()).collect();
        assert!(paths.contains(&"b"));
        assert!(results[0].1 >= results.last().unwrap().1);
    }
}