Skip to main content

gobby_code/search/
rrf.rs

1//! Reciprocal Rank Fusion: merges ranked result lists from multiple sources.
2//!
3//! score(rank) = 1.0 / (K + rank) where K = 60.
4//! Ports logic from src/gobby/code_index/searcher.py.
5
6/// Merged result: (symbol_id, combined_score, source_names).
7pub type MergedResult = (String, f64, Vec<String>);
8
9/// Merge multiple ranked lists using Reciprocal Rank Fusion.
10///
11/// Each source is a `(name, ranked_ids)` pair where `ranked_ids` is ordered
12/// by relevance (index 0 = most relevant).
13///
14/// Returns `(id, score, sources)` sorted by score descending.
15pub fn merge(sources: Vec<(&str, Vec<String>)>) -> Vec<MergedResult> {
16    gobby_core::search::rrf_merge(sources)
17        .into_iter()
18        .map(|result| (result.id, result.score, result.sources))
19        .collect()
20}
21
22#[cfg(test)]
23mod tests {
24    use super::*;
25
26    #[test]
27    fn test_merge_single_source() {
28        let results = merge(vec![("fts", vec!["a".into(), "b".into(), "c".into()])]);
29        assert_eq!(results.len(), 3);
30        // First result should have highest score
31        assert_eq!(results[0].0, "a");
32        assert!(results[0].1 > results[1].1);
33        assert!(results[1].1 > results[2].1);
34    }
35
36    #[test]
37    fn test_merge_two_sources_same_ids() {
38        let results = merge(vec![
39            ("fts", vec!["a".into(), "b".into()]),
40            ("graph", vec!["a".into(), "c".into()]),
41        ]);
42        // "a" appears in both sources at rank 0, so it gets two rank-zero contributions.
43        let a_result = results.iter().find(|r| r.0 == "a").unwrap();
44        let expected = 2.0 * (1.0 / 60.0);
45        assert!((a_result.1 - expected).abs() < 1e-10);
46        assert_eq!(a_result.2.len(), 2);
47        // "a" should be ranked first
48        assert_eq!(results[0].0, "a");
49    }
50
51    #[test]
52    fn test_merge_sorts_sources_deterministically() {
53        let results = merge(vec![
54            ("semantic", vec!["b".into(), "a".into()]),
55            ("fts", vec!["b".into()]),
56        ]);
57
58        assert_eq!(results[0].0, "b");
59        assert_eq!(
60            results[0].2,
61            vec!["fts".to_string(), "semantic".to_string()]
62        );
63        assert_eq!(results[1].0, "a");
64    }
65
66    #[test]
67    fn test_merge_two_sources_disjoint() {
68        let results = merge(vec![("fts", vec!["a".into()]), ("graph", vec!["b".into()])]);
69        assert_eq!(results.len(), 2);
70        // Both have same score (rank 0 in their respective source)
71        assert!((results[0].1 - results[1].1).abs() < 1e-10);
72        // Each should have exactly 1 source
73        assert_eq!(results[0].2.len(), 1);
74        assert_eq!(results[1].2.len(), 1);
75    }
76
77    #[test]
78    fn test_merge_empty_sources() {
79        let results = merge(vec![]);
80        assert!(results.is_empty());
81    }
82
83    #[test]
84    fn test_merge_empty_id_lists() {
85        let results = merge(vec![("fts", vec![]), ("graph", vec![])]);
86        assert!(results.is_empty());
87    }
88
89    #[test]
90    fn merge_delegates_to_gobby_core_rrf() {
91        let sources = vec![
92            (
93                "fts",
94                vec!["a".to_string(), "a".to_string(), "b".to_string()],
95            ),
96            ("semantic", vec!["b".to_string()]),
97        ];
98        let results = merge(sources.clone());
99        let expected = gobby_core::search::rrf_merge(sources);
100
101        assert_eq!(results.len(), expected.len());
102        for (actual, expected) in results.iter().zip(expected.iter()) {
103            assert_eq!(actual.0, expected.id);
104            assert!((actual.1 - expected.score).abs() < 1e-10);
105            assert_eq!(actual.2, expected.sources);
106        }
107
108        let source = include_str!("rrf.rs");
109        let delegate = ["gobby_core", "::search::rrf_merge"].concat();
110        let local_const = ["const ", "RRF_K"].concat();
111        assert!(source.contains(&delegate));
112        assert!(!source.contains(&local_const));
113    }
114}