1use std::path::Path;
2use std::sync::Arc;
3use std::time::Instant;
4
5use rustc_hash::{FxHashMap, FxHashSet};
6
7use crate::config::bm25::BM25;
8use crate::config::limits::{LIMITS, PPR};
9use crate::config::scoring::{EGO, pit, rrf};
10use crate::config::tokenization::TOKENIZATION;
11use crate::edges;
12use crate::filtering;
13use crate::graph::{self, Graph};
14use crate::mode::{PipelineConfig, ScoringKind};
15use crate::ppr::personalized_pagerank;
16use crate::types::{DiffHunk, Fragment, FragmentId, extract_identifier_list};
17
18pub struct ScoringResult {
19 pub rel_scores: FxHashMap<FragmentId, f64>,
20 pub filtered_fragments: Vec<Fragment>,
21 pub graph: Graph,
22 pub graph_build_ms: f64,
27 pub ppr_truncated: bool,
30 pub ppr_forward_pushes: usize,
31 pub ppr_backward_pushes: usize,
32}
33
34fn finish_scoring(
41 fragments: &[Fragment],
42 core_ids: &FxHashSet<FragmentId>,
43 rel_scores: &FxHashMap<FragmentId, f64>,
44 graph: &Graph,
45) -> Vec<Fragment> {
46 let filtered = filtering::filter_unrelated_fragments(fragments, core_ids, graph);
47 let filtered = filtering::filter_positive_relevance(filtered, core_ids, rel_scores);
48 let filtered = filtering::filter_core_slice_context(filtered, core_ids);
49 filtering::cap_context_fragments(filtered, core_ids, rel_scores)
50}
51
52impl Default for ScoringResult {
53 fn default() -> Self {
54 Self {
55 rel_scores: FxHashMap::default(),
56 filtered_fragments: Vec::new(),
57 graph: Graph::new(),
58 graph_build_ms: 0.0,
59 ppr_truncated: false,
60 ppr_forward_pushes: 0,
61 ppr_backward_pushes: 0,
62 }
63 }
64}
65
66pub fn create_scoring_strategy(config: &PipelineConfig) -> Box<dyn ScoringStrategy> {
67 match config.scoring {
68 ScoringKind::Ego => Box::new(EgoGraphScoring::new(config.ego_depth)),
69 ScoringKind::Ppr => Box::new(PPRScoring::new(config.ppr_alpha)),
70 ScoringKind::Bm25 => Box::new(BM25Scoring),
71 ScoringKind::Rrf => Box::new(RrfFusionScoring::new(config.ego_depth)),
72 ScoringKind::Pit => Box::new(PitFusionScoring::new(config.ego_depth)),
73 }
74}
75
76pub trait ScoringStrategy: Send + Sync {
77 fn score_and_filter(
78 &self,
79 all_fragments: &[Fragment],
80 core_ids: &FxHashSet<FragmentId>,
81 hunks: &[DiffHunk],
82 repo_root: Option<&Path>,
83 seed_weights: Option<&FxHashMap<FragmentId, f64>>,
84 discovered_paths: Option<&FxHashSet<Arc<str>>>,
85 ) -> ScoringResult;
86}
87
88pub struct PPRScoring {
89 pub alpha: f64,
90}
91
92impl PPRScoring {
93 pub fn new(alpha: f64) -> Self {
94 Self { alpha }
95 }
96}
97
98impl ScoringStrategy for PPRScoring {
99 fn score_and_filter(
100 &self,
101 all_fragments: &[Fragment],
102 core_ids: &FxHashSet<FragmentId>,
103 hunks: &[DiffHunk],
104 repo_root: Option<&Path>,
105 seed_weights: Option<&FxHashMap<FragmentId, f64>>,
106 _discovered_paths: Option<&FxHashSet<Arc<str>>>,
107 ) -> ScoringResult {
108 let skip_expensive = all_fragments.len() > LIMITS.skip_expensive_threshold;
109 let t_graph = Instant::now();
110 let capped = edges::collect_capped_edges(all_fragments, repo_root, skip_expensive);
111 let mut g = graph::build_graph_capped(all_fragments, capped);
112 let graph_build_ms = t_graph.elapsed().as_secs_f64() * 1000.0;
113 let ppr = personalized_pagerank(
114 &mut g,
115 core_ids,
116 self.alpha,
117 PPR.convergence_tolerance,
118 PPR.forward_blend,
119 seed_weights,
120 );
121 let mut rel_scores = ppr.scores;
122 if ppr.truncated {
123 tracing::warn!(
124 "PPR push-cap hit on {} nodes (fwd_pushes={}, bwd_pushes={}); rel_scores biased",
125 g.node_count(),
126 ppr.forward_pushes,
127 ppr.backward_pushes,
128 );
129 }
130 filtering::apply_hunk_proximity_bonus(&mut rel_scores, core_ids, all_fragments, hunks);
131
132 let filtered = finish_scoring(all_fragments, core_ids, &rel_scores, &g);
133
134 ScoringResult {
135 rel_scores,
136 filtered_fragments: filtered,
137 graph: g,
138 graph_build_ms,
139 ppr_truncated: ppr.truncated,
140 ppr_forward_pushes: ppr.forward_pushes,
141 ppr_backward_pushes: ppr.backward_pushes,
142 }
143 }
144}
145
146pub struct EgoGraphScoring {
147 pub max_depth: usize,
148}
149
150impl EgoGraphScoring {
151 pub fn new(max_depth: usize) -> Self {
152 Self { max_depth }
153 }
154}
155
156impl ScoringStrategy for EgoGraphScoring {
157 fn score_and_filter(
158 &self,
159 all_fragments: &[Fragment],
160 core_ids: &FxHashSet<FragmentId>,
161 _hunks: &[DiffHunk],
162 repo_root: Option<&Path>,
163 _seed_weights: Option<&FxHashMap<FragmentId, f64>>,
164 _discovered_paths: Option<&FxHashSet<Arc<str>>>,
165 ) -> ScoringResult {
166 let skip_expensive = all_fragments.len() > LIMITS.skip_expensive_threshold;
167 let t_graph = Instant::now();
168 let capped = edges::collect_capped_edges(all_fragments, repo_root, skip_expensive);
169 let g = graph::build_graph_capped(all_fragments, capped);
170 let graph_build_ms = t_graph.elapsed().as_secs_f64() * 1000.0;
171 let mut rel_scores = g.ego_graph(core_ids, self.max_depth);
172
173 let diff_idents: FxHashSet<String> = all_fragments
174 .iter()
175 .filter(|f| core_ids.contains(&f.id))
176 .flat_map(|f| f.identifiers.iter().cloned())
177 .collect();
178
179 if !diff_idents.is_empty() {
180 for frag in all_fragments {
181 if core_ids.contains(&frag.id) || !rel_scores.contains_key(&frag.id) {
182 continue;
183 }
184 let overlap = frag.identifiers.intersection(&diff_idents).count();
185 if overlap > 0 {
186 let bonus = EGO.identifier_overlap_epsilon
187 * overlap.min(EGO.identifier_overlap_cap) as f64
188 / EGO.identifier_overlap_cap as f64;
189 *rel_scores.get_mut(&frag.id).unwrap() += bonus;
190 }
191 }
192 }
193
194 let filtered = finish_scoring(all_fragments, core_ids, &rel_scores, &g);
195
196 ScoringResult {
197 rel_scores,
198 filtered_fragments: filtered,
199 graph: g,
200 graph_build_ms,
201 ..Default::default()
202 }
203 }
204}
205
206pub struct BM25Scoring;
207
208impl ScoringStrategy for BM25Scoring {
209 fn score_and_filter(
210 &self,
211 all_fragments: &[Fragment],
212 core_ids: &FxHashSet<FragmentId>,
213 _hunks: &[DiffHunk],
214 _repo_root: Option<&Path>,
215 _seed_weights: Option<&FxHashMap<FragmentId, f64>>,
216 _discovered_paths: Option<&FxHashSet<Arc<str>>>,
217 ) -> ScoringResult {
218 let query_tokens: Vec<String> = all_fragments
219 .iter()
220 .filter(|f| core_ids.contains(&f.id))
221 .flat_map(|f| {
222 extract_identifier_list(&f.content, TOKENIZATION.query_min_identifier_length)
223 })
224 .collect();
225 let query_set: FxHashSet<String> = query_tokens.into_iter().collect();
226
227 let docs: Vec<(FragmentId, Vec<String>)> = all_fragments
228 .iter()
229 .filter(|f| !core_ids.contains(&f.id))
230 .map(|f| {
231 (
232 f.id.clone(),
233 extract_identifier_list(&f.content, TOKENIZATION.query_min_identifier_length),
234 )
235 })
236 .collect();
237
238 let n_docs = docs.len().max(1);
239 let avgdl = docs.iter().map(|(_, d)| d.len()).sum::<usize>() as f64 / n_docs as f64;
240
241 let mut df: FxHashMap<String, usize> = FxHashMap::default();
242 for (_, doc) in &docs {
243 let unique: FxHashSet<&str> = doc.iter().map(|s| s.as_str()).collect();
244 for term in unique {
245 *df.entry(term.to_string()).or_insert(0) += 1;
246 }
247 }
248
249 let idf: FxHashMap<String, f64> = query_set
250 .iter()
251 .map(|t| {
252 let d = df.get(t).copied().unwrap_or(0) as f64;
253 let val =
254 ((n_docs as f64 - d + BM25.idf_smoothing) / (d + BM25.idf_smoothing)).ln_1p();
255 (t.clone(), val)
256 })
257 .collect();
258
259 let mut rel_scores: FxHashMap<FragmentId, f64> = FxHashMap::default();
260 for frag in all_fragments {
261 if core_ids.contains(&frag.id) {
262 rel_scores.insert(frag.id.clone(), 1.0);
263 }
264 }
265 for (fid, doc) in &docs {
266 let dl = doc.len() as f64;
267 let mut tf: FxHashMap<&str, u32> = FxHashMap::default();
268 for t in doc {
269 *tf.entry(t.as_str()).or_insert(0) += 1;
270 }
271 let mut score = 0.0;
272 for t in &query_set {
273 let freq = tf.get(t.as_str()).copied().unwrap_or(0) as f64;
274 if freq == 0.0 {
275 continue;
276 }
277 let idf_val = idf.get(t).copied().unwrap_or(0.0);
278 score += idf_val * (freq * BM25.k1)
279 / (freq + BM25.k1 * (1.0 - BM25.b + BM25.b * dl / avgdl));
280 }
281 if score > 0.0 {
282 rel_scores.insert(fid.clone(), score);
283 }
284 }
285
286 let max_score = rel_scores.values().copied().fold(0.0f64, f64::max);
287 if max_score > 0.0 {
288 for v in rel_scores.values_mut() {
289 *v /= max_score;
290 }
291 }
292
293 let filtered =
297 filtering::filter_positive_relevance(all_fragments.to_vec(), core_ids, &rel_scores);
298 let filtered = filtering::cap_context_fragments(filtered, core_ids, &rel_scores);
299
300 ScoringResult {
301 rel_scores,
302 filtered_fragments: filtered,
303 ..Default::default()
304 }
305 }
306}
307
308pub struct RrfFusionScoring {
317 pub ego_depth: usize,
318 pub k: f64,
319}
320
321impl RrfFusionScoring {
322 pub fn new(ego_depth: usize) -> Self {
323 Self {
324 ego_depth,
325 k: rrf().k,
326 }
327 }
328}
329
330fn rank_positions(
342 rel: &FxHashMap<FragmentId, f64>,
343 admitted: &FxHashSet<FragmentId>,
344 core_ids: &FxHashSet<FragmentId>,
345) -> FxHashMap<FragmentId, usize> {
346 let mut ranked: Vec<(&FragmentId, f64)> = rel
347 .iter()
348 .filter(|(fid, score)| **score > 0.0 && !core_ids.contains(*fid) && admitted.contains(*fid))
349 .map(|(fid, score)| (fid, *score))
350 .collect();
351 ranked.sort_by(|(ida, sa), (idb, sb)| sb.total_cmp(sa).then_with(|| ida.cmp(idb)));
354 ranked
355 .into_iter()
356 .enumerate()
357 .map(|(i, (fid, _))| (fid.clone(), i + 1))
358 .collect()
359}
360
361fn fuse_reciprocal_ranks(
362 components: &[(&FxHashMap<FragmentId, f64>, &FxHashSet<FragmentId>)],
363 core_ids: &FxHashSet<FragmentId>,
364 k: f64,
365) -> FxHashMap<FragmentId, f64> {
366 let mut fused: FxHashMap<FragmentId, f64> = FxHashMap::default();
367 for (rel, admitted) in components {
368 for (fid, rank) in rank_positions(rel, admitted, core_ids) {
369 *fused.entry(fid).or_insert(0.0) += 1.0 / (k + rank as f64);
370 }
371 }
372
373 let max_fused = fused.values().copied().fold(0.0f64, f64::max);
374 if max_fused > 0.0 {
375 for v in fused.values_mut() {
376 *v /= max_fused;
377 }
378 }
379 for fid in core_ids {
387 fused.insert(fid.clone(), 1.0);
388 }
389 fused
390}
391
392impl ScoringStrategy for RrfFusionScoring {
393 fn score_and_filter(
394 &self,
395 all_fragments: &[Fragment],
396 core_ids: &FxHashSet<FragmentId>,
397 hunks: &[DiffHunk],
398 repo_root: Option<&Path>,
399 seed_weights: Option<&FxHashMap<FragmentId, f64>>,
400 discovered_paths: Option<&FxHashSet<Arc<str>>>,
401 ) -> ScoringResult {
402 let ego = EgoGraphScoring::new(self.ego_depth).score_and_filter(
403 all_fragments,
404 core_ids,
405 hunks,
406 repo_root,
407 seed_weights,
408 discovered_paths,
409 );
410 let lexical = BM25Scoring.score_and_filter(
411 all_fragments,
412 core_ids,
413 hunks,
414 repo_root,
415 seed_weights,
416 discovered_paths,
417 );
418
419 let ego_admitted: FxHashSet<FragmentId> = ego
420 .filtered_fragments
421 .iter()
422 .map(|f| f.id.clone())
423 .collect();
424 let lexical_admitted: FxHashSet<FragmentId> = lexical
425 .filtered_fragments
426 .iter()
427 .map(|f| f.id.clone())
428 .collect();
429
430 let rel_scores = fuse_reciprocal_ranks(
431 &[
432 (&ego.rel_scores, &ego_admitted),
433 (&lexical.rel_scores, &lexical_admitted),
434 ],
435 core_ids,
436 self.k,
437 );
438
439 let union_ids: FxHashSet<FragmentId> =
440 ego_admitted.union(&lexical_admitted).cloned().collect();
441
442 let union: Vec<Fragment> = all_fragments
443 .iter()
444 .filter(|f| union_ids.contains(&f.id))
445 .cloned()
446 .collect();
447
448 let filtered = finish_scoring(&union, core_ids, &rel_scores, &ego.graph);
466
467 ScoringResult {
468 rel_scores,
469 filtered_fragments: filtered,
470 graph: ego.graph,
471 graph_build_ms: ego.graph_build_ms,
472 ..Default::default()
473 }
474 }
475}
476
477pub struct PitFusionScoring {
497 pub ego_depth: usize,
498 pub blend: f64,
499 pub agreement_bonus: f64,
500 pub agreement_top_k: usize,
501}
502
503impl PitFusionScoring {
504 pub fn new(ego_depth: usize) -> Self {
505 let cfg = pit();
506 Self {
507 ego_depth,
508 blend: cfg.blend,
509 agreement_bonus: cfg.agreement_bonus,
510 agreement_top_k: cfg.agreement_top_k,
511 }
512 }
513}
514
515fn percentiles(
533 rel: &FxHashMap<FragmentId, f64>,
534 admitted: &FxHashSet<FragmentId>,
535 core_ids: &FxHashSet<FragmentId>,
536 top_k: usize,
537) -> (FxHashMap<FragmentId, f64>, FxHashSet<FragmentId>) {
538 let mut ranked: Vec<(&FragmentId, f64)> = rel
539 .iter()
540 .filter(|(fid, score)| **score > 0.0 && !core_ids.contains(*fid))
541 .map(|(fid, score)| (fid, *score))
542 .collect();
543 ranked.sort_by(|(ida, sa), (idb, sb)| sb.total_cmp(sa).then_with(|| ida.cmp(idb)));
546
547 let n = ranked.len();
548 let mut out: FxHashMap<FragmentId, f64> = FxHashMap::default();
549 let mut top: FxHashSet<FragmentId> = FxHashSet::default();
550 if n == 0 {
551 return (out, top);
552 }
553
554 if std::env::var("DIFFCTX_PIT_TRANSFORM").as_deref() == Ok("maxnorm") {
562 let denom = ranked
563 .iter()
564 .map(|(_, s)| *s)
565 .fold(0.0f64, f64::max)
566 .max(f64::MIN_POSITIVE);
567 for (fid, score) in &ranked {
568 if admitted.contains(*fid) {
569 out.insert((*fid).clone(), *score / denom);
570 }
571 }
572 } else {
573 let mut i = 0;
574 while i < n {
575 let mut j = i;
577 while j + 1 < n && ranked[j + 1].1.to_bits() == ranked[i].1.to_bits() {
578 j += 1;
579 }
580 let mean_pos = (i + j) as f64 / 2.0;
582 let percentile = 1.0 - mean_pos / n as f64;
583 for (fid, _) in &ranked[i..=j] {
584 if admitted.contains(*fid) {
585 out.insert((*fid).clone(), percentile);
586 }
587 }
588 i = j + 1;
589 }
590 }
591
592 for (fid, _) in ranked
596 .iter()
597 .filter(|(fid, _)| admitted.contains(*fid))
598 .take(top_k)
599 {
600 top.insert((*fid).clone());
601 }
602 (out, top)
603}
604
605fn quantile_map_to(
627 reference: &[f64],
628 fused: &FxHashMap<FragmentId, f64>,
629) -> FxHashMap<FragmentId, f64> {
630 let mut order: Vec<(&FragmentId, f64)> = fused
635 .iter()
636 .filter(|(_, s)| **s > 0.0)
637 .map(|(f, s)| (f, *s))
638 .collect();
639 if reference.is_empty() || order.is_empty() {
640 return order.into_iter().map(|(f, s)| (f.clone(), s)).collect();
641 }
642 order.sort_by(|(ida, sa), (idb, sb)| sb.total_cmp(sa).then_with(|| ida.cmp(idb)));
643
644 let m = order.len();
645 let n = reference.len();
646 let mut out = FxHashMap::default();
647 let mut i = 0;
648 while i < m {
649 let mut j = i;
653 while j + 1 < m && order[j + 1].1.to_bits() == order[i].1.to_bits() {
654 j += 1;
655 }
656 let mid_rank = (i + j) as f64 / 2.0;
657 let idx = if m == 1 {
658 n - 1
659 } else {
660 let pos = ((m - 1) as f64 - mid_rank) / (m - 1) as f64;
661 ((pos * (n - 1) as f64).round() as usize).min(n - 1)
662 };
663 for (fid, _) in &order[i..=j] {
664 out.insert((*fid).clone(), reference[idx]);
665 }
666 i = j + 1;
667 }
668 out
669}
670
671impl ScoringStrategy for PitFusionScoring {
672 fn score_and_filter(
673 &self,
674 all_fragments: &[Fragment],
675 core_ids: &FxHashSet<FragmentId>,
676 hunks: &[DiffHunk],
677 repo_root: Option<&Path>,
678 seed_weights: Option<&FxHashMap<FragmentId, f64>>,
679 discovered_paths: Option<&FxHashSet<Arc<str>>>,
680 ) -> ScoringResult {
681 let ego = EgoGraphScoring::new(self.ego_depth).score_and_filter(
682 all_fragments,
683 core_ids,
684 hunks,
685 repo_root,
686 seed_weights,
687 discovered_paths,
688 );
689 let lexical = BM25Scoring.score_and_filter(
690 all_fragments,
691 core_ids,
692 hunks,
693 repo_root,
694 seed_weights,
695 discovered_paths,
696 );
697
698 let ego_admitted: FxHashSet<FragmentId> = ego
699 .filtered_fragments
700 .iter()
701 .map(|f| f.id.clone())
702 .collect();
703 let lexical_admitted: FxHashSet<FragmentId> = lexical
704 .filtered_fragments
705 .iter()
706 .map(|f| f.id.clone())
707 .collect();
708
709 let (ego_pct, ego_top) = percentiles(
710 &ego.rel_scores,
711 &ego_admitted,
712 core_ids,
713 self.agreement_top_k,
714 );
715 let (lex_pct, lex_top) = percentiles(
716 &lexical.rel_scores,
717 &lexical_admitted,
718 core_ids,
719 self.agreement_top_k,
720 );
721
722 let mut rel_scores: FxHashMap<FragmentId, f64> = FxHashMap::default();
723 for fid in ego_pct.keys().chain(lex_pct.keys()) {
724 if rel_scores.contains_key(fid) {
725 continue;
726 }
727 let e = ego_pct.get(fid).copied().unwrap_or(0.0);
733 let l = lex_pct.get(fid).copied().unwrap_or(0.0);
734 let mut score = self.blend * e + (1.0 - self.blend) * l;
735 if ego_top.contains(fid) && lex_top.contains(fid) {
736 score += self.agreement_bonus;
737 }
738 rel_scores.insert(fid.clone(), score);
739 }
740
741 if std::env::var("DIFFCTX_PIT_SHAPE").as_deref() == Ok("flat") {
745 let max_fused = rel_scores.values().copied().fold(0.0f64, f64::max);
746 if max_fused > 0.0 {
747 for v in rel_scores.values_mut() {
748 *v /= max_fused;
749 }
750 }
751 for fid in core_ids {
752 rel_scores.insert(fid.clone(), 1.0);
753 }
754 } else {
755 let mut reference: Vec<f64> = ego
756 .rel_scores
757 .iter()
758 .filter(|(fid, s)| {
759 **s > 0.0 && !core_ids.contains(*fid) && ego_admitted.contains(*fid)
760 })
761 .map(|(_, s)| *s)
762 .collect();
763 reference.sort_by(f64::total_cmp);
764 rel_scores = quantile_map_to(&reference, &rel_scores);
765
766 for fid in core_ids {
770 if let Some(s) = ego.rel_scores.get(fid) {
771 rel_scores.insert(fid.clone(), *s);
772 }
773 }
774 }
775
776 let union_ids: FxHashSet<FragmentId> =
777 ego_admitted.union(&lexical_admitted).cloned().collect();
778 let union: Vec<Fragment> = all_fragments
779 .iter()
780 .filter(|f| union_ids.contains(&f.id))
781 .cloned()
782 .collect();
783 let filtered = finish_scoring(&union, core_ids, &rel_scores, &ego.graph);
784
785 ScoringResult {
786 rel_scores,
787 filtered_fragments: filtered,
788 graph: ego.graph,
789 graph_build_ms: ego.graph_build_ms,
790 ..Default::default()
791 }
792 }
793}
794
795#[cfg(test)]
796mod tests {
797 use super::*;
798 use crate::types::FragmentKind;
799
800 fn fid(path: &str, start: u32) -> FragmentId {
801 FragmentId::new(Arc::from(path), start, start + 4)
802 }
803
804 fn scores(entries: &[(FragmentId, f64)]) -> FxHashMap<FragmentId, f64> {
805 entries.iter().cloned().collect()
806 }
807
808 fn ballot(rel: &FxHashMap<FragmentId, f64>) -> FxHashSet<FragmentId> {
812 rel.keys().cloned().collect()
813 }
814
815 #[test]
820 fn agreement_between_both_signals_outranks_a_single_signal_top_hit() {
821 let agreed = fid("agreed.rs", 1);
822 let ego_only = fid("ego_only.rs", 1);
823 let cores: FxHashSet<FragmentId> = FxHashSet::default();
824
825 let ego = scores(&[(ego_only.clone(), 0.9), (agreed.clone(), 0.5)]);
828 let lexical = scores(&[(agreed.clone(), 0.5)]);
829
830 let fused = fuse_reciprocal_ranks(
831 &[(&ego, &ballot(&ego)), (&lexical, &ballot(&lexical))],
832 &cores,
833 60.0,
834 );
835 assert!(
836 fused[&agreed] > fused[&ego_only],
837 "agreement lost to a single-signal top hit: {:?} vs {:?}",
838 fused[&agreed],
839 fused[&ego_only]
840 );
841 }
842
843 #[test]
847 fn only_the_rank_order_of_a_component_matters_not_its_scale() {
848 let a = fid("a.rs", 1);
849 let b = fid("b.rs", 1);
850 let cores: FxHashSet<FragmentId> = FxHashSet::default();
851 let lexical = scores(&[(a.clone(), 0.4), (b.clone(), 0.1)]);
852
853 let modest = scores(&[(a.clone(), 0.6), (b.clone(), 0.4)]);
854 let enormous = scores(&[(a.clone(), 6_000.0), (b.clone(), 4_000.0)]);
855
856 assert_eq!(
857 fuse_reciprocal_ranks(
858 &[(&modest, &ballot(&modest)), (&lexical, &ballot(&lexical))],
859 &cores,
860 60.0
861 ),
862 fuse_reciprocal_ranks(
863 &[
864 (&enormous, &ballot(&enormous)),
865 (&lexical, &ballot(&lexical))
866 ],
867 &cores,
868 60.0
869 ),
870 "rescaling one component changed the fused scores"
871 );
872 }
873
874 #[test]
878 fn fused_scores_are_normalised_and_cores_anchor_the_top() {
879 let core = fid("changed.rs", 1);
880 let ctx = fid("ctx.rs", 1);
881 let cores: FxHashSet<FragmentId> = std::iter::once(core.clone()).collect();
882 let ego = scores(&[(core.clone(), 1.0), (ctx.clone(), 0.3)]);
883 let lexical = scores(&[(ctx.clone(), 0.2)]);
884
885 let fused = fuse_reciprocal_ranks(
886 &[(&ego, &ballot(&ego)), (&lexical, &ballot(&lexical))],
887 &cores,
888 60.0,
889 );
890 assert_eq!(fused[&core], 1.0, "core is not anchored at the top");
891 for (id, score) in &fused {
892 assert!(
893 (0.0..=1.0).contains(score),
894 "{id:?} scored {score} outside [0, 1]"
895 );
896 }
897 }
898
899 #[test]
903 fn cores_do_not_occupy_rank_positions() {
904 let core = fid("changed.rs", 1);
905 let ctx = fid("ctx.rs", 1);
906 let cores: FxHashSet<FragmentId> = std::iter::once(core.clone()).collect();
907
908 let with_core = scores(&[(core.clone(), 1.0), (ctx.clone(), 0.3)]);
909 let without_core = scores(&[(ctx.clone(), 0.3)]);
910
911 assert_eq!(
912 rank_positions(&with_core, &ballot(&with_core), &cores).get(&ctx),
913 rank_positions(&without_core, &ballot(&without_core), &cores).get(&ctx),
914 "a core shifted the rank of a context fragment"
915 );
916 }
917
918 #[test]
921 fn equal_component_scores_rank_deterministically() {
922 let cores: FxHashSet<FragmentId> = FxHashSet::default();
923 let ids: Vec<FragmentId> = (0..8).map(|i| fid(&format!("f{i}.rs"), 1)).collect();
924 let tied: FxHashMap<FragmentId, f64> = ids.iter().map(|i| (i.clone(), 0.5)).collect();
925
926 let baseline = rank_positions(&tied, &ballot(&tied), &cores);
927 for _ in 0..8 {
928 let again: FxHashMap<FragmentId, f64> =
929 ids.iter().rev().map(|i| (i.clone(), 0.5)).collect();
930 assert_eq!(rank_positions(&again, &ballot(&again), &cores), baseline);
931 }
932
933 let mut sorted = ids.clone();
935 sorted.sort();
936 for (expected_rank, id) in sorted.iter().enumerate() {
937 assert_eq!(baseline[id], expected_rank + 1);
938 }
939 }
940
941 #[test]
944 fn non_positive_component_scores_are_not_ranked() {
945 let kept = fid("kept.rs", 1);
946 let zero = fid("zero.rs", 1);
947 let cores: FxHashSet<FragmentId> = FxHashSet::default();
948 let component = scores(&[(kept.clone(), 0.3), (zero.clone(), 0.0)]);
949
950 let ranks = rank_positions(&component, &ballot(&component), &cores);
951 assert!(ranks.contains_key(&kept));
952 assert!(
953 !ranks.contains_key(&zero),
954 "a zero-scored fragment was ranked"
955 );
956
957 let fused = fuse_reciprocal_ranks(&[(&component, &ballot(&component))], &cores, 60.0);
958 assert!(
959 !fused.contains_key(&zero),
960 "a zero-scored fragment was fused"
961 );
962 }
963
964 #[test]
967 fn a_larger_k_flattens_the_gap_between_adjacent_ranks() {
968 let first = fid("first.rs", 1);
969 let second = fid("second.rs", 1);
970 let cores: FxHashSet<FragmentId> = FxHashSet::default();
971 let component = scores(&[(first.clone(), 0.9), (second.clone(), 0.8)]);
972
973 let sharp = fuse_reciprocal_ranks(&[(&component, &ballot(&component))], &cores, 1.0);
974 let flat = fuse_reciprocal_ranks(&[(&component, &ballot(&component))], &cores, 60.0);
975
976 assert!(
978 flat[&second] > sharp[&second],
979 "k did not flatten adjacent ranks: {} vs {}",
980 flat[&second],
981 sharp[&second]
982 );
983 }
984
985 #[test]
992 fn a_component_cannot_vote_for_what_its_own_filters_rejected() {
993 let good = fid("good.rs", 1);
994 let rejected = fid("garbage.rs", 1);
995 let cores: FxHashSet<FragmentId> = FxHashSet::default();
996 let component = scores(&[(rejected.clone(), 0.9), (good.clone(), 0.1)]);
999 let admitted: FxHashSet<FragmentId> = std::iter::once(good.clone()).collect();
1000
1001 let fused = fuse_reciprocal_ranks(&[(&component, &admitted)], &cores, 60.0);
1002
1003 assert!(
1004 !fused.contains_key(&rejected),
1005 "a fragment the component filtered out still earned fused mass {:?}",
1006 fused.get(&rejected)
1007 );
1008 assert!(
1009 fused.contains_key(&good),
1010 "the admitted fragment lost its vote"
1011 );
1012 }
1013
1014 #[test]
1015 fn a_percentile_is_a_position_in_the_full_population_not_the_admitted_subset() {
1016 let cores: FxHashSet<FragmentId> = FxHashSet::default();
1017 let mut entries: Vec<(FragmentId, f64)> = (0..9)
1021 .map(|i| (fid("strong.rs", i + 1), 1.0 - i as f64 * 0.05))
1022 .collect();
1023 let weak = fid("weak.rs", 100);
1024 entries.push((weak.clone(), 0.01));
1025 let rel = scores(&entries);
1026 let admitted: FxHashSet<FragmentId> = std::iter::once(weak.clone()).collect();
1027
1028 let (pct, _) = percentiles(&rel, &admitted, &cores, 5);
1029
1030 assert_eq!(pct.len(), 1, "only admitted fragments may receive a value");
1031 let p = pct[&weak];
1032 assert!(
1033 p <= 0.2,
1034 "the weakest of ten scored {p}, reading as strong because the CDF \
1035 was estimated over the admitted subset"
1036 );
1037 }
1038
1039 #[test]
1040 fn a_rejected_fragment_gets_no_percentile_at_all() {
1041 let cores: FxHashSet<FragmentId> = FxHashSet::default();
1042 let good = fid("good.rs", 1);
1043 let rejected = fid("garbage.rs", 1);
1044 let rel = scores(&[(rejected.clone(), 0.9), (good.clone(), 0.1)]);
1045 let admitted: FxHashSet<FragmentId> = std::iter::once(good.clone()).collect();
1046
1047 let (pct, top) = percentiles(&rel, &admitted, &cores, 5);
1048
1049 assert!(
1050 !pct.contains_key(&rejected),
1051 "the component's veto was lost"
1052 );
1053 assert!(
1054 !top.contains(&rejected),
1055 "a rejected fragment cannot sit in the component's top-k"
1056 );
1057 assert!(pct.contains_key(&good));
1058 }
1059
1060 #[test]
1061 fn quantile_map_restores_the_reference_distribution_in_fused_order() {
1062 let reference = vec![0.01, 0.02, 0.05, 0.4, 1.9];
1064 let a = fid("a.rs", 1);
1065 let b = fid("b.rs", 1);
1066 let c = fid("c.rs", 1);
1067 let fused = scores(&[(a.clone(), 0.9), (b.clone(), 0.5), (c.clone(), 0.1)]);
1070
1071 let mapped = quantile_map_to(&reference, &fused);
1072
1073 assert_eq!(mapped[&a], 1.9, "the fused top must take the reference max");
1074 assert_eq!(
1075 mapped[&c], 0.01,
1076 "the fused bottom must take the reference min"
1077 );
1078 assert!(
1079 mapped[&a] > mapped[&b] && mapped[&b] > mapped[&c],
1080 "the fused order was not preserved"
1081 );
1082 }
1083
1084 #[test]
1085 fn quantile_map_over_the_same_population_is_the_identity_on_values() {
1086 let ids: Vec<FragmentId> = (0..5).map(|i| fid("f.rs", i + 1)).collect();
1090 let ego_scores = [0.02, 0.07, 0.11, 0.55, 0.9];
1091 let mut reference: Vec<f64> = ego_scores.to_vec();
1092 reference.sort_by(f64::total_cmp);
1093 let fused = scores(
1095 &ids.iter()
1096 .zip([0.2, 0.4, 0.6, 0.8, 1.0])
1097 .map(|(id, p)| (id.clone(), p))
1098 .collect::<Vec<_>>(),
1099 );
1100
1101 let mapped = quantile_map_to(&reference, &fused);
1102
1103 for (id, expected) in ids.iter().zip(ego_scores) {
1104 assert_eq!(
1105 mapped[id], expected,
1106 "same-population quantile map must reproduce the component's own values"
1107 );
1108 }
1109 }
1110
1111 #[test]
1112 fn a_strategy_is_created_for_every_scoring_kind() {
1113 for mode in [
1115 crate::mode::ScoringMode::Ego,
1116 crate::mode::ScoringMode::Ppr,
1117 crate::mode::ScoringMode::Bm25,
1118 crate::mode::ScoringMode::Rrf,
1119 ] {
1120 let config = PipelineConfig::from_mode(mode);
1121 let strategy = create_scoring_strategy(&config);
1122 let empty: Vec<Fragment> = Vec::new();
1123 let result =
1124 strategy.score_and_filter(&empty, &FxHashSet::default(), &[], None, None, None);
1125 assert!(
1126 result.filtered_fragments.is_empty(),
1127 "{mode:?} invented fragments from an empty universe"
1128 );
1129 }
1130 let _ = FragmentKind::Function;
1132 }
1133}