Skip to main content

_diffctx/
scoring.rs

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    /// Wall time spent constructing the typed dependency graph (edge
23    /// builders + dedup + hub suppression + per-source cap). Reported
24    /// separately so `scoring_ms` stays pure rank computation. Zero for
25    /// BM25 (no graph built).
26    pub graph_build_ms: f64,
27    /// PPR push-iteration was cut by `max_pushes_cap` before convergence.
28    /// Always false for non-PPR strategies (EGO/BM25).
29    pub ppr_truncated: bool,
30    pub ppr_forward_pushes: usize,
31    pub ppr_backward_pushes: usize,
32}
33
34/// The guard chain every graph-backed strategy ends with: drop what the graph
35/// says is unrelated, drop what scored zero, then cap per file.
36///
37/// EGO, PPR and RRF each spelled this out. BM25 spelled out its own copy of the
38/// middle step instead of calling it — semantically the same today, and exactly
39/// the shape that let the two test-file classifiers drift apart (#182).
40fn 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        // Deliberately NOT `finish_scoring`: there is no graph here, so the
294        // structural guard has nothing to judge with. The other two steps are
295        // the shared ones rather than a local re-spelling of the same predicate.
296        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
308/// Reciprocal-rank fusion of the structural (EGO) and lexical (BM25) signals.
309///
310/// The two are complementary rather than redundant: on genuine
311/// retrieval the lexical component alone outranks the deployed
312/// graph+lexical mixture, and their score-free union raises the reachable
313/// recall well above either — a miscalibrated-mixture signature. RRF
314/// fuses on ranks only, so neither component's score scale can dominate
315/// the other, which is exactly the failure the weighted mixture had.
316pub 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
330/// A component's ballot: the fragments it admitted, ordered by its own score.
331///
332/// `admitted` is the component's `filtered_fragments`, not its whole score map.
333/// Rank fusion is defined over the result lists retrievers return, and each
334/// component's guards are the only place an *absolute* judgement survives —
335/// a rank cannot express "this scored near zero". Ranking the full score map
336/// instead lets a component vote for fragments its own filters rejected, and
337/// reciprocal rank then promotes that tail: BM25 scores anything sharing a
338/// generic token, so garbage landed at a respectable rank and earned real
339/// fused mass. Measured on the oracle corpus, that cost 97 cases against EGO
340/// on precision (`forbidden_rate >= 90%` on 91 of them) while recall held.
341fn 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    // Ties broken by id so the rank list — and therefore every fused
352    // score — is independent of hash-map iteration order.
353    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    // Cores sit at the top of the scale, matching every other strategy —
380    // downstream `r_cap` normalisation and the absolute relevance gates read
381    // these values, so the fused range has to stay [0, 1]. Note they do not
382    // strictly dominate: max-normalisation already put the best non-core at
383    // exactly 1.0, so it ties with the cores rather than sitting below them.
384    // `r_cap` excludes cores when it computes its spread, so the tie is benign
385    // there — but do not read this as a guarantee that cores rank first.
386    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        // The union re-admits paths that EGO's structural guards dropped
449        // (hub noise, generic-config-only code), because BM25 has no graph
450        // to judge them by. The guards are re-applied and the per-file cap
451        // recomputed against the fused scores, since each component capped
452        // against its own.
453        //
454        // Measured caveat, not a claim of soundness: re-applying them does NOT
455        // make the union a net win. On the oracle corpus RRF loses 97 cases to
456        // EGO and gains 18, all on precision, and restricting candidates to
457        // EGO's admitted set recovers only 28 of the 82 (#125).
458        //
459        // A second, unmeasured degree of freedom lives here: each component
460        // already applied `cap_context_fragments` (30/file) against its own
461        // scores before voting, so a fragment the fused ranking would have kept
462        // can have been capped away before it ever reached the ballot. The cap
463        // is per file and the losses are cross-file, so this is unlikely to be
464        // the 97 — but it has not been isolated.
465        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
477/// Percentile fusion: the same two signals as RRF, blended on their empirical
478/// distribution position rather than on rank alone.
479///
480/// RRF converts each component to a pure rank, which throws away the magnitude
481/// that says "this scored near zero". Measured on the oracle corpus that costs
482/// 97 cases against EGO and wins 18, all on precision: BM25 gives any
483/// generic-token match a small positive score, and `1/(k + rank)` promotes that
484/// noise to real fused mass (#125).
485///
486/// The probability-integral transform keeps the position. A fragment in the 5th
487/// percentile of a component contributes 0.05 from it, not `1/(k + 12)`. Two
488/// signals that disagree therefore cannot manufacture a strong candidate out of
489/// two weak opinions, which is precisely what the rank form allowed.
490///
491/// `score = blend * PIT(ego) + (1 - blend) * PIT(bm25) + bonus * [both in top-k]`
492///
493/// The agreement term is what fusion is actually for — a fragment both signals
494/// rank highly is more trustworthy than either alone — and it is additive and
495/// small so it breaks ties rather than deciding the ranking.
496pub 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
515/// Empirical-CDF position in `[0, 1]` for every admitted, positively-scored
516/// fragment, plus the set that sits in the component's own top-k.
517///
518/// Ties share a percentile: two fragments a component cannot separate must not
519/// be separated here either, or the blend would invent a preference the signal
520/// never expressed.
521///
522/// The CDF is estimated over everything the component scored positively, while
523/// only the fragments it *admitted* receive a value. Those are two different
524/// questions and conflating them was a defect: a percentile read off a
525/// component's admitted set is a position within that set, and the two
526/// components' admitted sets differ by an order of magnitude on a real repo
527/// (BM25 admits a handful, EGO hundreds). Blending a position-among-6 with a
528/// position-among-300 as if they were the same quantity is not a fusion of the
529/// two signals. The admission veto itself is kept — a fragment a component
530/// rejected still contributes nothing from it (#125, `091c4db3`) — because the
531/// component's own guards are the only place an absolute judgement survives.
532fn 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    // Descending by score, id as the tie-break so the traversal is independent
544    // of hash-map iteration order.
545    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    // Ablation (not a shipped mode): `DIFFCTX_PIT_TRANSFORM=maxnorm` fuses the
555    // components on their own score shape, rescaled to a common [0, 1], instead
556    // of on distributional position. It is the control that isolates the
557    // transform, because a linear rescale leaves every downstream
558    // magnitude-reading rule invariant: `r_cap = median + sigma*std` scales with
559    // the data, so `rel / r_cap` is unchanged. The percentile does not have that
560    // property, which is the whole point of comparing them.
561    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            // One run of equal scores shares the mean percentile of the run.
576            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            // Position 0 is the strongest, so invert: the best fragment gets ~1.0.
581            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    // Top-k is drawn from the admitted fragments: the agreement bonus asks
593    // "do both components rank this highly", and a fragment a component
594    // rejected is not ranked highly by it.
595    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
605/// Map a fused ranking back onto the reference component's own score
606/// distribution, preserving order.
607///
608/// Selection does not read `rel` as an ordering alone. `compute_r_cap` takes
609/// `median + sigma*std` of the score cloud and the utility uses
610/// `(rel / r_cap).min(1.0)`, so the *shape* of the distribution decides how many
611/// candidates saturate. EGO's raw scores are strongly right-skewed (hop decay
612/// puts most mass near zero), which makes `r_cap` small and the saturation
613/// meaningful. A percentile is uniform on [0, 1] by construction: its median is
614/// ~0.5 and `median + 2*std` lands above the maximum, so nothing saturates and
615/// the selector runs in a regime nothing was calibrated for.
616///
617/// That is a unit mismatch, not a tuning problem, so the fix is to restore the
618/// units rather than to re-tune `r_cap_sigma` and `tau` per transform. Fusion
619/// then decides the order — which is what fusion is for — and the selector sees
620/// the score cloud it was built against.
621///
622/// A consequence worth stating because it doubles as the correctness gate: at
623/// `blend = 1.0` with no agreement bonus the fused order is EGO's order over
624/// EGO's own admitted set, so this maps every fragment back to its exact EGO
625/// score and the mode must reproduce EGO bit-for-bit.
626fn quantile_map_to(
627    reference: &[f64],
628    fused: &FxHashMap<FragmentId, f64>,
629) -> FxHashMap<FragmentId, f64> {
630    // A fused score of zero means no component endorsed the fragment.
631    // Mapping it anyway would hand it a positive reference value and resurrect
632    // a candidate `filter_positive_relevance` is required to drop — at
633    // blend=1.0 that alone broke the EGO-equivalence gate (390 vs 371).
634    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        // A run of equal fused scores is a tie the fusion never resolved, so
650        // the run shares one mapped value — its midpoint slot — rather than
651        // being fanned across adjacent reference values by id order.
652        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            // A fragment only one component admitted contributes 0 from the
728            // other — that is the point. Under RRF an absent component was
729            // simply silent; here it is an explicit "this signal ranks you at
730            // the bottom", which is what stops one weak opinion carrying a
731            // fragment.
732            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        // `DIFFCTX_PIT_SHAPE=flat` keeps the fused percentiles as the scores the
742        // selector sees. That is the pre-`quantile_map_to` behaviour, retained
743        // so the transform's cost stays measurable rather than only argued.
744        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            // Cores keep EGO's own values rather than a pinned 1.0. `r_cap`
767            // excludes cores, but the utility does not, and a synthetic 1.0 on
768            // EGO's raw scale is a different number from the one EGO assigns.
769            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    /// Every scored fragment counts as admitted, which is the pre-#125-fix
809    /// behaviour these properties were written against. Tests that care about
810    /// the admission gate itself build the ballot explicitly.
811    fn ballot(rel: &FxHashMap<FragmentId, f64>) -> FxHashSet<FragmentId> {
812        rel.keys().cloned().collect()
813    }
814
815    /// The property RRF exists for, and the reason `k` damps the top of each
816    /// list: agreement between the two signals beats a single signal's first
817    /// place. A weighted mixture cannot express this without calibrating the two
818    /// score scales against each other — which is the failure RRF replaces.
819    #[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        // `ego_only` is rank 1 in ego and absent from bm25; `agreed` is only
826        // rank 2 in each, but present in both.
827        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    /// Only ranks may cross between the components. If a raw score leaked in,
844    /// one signal's scale could dominate the other — the miscalibrated-mixture
845    /// behaviour the mode was added to avoid.
846    #[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    /// Downstream `r_cap` normalisation and the absolute relevance gates read
875    /// these values, so the fused range has to stay within [0, 1] with the cores
876    /// at the top.
877    #[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    /// Cores are ranked separately (they are always placed first), so letting
900    /// them consume rank slots would push every context fragment down and change
901    /// the fused scores for reasons unrelated to relevance.
902    #[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    /// Hash-map iteration order must not reach the fused scores: equal
919    /// component scores are ranked by fragment id.
920    #[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        // Ascending id order is the tie-break, so ranks follow the sorted ids.
934        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    /// A zero or negative component score is "not a candidate", not "ranked
942    /// last": including it would hand it a reciprocal-rank contribution.
943    #[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    /// `DIFFCTX_RRF_K` is a documented knob; a larger k flattens the reciprocal
965    /// curve, which is what makes agreement outweigh a single top hit.
966    #[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        // Both are max-normalised, so compare the runner-up's share of the top.
977        assert!(
978            flat[&second] > sharp[&second],
979            "k did not flatten adjacent ranks: {} vs {}",
980            flat[&second],
981            sharp[&second]
982        );
983    }
984
985    /// The gate that ranks cannot express. A component scores far more
986    /// fragments than it admits — BM25 gives anything sharing a generic token a
987    /// small positive score — and a rank list is purely ordinal, so the moment a
988    /// rejected fragment appears in it the reciprocal rank hands it real fused
989    /// mass. Fusing whole score maps instead of the returned result lists cost
990    /// 97 oracle cases on precision.
991    #[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        // The rejected fragment outscores the admitted one, so if it is ranked
997        // at all it takes rank 1 and the top of the normalised scale with it.
998        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        // Nine strong fragments the component scored but did not admit, plus
1018        // one weak admitted straggler. Its percentile must say "bottom of the
1019        // component's world", not "top of the admitted set of one".
1020        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        // Reference: a skewed cloud like EGO's (mass near zero).
1063        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        // Fused scores are uniform-ish percentiles; only their order may
1068        // survive the mapping.
1069        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        // blend=1.0, bonus=0: the fused order IS ego's order over ego's own
1087        // admitted set, so mapping back onto ego's sorted scores must return
1088        // exactly those scores — the property the corpus gate checks end to end.
1089        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        // Fused percentiles in the same order as the ego scores.
1094        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        // A new mode that forgets its arm here would silently score as another.
1114        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        // Keeps `FragmentKind` in scope for the helper above.
1131        let _ = FragmentKind::Function;
1132    }
1133}