engram-memory 0.6.0

//! Weighted Reciprocal Rank Fusion for hybrid search.
//!
//! Fuses dense-vector and full-text hits via weighted Reciprocal Rank Fusion.
//! Both input slices arrive sorted best-first (HNSW by descending similarity,
//! FTS5 by ascending rank), so only each hit's rank position contributes —
//! the raw similarity/rank scores are intentionally ignored. A memory present
//! in both lists accumulates both weighted reciprocal-rank terms.

use std::collections::HashMap;

use crate::config::SearchConfig;

pub(crate) fn merge_results(
    vector_results: &[(String, f32)],
    sparse_results: &[(String, f64)],
    search: &SearchConfig,
) -> Vec<(String, f64)> {
    let k = search.rrf_k as f64;
    let mut combined: HashMap<String, f64> = HashMap::new();
    for (rank, (memory_id, _)) in vector_results.iter().enumerate() {
        *combined.entry(memory_id.clone()).or_insert(0.0) +=
            search.vector_weight / (k + (rank + 1) as f64);
    }
    for (rank, (memory_id, _)) in sparse_results.iter().enumerate() {
        *combined.entry(memory_id.clone()).or_insert(0.0) +=
            search.sparse_weight / (k + (rank + 1) as f64);
    }
    let mut results: Vec<(String, f64)> = combined.into_iter().collect();
    results.sort_by(|a, b| b.1.total_cmp(&a.1));
    results
}

pub(crate) fn limit_results(results: Vec<(String, f64)>, top_k: usize) -> Vec<(String, f64)> {
    results.into_iter().take(top_k).collect()
}

pub(crate) const SHADOW_K_SET: [usize; 4] = [3, 5, 10, 20];

/// Counterfactual RRF-mass coverage at each cutoff in `k_set`: the fraction of
/// the total fused score captured by the top-k hits. Pure and side-effect-free —
/// the result is logged for offline analysis and never feeds the served ranking.
pub(crate) fn shadow_rewards_for_k_set(
    merged: &[(String, f64)],
    k_set: &[usize],
) -> Vec<(usize, f64)> {
    let total: f64 = merged.iter().map(|(_, score)| score).sum();
    if merged.is_empty() || total <= 0.0 {
        return Vec::new();
    }
    k_set
        .iter()
        .map(|&k| {
            let covered: f64 = merged.iter().take(k).map(|(_, score)| score).sum();
            (k, covered / total)
        })
        .collect()
}

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

    fn vector_hit(id: &str) -> (String, f32) {
        (id.to_string(), 0.0)
    }

    fn sparse_hit(id: &str) -> (String, f64) {
        (id.to_string(), 0.0)
    }

    #[test]
    fn merge_ranks_by_position_not_score() {
        let vector = vec![vector_hit("a"), vector_hit("b")];
        let merged = merge_results(&vector, &[], &SearchConfig::default());
        assert_eq!(merged[0].0, "a");
        assert_eq!(merged[1].0, "b");
        assert!(merged[0].1 > merged[1].1);
    }

    #[test]
    fn merge_sums_contributions_for_shared_id() {
        let search = SearchConfig::default();
        let vector = vec![vector_hit("a"), vector_hit("b")];
        let sparse = vec![sparse_hit("a"), sparse_hit("c")];
        let merged = merge_results(&vector, &sparse, &search);

        let k = search.rrf_k as f64;
        let score_a = merged.iter().find(|(id, _)| id == "a").unwrap().1;
        let expected_a = search.vector_weight / (k + 1.0) + search.sparse_weight / (k + 1.0);
        assert!((score_a - expected_a).abs() < 1e-12);

        let score_b = merged.iter().find(|(id, _)| id == "b").unwrap().1;
        assert!(score_a > score_b);
    }

    // Asymmetric weights must drive the final ranking, not just rank position.
    // Construction: X is rank-1 in the dense-vector list and absent from sparse;
    // Y is rank-1 in sparse and far down the vector list (rank-9). With heavy
    // vector weighting (0.9 / 0.1) the dense source dominates and X leads; the
    // symmetric default (0.7 / 0.3) lets Y's strong sparse hit win instead, so
    // the two configs produce opposite orderings.
    //
    // Exact RRF (k = 60), heavy-vector 0.9 / 0.1:
    //   X = 0.9 / (60 + 1)              = 0.014754098360655738
    //   Y = 0.9 / (60 + 9) + 0.1/(60+1) = 0.014682822523164649  → X > Y
    // Symmetric default 0.7 / 0.3:
    //   X = 0.7 / (60 + 1)              = 0.011475409836065573
    //   Y = 0.7 / (60 + 9) + 0.3/(60+1) = 0.015062960323117129  → Y > X
    //
    // Y sits at vector rank-9 because, sharing both lists, it accrues the heavy
    // vector term too; only past rank-8 does X's single rank-1 vector hit
    // overtake Y's rank-9 vector hit plus its rank-1 sparse hit at 0.9 / 0.1.
    #[test]
    fn asymmetric_weights_drive_final_ranking() {
        let vector = vec![
            vector_hit("x"),
            vector_hit("v2"),
            vector_hit("v3"),
            vector_hit("v4"),
            vector_hit("v5"),
            vector_hit("v6"),
            vector_hit("v7"),
            vector_hit("v8"),
            vector_hit("y"),
        ];
        let sparse = vec![sparse_hit("y")];

        let heavy_vector = SearchConfig {
            rrf_k: 60,
            vector_weight: 0.9,
            sparse_weight: 0.1,
        };
        let merged = merge_results(&vector, &sparse, &heavy_vector);

        let score_x = merged.iter().find(|(id, _)| id == "x").unwrap().1;
        let score_y = merged.iter().find(|(id, _)| id == "y").unwrap().1;
        let expected_x = 0.9 / 61.0;
        let expected_y = 0.9 / 69.0 + 0.1 / 61.0;
        assert!((score_x - expected_x).abs() < 1e-12);
        assert!((score_y - expected_y).abs() < 1e-12);
        assert!(
            score_x > score_y,
            "heavy vector weighting must let the vector rank-1 id win: x={score_x}, y={score_y}"
        );
        let x_position = merged.iter().position(|(id, _)| id == "x").unwrap();
        let y_position = merged.iter().position(|(id, _)| id == "y").unwrap();
        assert!(x_position < y_position, "x must outrank y under 0.9 / 0.1");

        let symmetric = SearchConfig::default();
        let merged_symmetric = merge_results(&vector, &sparse, &symmetric);
        let symmetric_x = merged_symmetric.iter().find(|(id, _)| id == "x").unwrap().1;
        let symmetric_y = merged_symmetric.iter().find(|(id, _)| id == "y").unwrap().1;
        assert!(
            symmetric_y > symmetric_x,
            "the symmetric default must produce the opposite order: x={symmetric_x}, y={symmetric_y}"
        );
    }

    #[test]
    fn merge_empty_inputs_yield_empty() {
        let merged = merge_results(&[], &[], &SearchConfig::default());
        assert!(merged.is_empty());
    }

    #[test]
    fn merge_scores_sparse_only_doc_absent_from_vector() {
        let search = SearchConfig::default();
        let sparse = vec![sparse_hit("only_sparse")];
        let merged = merge_results(&[], &sparse, &search);

        let k = search.rrf_k as f64;
        assert_eq!(merged.len(), 1);
        assert_eq!(merged[0].0, "only_sparse");
        let expected = search.sparse_weight / (k + 1.0);
        assert!((merged[0].1 - expected).abs() < 1e-12);
    }

    fn scored(id: &str, score: f64) -> (String, f64) {
        (id.to_string(), score)
    }

    #[test]
    fn shadow_rewards_are_monotonic_non_decreasing() {
        let merged = vec![
            scored("a", 0.5),
            scored("b", 0.3),
            scored("c", 0.1),
            scored("d", 0.1),
        ];
        let rewards = shadow_rewards_for_k_set(&merged, &SHADOW_K_SET);
        assert_eq!(rewards.len(), SHADOW_K_SET.len());
        for window in rewards.windows(2) {
            assert!(
                window[1].1 >= window[0].1,
                "reward must not decrease as k grows: {window:?}"
            );
        }
    }

    #[test]
    fn shadow_reward_is_one_when_k_covers_all() {
        let merged = vec![scored("a", 0.6), scored("b", 0.4)];
        let rewards = shadow_rewards_for_k_set(&merged, &[10]);
        assert_eq!(rewards.len(), 1);
        assert_eq!(rewards[0].0, 10);
        assert!((rewards[0].1 - 1.0).abs() < 1e-12);
    }

    #[test]
    fn shadow_rewards_empty_merged_yields_empty() {
        let rewards = shadow_rewards_for_k_set(&[], &SHADOW_K_SET);
        assert!(rewards.is_empty());
    }

    #[test]
    fn shadow_rewards_non_positive_total_yields_empty() {
        let merged = vec![scored("a", 0.0), scored("b", 0.0)];
        let rewards = shadow_rewards_for_k_set(&merged, &SHADOW_K_SET);
        assert!(rewards.is_empty());
    }

    #[test]
    fn shadow_reward_single_element_is_one() {
        let merged = vec![scored("a", 0.42)];
        let rewards = shadow_rewards_for_k_set(&merged, &SHADOW_K_SET);
        assert_eq!(rewards.len(), SHADOW_K_SET.len());
        for (_, reward) in rewards {
            assert!((reward - 1.0).abs() < 1e-12);
        }
    }
}