Skip to main content

_diffctx/
select.rs

1use std::cmp::Ordering;
2use std::collections::BinaryHeap;
3use std::sync::Arc;
4
5use rustc_hash::{FxHashMap, FxHashSet};
6
7use crate::config::limits::UTILITY;
8use crate::config::selection::selection;
9use crate::interval::IntervalIndex;
10use crate::types::{Fragment, FragmentId};
11use crate::utility::needs::InformationNeed;
12use crate::utility::scoring::{
13    UtilityState, apply_fragment, compute_density, marginal_gain, utility_value,
14};
15
16const SENTINEL_TOKEN_COUNT: u32 = 1_000_000_000;
17
18#[derive(Debug, Clone, Copy, PartialEq, Eq)]
19pub enum SelectionReason {
20    TopK,
21    NoCandidates,
22    BudgetExhausted,
23    NoUtility,
24    StoppedByTau,
25    BestSingleton,
26}
27
28impl SelectionReason {
29    pub fn as_str(&self) -> &'static str {
30        match self {
31            Self::TopK => "topk",
32            Self::NoCandidates => "no_candidates",
33            Self::BudgetExhausted => "budget_exhausted",
34            Self::NoUtility => "no_utility",
35            Self::StoppedByTau => "stopped_by_tau",
36            Self::BestSingleton => "best_singleton",
37        }
38    }
39}
40
41pub struct SelectionResult {
42    pub selected: Vec<Fragment>,
43    pub reason: SelectionReason,
44    pub used_tokens: u32,
45    pub utility: f64,
46    /// Greedy iterations actually executed (number of `apply_fragment`
47    /// calls in `run_greedy_loop_heap`). Diagnoses lazy-heap blowup:
48    /// expected ≈ output size, pathological ≫ output size when
49    /// stale-version rejections dominate.
50    pub greedy_iters: usize,
51    /// Additive certificate for adaptive stopping: an upper bound
52    /// (`tau * peak_density * remaining_budget`) on the utility that
53    /// continuing the same greedy to the feasibility frontier could
54    /// still have added. 0 when the loop ended for any other reason.
55    pub stopping_certificate: f64,
56}
57
58struct HeapEntry {
59    neg_density: f64,
60    frag_id: FragmentId,
61    version: u32,
62}
63
64impl PartialEq for HeapEntry {
65    fn eq(&self, other: &Self) -> bool {
66        self.neg_density.to_bits() == other.neg_density.to_bits() && self.frag_id == other.frag_id
67    }
68}
69
70impl Eq for HeapEntry {}
71
72impl PartialOrd for HeapEntry {
73    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
74        Some(self.cmp(other))
75    }
76}
77
78impl Ord for HeapEntry {
79    fn cmp(&self, other: &Self) -> Ordering {
80        other
81            .neg_density
82            .total_cmp(&self.neg_density)
83            .then_with(|| other.frag_id.cmp(&self.frag_id))
84    }
85}
86
87struct SelectionState {
88    selected: Vec<Fragment>,
89    selected_ids: IntervalIndex,
90    remaining_budget: u32,
91    utility_state: UtilityState,
92}
93
94fn drop_redundant_signatures(candidates: &[Fragment], budget: u32) -> Vec<Fragment> {
95    let mut full_token_by_loc: FxHashMap<(Arc<str>, u32), u32> = FxHashMap::default();
96    for f in candidates {
97        if !f.kind.is_signature() {
98            // Keep the LARGEST co-located full fragment, not the last one seen.
99            // Two non-signature fragments can share a start line (a class header
100            // `Definition` at [10,12] and the full class at [10,300]), and with a
101            // plain `insert` whichever came last in the candidate vec won the
102            // slot. When the small header won, the class's stub was filtered out
103            // as "redundant" precisely when the full class did not fit and the
104            // stub was its only affordable representation — and the outcome
105            // depended on vec order rather than on anything meaningful.
106            full_token_by_loc
107                .entry((f.id.path.clone(), f.start_line()))
108                .and_modify(|t| *t = (*t).max(f.token_count))
109                .or_insert(f.token_count);
110        }
111    }
112    candidates
113        .iter()
114        .filter(|f| {
115            if !f.kind.is_signature() {
116                return true;
117            }
118            let key = (f.id.path.clone(), f.start_line());
119            full_token_by_loc
120                .get(&key)
121                .copied()
122                .unwrap_or(SENTINEL_TOKEN_COUNT)
123                > budget
124        })
125        .cloned()
126        .collect()
127}
128
129fn compute_r_cap(
130    rel: &FxHashMap<FragmentId, f64>,
131    core_ids: Option<&FxHashSet<FragmentId>>,
132) -> f64 {
133    let values: Vec<f64> = rel
134        .iter()
135        .filter(|(fid, v)| **v > 0.0 && core_ids.map_or(true, |c| !c.contains(*fid)))
136        .map(|(_, v)| *v)
137        .collect();
138
139    if values.len() < 2 {
140        return if let Some(&v) = values.first() {
141            v.max(selection().r_cap_min)
142        } else {
143            1.0
144        };
145    }
146
147    let mut sorted = values.clone();
148    sorted.sort_by(|a, b| a.total_cmp(b));
149    let mid = sorted.len() / 2;
150    let med = if sorted.len() % 2 == 0 {
151        (sorted[mid - 1] + sorted[mid]) / 2.0
152    } else {
153        sorted[mid]
154    };
155
156    let mean: f64 = values.iter().sum::<f64>() / values.len() as f64;
157    let variance: f64 =
158        values.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / (values.len() - 1) as f64;
159    let std = variance.sqrt();
160
161    (med + UTILITY.r_cap_sigma * std).max(1e-9)
162}
163
164fn build_signature_lookup(
165    fragments: &[Fragment],
166    core_fragments: &[Fragment],
167    core_excerpts: Option<&FxHashMap<FragmentId, Fragment>>,
168) -> FxHashMap<FragmentId, Fragment> {
169    let mut sig_by_loc: FxHashMap<(Arc<str>, u32), Fragment> = FxHashMap::default();
170    for f in fragments {
171        if f.kind.is_signature() {
172            sig_by_loc.insert((f.id.path.clone(), f.start_line()), f.clone());
173        }
174    }
175    let mut sig_lookup = FxHashMap::default();
176    for cf in core_fragments {
177        let key = (cf.id.path.clone(), cf.start_line());
178        if let Some(sig) = sig_by_loc.get(&key) {
179            sig_lookup.insert(cf.id.clone(), sig.clone());
180            continue;
181        }
182        // Kinds without a signature (chunk, section) fall back to the excerpt
183        // around the hunk; without it an oversized core is skipped outright and
184        // the change signal disappears from the output (#103).
185        if let Some(excerpt) = core_excerpts.and_then(|e| e.get(&cf.id)) {
186            sig_lookup.insert(cf.id.clone(), excerpt.clone());
187        }
188    }
189    sig_lookup
190}
191
192fn select_core_fragments(
193    core_fragments: &[Fragment],
194    rel: &FxHashMap<FragmentId, f64>,
195    needs: &[InformationNeed],
196    state: &mut SelectionState,
197    budget_tokens: u32,
198    sig_lookup: &FxHashMap<FragmentId, Fragment>,
199) {
200    let core_budget = (budget_tokens as f64 * selection().core_budget_fraction) as u32;
201    // Counter for cores placed; the first pass keeps `core_used <= core_budget`,
202    // but the rescue pass below intentionally allows it to exceed `core_budget`
203    // up to `budget_tokens`. Don't assume the tighter bound past this scope.
204    let mut core_used = 0u32;
205
206    let mut sorted_core: Vec<&Fragment> = core_fragments.iter().collect();
207    sorted_core.sort_by(|a, b| {
208        let ra = rel.get(&a.id).copied().unwrap_or(0.0);
209        let rb = rel.get(&b.id).copied().unwrap_or(0.0);
210        rb.total_cmp(&ra)
211    });
212
213    let place_fragment =
214        |frag: &Fragment, core_used: &mut u32, state: &mut SelectionState, rel_score: f64| {
215            state.selected.push(frag.clone());
216            state.selected_ids.add_id(&frag.id);
217            state.remaining_budget = state.remaining_budget.saturating_sub(frag.token_count);
218            *core_used += frag.token_count;
219            apply_fragment(frag, rel_score, needs, &mut state.utility_state);
220        };
221
222    let mut skipped: Vec<&Fragment> = Vec::new();
223    for frag in &sorted_core {
224        if state.selected_ids.is_superset_of(frag) {
225            continue;
226        }
227        if core_used + frag.token_count > core_budget {
228            if let Some(sig) = sig_lookup.get(&frag.id) {
229                if !state.selected_ids.contains(&sig.id)
230                    && core_used + sig.token_count <= core_budget
231                {
232                    let rel_score = rel.get(&frag.id).copied().unwrap_or(0.0);
233                    place_fragment(sig, &mut core_used, state, rel_score);
234                    continue;
235                }
236            }
237            skipped.push(frag);
238            continue;
239        }
240
241        let rel_score = rel.get(&frag.id).copied().unwrap_or(0.0);
242        place_fragment(frag, &mut core_used, state, rel_score);
243    }
244
245    // Bug #2 fix: cores that didn't fit the core_budget reservation must not be
246    // demoted to ordinary greedy candidates without a chance to be placed first.
247    // Sweep skipped cores cheapest-first against the *full* remaining budget
248    // (not just the core slice) so seeds aren't dropped purely because the
249    // highest-relevance core happened to be heavy.
250    if !skipped.is_empty() {
251        skipped.sort_by(|a, b| a.token_count.cmp(&b.token_count));
252        for frag in skipped {
253            if state.remaining_budget == 0 {
254                break;
255            }
256            if state.selected_ids.is_superset_of(frag) {
257                continue;
258            }
259            if frag.token_count <= state.remaining_budget {
260                let rel_score = rel.get(&frag.id).copied().unwrap_or(0.0);
261                place_fragment(frag, &mut core_used, state, rel_score);
262            } else if let Some(sig) = sig_lookup.get(&frag.id) {
263                if !state.selected_ids.contains(&sig.id)
264                    && sig.token_count <= state.remaining_budget
265                {
266                    let rel_score = rel.get(&frag.id).copied().unwrap_or(0.0);
267                    place_fragment(sig, &mut core_used, state, rel_score);
268                }
269            }
270        }
271    }
272}
273
274fn build_initial_heap(
275    candidates: &[Fragment],
276    rel: &FxHashMap<FragmentId, f64>,
277    needs: &[InformationNeed],
278    state: &UtilityState,
279    id_to_frag: &mut FxHashMap<FragmentId, Fragment>,
280) -> BinaryHeap<HeapEntry> {
281    let mut heap = BinaryHeap::new();
282    for frag in candidates {
283        if frag.token_count > 0 {
284            let density = compute_density(
285                frag,
286                rel.get(&frag.id).copied().unwrap_or(0.0),
287                needs,
288                state,
289            );
290            heap.push(HeapEntry {
291                neg_density: -density,
292                frag_id: frag.id.clone(),
293                version: 0,
294            });
295            id_to_frag.insert(frag.id.clone(), frag.clone());
296        }
297    }
298    heap
299}
300
301fn find_best_candidate_heap(
302    heap: &mut BinaryHeap<HeapEntry>,
303    current_version: u32,
304    id_to_frag: &FxHashMap<FragmentId, Fragment>,
305    selected_ids: &IntervalIndex,
306    remaining_budget: u32,
307    rel: &FxHashMap<FragmentId, f64>,
308    needs: &[InformationNeed],
309    state: &UtilityState,
310) -> (Option<Fragment>, f64, u32) {
311    let cv = current_version;
312    while let Some(entry) = heap.pop() {
313        let frag = match id_to_frag.get(&entry.frag_id) {
314            Some(f) => f,
315            None => continue,
316        };
317        if frag.token_count > remaining_budget {
318            continue;
319        }
320        if selected_ids.overlaps(frag) {
321            continue;
322        }
323        if entry.version < cv {
324            let new_density = compute_density(
325                frag,
326                rel.get(&frag.id).copied().unwrap_or(0.0),
327                needs,
328                state,
329            );
330            heap.push(HeapEntry {
331                neg_density: -new_density,
332                frag_id: frag.id.clone(),
333                version: cv,
334            });
335            continue;
336        }
337        let actual_density = -entry.neg_density;
338        if actual_density <= 0.0 {
339            return (None, 0.0, cv);
340        }
341        return (Some(frag.clone()), actual_density, cv + 1);
342    }
343    (None, 0.0, cv)
344}
345
346fn find_best_singleton(
347    non_core: &[Fragment],
348    base_selected_ids: &IntervalIndex,
349    base_budget: u32,
350    rel: &FxHashMap<FragmentId, f64>,
351    needs: &[InformationNeed],
352    base_state: &UtilityState,
353) -> (Option<Fragment>, f64) {
354    let mut best_singleton = None;
355    let mut best_gain = 0.0;
356    for f in non_core {
357        if f.token_count > base_budget {
358            continue;
359        }
360        if base_selected_ids.overlaps(f) {
361            continue;
362        }
363        let gain = marginal_gain(f, rel.get(&f.id).copied().unwrap_or(0.0), needs, base_state);
364        if gain > best_gain {
365            best_gain = gain;
366            best_singleton = Some(f.clone());
367        }
368    }
369    (best_singleton, best_gain)
370}
371
372/// Paper-aligned strict Khuller H₁: `argmax_{f ∈ F, |f| ≤ B} U({f})`.
373/// Iterates the FULL ground set (including core), gates on the FULL
374/// budget B (not the residual after partial-core packing), and evaluates
375/// each candidate as a lone selection against an empty utility state.
376///
377/// This is the comparator that, together with the greedy chain, gives
378/// the `(1-1/e)/2` approximation guarantee from Khuller-Moss-Naor 1999
379/// for monotone submodular maximization under a knapsack constraint.
380/// Restricting H₁ to `non_core` with budget `B − cost(packed_core)`
381/// (the prior `find_best_singleton`) is a strict relaxation that
382/// excludes any single high-utility fragment with `|f| > B − β_core·B`.
383fn find_best_singleton_full_set(
384    fragments: &[Fragment],
385    budget_tokens: u32,
386    rel: &FxHashMap<FragmentId, f64>,
387    needs: &[InformationNeed],
388    empty_state: &UtilityState,
389) -> (Option<Fragment>, f64) {
390    let mut best = None;
391    let mut best_gain = 0.0;
392    for f in fragments {
393        if f.token_count == 0 || f.token_count > budget_tokens {
394            continue;
395        }
396        let gain = marginal_gain(
397            f,
398            rel.get(&f.id).copied().unwrap_or(0.0),
399            needs,
400            empty_state,
401        );
402        if gain > best_gain {
403            best_gain = gain;
404            best = Some(f.clone());
405        }
406    }
407    (best, best_gain)
408}
409
410fn init_selection_state(
411    core_ids: &FxHashSet<FragmentId>,
412    rel: &FxHashMap<FragmentId, f64>,
413    budget_tokens: u32,
414    file_importance: Option<&FxHashMap<Arc<str>, f64>>,
415) -> SelectionState {
416    let mut utility_state = UtilityState::default();
417    utility_state.r_cap = compute_r_cap(rel, Some(core_ids));
418    utility_state.changed_dirs = core_ids
419        .iter()
420        .filter_map(|cid| {
421            std::path::Path::new(cid.path.as_ref())
422                .parent()
423                .map(|p| p.to_path_buf())
424        })
425        .collect();
426    if let Some(fi) = file_importance {
427        utility_state.file_importance.clone_from(fi);
428    }
429    SelectionState {
430        selected: Vec::new(),
431        selected_ids: IntervalIndex::new(),
432        remaining_budget: budget_tokens,
433        utility_state,
434    }
435}
436
437fn run_greedy_loop_heap(
438    heap: &mut BinaryHeap<HeapEntry>,
439    id_to_frag: &FxHashMap<FragmentId, Fragment>,
440    state: &mut SelectionState,
441    rel: &FxHashMap<FragmentId, f64>,
442    needs: &[InformationNeed],
443    tau: f64,
444    _initial_budget: u32,
445) -> (usize, f64, usize) {
446    let mut current_version = 0u32;
447    let mut peak_density: f64 = 0.0;
448    let mut loop_iters: usize = 0;
449
450    while !heap.is_empty() && state.remaining_budget > 0 {
451        loop_iters += 1;
452        let (best_frag, best_density, new_version) = find_best_candidate_heap(
453            heap,
454            current_version,
455            id_to_frag,
456            &state.selected_ids,
457            state.remaining_budget,
458            rel,
459            needs,
460            &state.utility_state,
461        );
462        current_version = new_version;
463
464        let best_frag = match best_frag {
465            Some(f) => f,
466            None => break,
467        };
468        if best_density <= 0.0 {
469            break;
470        }
471
472        if best_density > peak_density {
473            peak_density = best_density;
474        } else if peak_density > 0.0 && best_density < tau * peak_density {
475            break;
476        }
477
478        state.selected.push(best_frag.clone());
479        state.selected_ids.add_id(&best_frag.id);
480        state.remaining_budget = state.remaining_budget.saturating_sub(best_frag.token_count);
481        let rel_score = rel.get(&best_frag.id).copied().unwrap_or(0.0);
482        apply_fragment(&best_frag, rel_score, needs, &mut state.utility_state);
483    }
484
485    let threshold = tau * peak_density;
486    (state.selected.len(), threshold, loop_iters)
487}
488
489fn setup_and_select_core(
490    fragments: &[Fragment],
491    core_ids: &FxHashSet<FragmentId>,
492    rel: &FxHashMap<FragmentId, f64>,
493    needs: &[InformationNeed],
494    budget_tokens: u32,
495    file_importance: Option<&FxHashMap<Arc<str>, f64>>,
496    core_excerpts: Option<&FxHashMap<FragmentId, Fragment>>,
497) -> (SelectionState, Vec<Fragment>, Vec<Fragment>, bool) {
498    let mut core_fragments: Vec<Fragment> = fragments
499        .iter()
500        .filter(|f| core_ids.contains(&f.id))
501        .cloned()
502        .collect();
503    core_fragments.sort_by(|a, b| {
504        let ta = if a.token_count > 0 {
505            a.token_count
506        } else {
507            SENTINEL_TOKEN_COUNT
508        };
509        let tb = if b.token_count > 0 {
510            b.token_count
511        } else {
512            SENTINEL_TOKEN_COUNT
513        };
514        ta.cmp(&tb)
515            .then(a.line_count().cmp(&b.line_count()))
516            .then(a.start_line().cmp(&b.start_line()))
517    });
518
519    let non_core_fragments: Vec<Fragment> = fragments
520        .iter()
521        .filter(|f| !core_ids.contains(&f.id))
522        .cloned()
523        .collect();
524
525    let sig_lookup = build_signature_lookup(fragments, &core_fragments, core_excerpts);
526    let mut state = init_selection_state(core_ids, rel, budget_tokens, file_importance);
527    select_core_fragments(
528        &core_fragments,
529        rel,
530        needs,
531        &mut state,
532        budget_tokens,
533        &sig_lookup,
534    );
535
536    let selected_core_ids: FxHashSet<FragmentId> =
537        state.selected.iter().map(|f| f.id.clone()).collect();
538    let skipped_core: Vec<FragmentId> = core_ids
539        .iter()
540        .filter(|id| !selected_core_ids.contains(*id))
541        .cloned()
542        .collect();
543
544    let mut non_core_with_skipped = non_core_fragments;
545    if !skipped_core.is_empty() {
546        let skipped_set: FxHashSet<FragmentId> = skipped_core.into_iter().collect();
547        for cf in &core_fragments {
548            if skipped_set.contains(&cf.id) {
549                non_core_with_skipped.push(cf.clone());
550            }
551        }
552    }
553
554    let should_return_early = state.remaining_budget == 0;
555    let selected_copy = state.selected.clone();
556    (
557        state,
558        non_core_with_skipped,
559        selected_copy,
560        should_return_early,
561    )
562}
563
564pub fn lazy_greedy_select(
565    fragments: Vec<Fragment>,
566    core_ids: &FxHashSet<FragmentId>,
567    rel: &FxHashMap<FragmentId, f64>,
568    needs: &[InformationNeed],
569    budget_tokens: u32,
570    tau: f64,
571    file_importance: Option<&FxHashMap<Arc<str>, f64>>,
572    core_excerpts: Option<&FxHashMap<FragmentId, Fragment>>,
573) -> SelectionResult {
574    if fragments.is_empty() {
575        return SelectionResult {
576            selected: Vec::new(),
577            reason: SelectionReason::NoCandidates,
578            used_tokens: 0,
579            utility: 0.0,
580            greedy_iters: 0,
581            stopping_certificate: 0.0,
582        };
583    }
584
585    let (mut state, non_core_fragments, _selected_core, should_return_early) =
586        setup_and_select_core(
587            &fragments,
588            core_ids,
589            rel,
590            needs,
591            budget_tokens,
592            file_importance,
593            core_excerpts,
594        );
595
596    if should_return_early {
597        let used = budget_tokens - state.remaining_budget;
598        return SelectionResult {
599            selected: state.selected,
600            reason: SelectionReason::BudgetExhausted,
601            used_tokens: used,
602            utility: utility_value(&state.utility_state),
603            greedy_iters: 0,
604            stopping_certificate: 0.0,
605        };
606    }
607
608    let base_state = state.utility_state.copy();
609    let base_selected = state.selected.clone();
610    let base_budget = state.remaining_budget;
611
612    let candidates: Vec<Fragment> = non_core_fragments
613        .iter()
614        .filter(|f| !state.selected_ids.overlaps(f))
615        .cloned()
616        .collect();
617    let candidates = drop_redundant_signatures(&candidates, state.remaining_budget);
618
619    let mut id_to_frag: FxHashMap<FragmentId, Fragment> = FxHashMap::default();
620    let mut heap = build_initial_heap(
621        &candidates,
622        rel,
623        needs,
624        &state.utility_state,
625        &mut id_to_frag,
626    );
627
628    let (_, threshold, greedy_iters) = run_greedy_loop_heap(
629        &mut heap,
630        &id_to_frag,
631        &mut state,
632        rel,
633        needs,
634        tau,
635        budget_tokens,
636    );
637
638    let greedy_utility = utility_value(&state.utility_state);
639
640    let mut base_selected_ids = IntervalIndex::new();
641    for f in &base_selected {
642        base_selected_ids.add_id(&f.id);
643    }
644
645    let (best_singleton, best_gain) = find_best_singleton(
646        &non_core_fragments,
647        &base_selected_ids,
648        base_budget,
649        rel,
650        needs,
651        &base_state,
652    );
653
654    let empty_state = init_selection_state(core_ids, rel, budget_tokens, file_importance);
655    let (full_singleton, full_singleton_gain) = find_best_singleton_full_set(
656        &fragments,
657        budget_tokens,
658        rel,
659        needs,
660        &empty_state.utility_state,
661    );
662
663    let mut best_alt_utility = greedy_utility;
664    let mut best_alt: Option<(Vec<Fragment>, u32)> = None;
665
666    if let Some(ref singleton) = best_singleton {
667        let u = utility_value(&base_state) + best_gain;
668        if u > best_alt_utility {
669            best_alt_utility = u;
670            let used = budget_tokens - (base_budget - singleton.token_count);
671            let mut sel = base_selected.clone();
672            sel.push(singleton.clone());
673            best_alt = Some((sel, used));
674        }
675    }
676
677    if let Some(ref full) = full_singleton {
678        let u = utility_value(&empty_state.utility_state) + full_singleton_gain;
679        // Additive on top of the core selection, never a replacement for it.
680        // This branch used to return `vec![full]` outright, so a single heavy
681        // fragment whose standalone utility beat the greedy chain's discarded
682        // every changed-code fragment — the one thing the output exists to
683        // carry. `ensure_changed_files_represented` could not reliably undo it
684        // either: it only had `budget - full.token_count` left and only picks a
685        // fragment that fits. The two utilities are also measured from
686        // different baselines (this one from an empty state, `greedy_utility`
687        // from the core base), so the comparison can only ever be a heuristic
688        // nudge — not grounds for dropping the core.
689        if u > best_alt_utility && full.token_count <= base_budget {
690            best_alt_utility = u;
691            let mut sel = base_selected.clone();
692            if !sel.iter().any(|f| f.id == full.id) {
693                sel.push(full.clone());
694            }
695            let used = budget_tokens - (base_budget - full.token_count);
696            best_alt = Some((sel, used));
697        }
698    }
699
700    if let Some((sel, used)) = best_alt {
701        return SelectionResult {
702            selected: sel,
703            reason: SelectionReason::BestSingleton,
704            used_tokens: used,
705            utility: best_alt_utility,
706            greedy_iters,
707            stopping_certificate: 0.0,
708        };
709    }
710
711    let used = budget_tokens - state.remaining_budget;
712    let reason = if state.remaining_budget == 0 {
713        SelectionReason::BudgetExhausted
714    } else if greedy_utility <= 0.0 {
715        SelectionReason::NoUtility
716    } else if state.selected.is_empty() || state.selected.len() == base_selected.len() {
717        SelectionReason::NoCandidates
718    } else if threshold > 0.0 && !heap.is_empty() {
719        SelectionReason::StoppedByTau
720    } else {
721        SelectionReason::NoCandidates
722    };
723
724    let stopping_certificate = if matches!(reason, SelectionReason::StoppedByTau) {
725        threshold * f64::from(state.remaining_budget)
726    } else {
727        0.0
728    };
729
730    SelectionResult {
731        selected: state.selected,
732        reason,
733        used_tokens: used,
734        utility: greedy_utility,
735        greedy_iters,
736        stopping_certificate,
737    }
738}
739
740#[cfg(test)]
741mod tests {
742    use super::*;
743    use crate::types::FragmentKind;
744
745    fn frag(path: &str, start: u32, end: u32, kind: FragmentKind, tokens: u32) -> Fragment {
746        let mut identifiers = FxHashSet::default();
747        identifiers.insert(format!("sym_{path}_{start}"));
748        Fragment {
749            id: FragmentId::new(Arc::from(path), start, end),
750            kind,
751            content: Arc::from(format!("// {path}:{start}-{end}\n")),
752            identifiers,
753            token_count: tokens,
754            symbol_name: Some(format!("sym_{start}")),
755        }
756    }
757
758    fn rel_map(frags: &[Fragment], score: f64) -> FxHashMap<FragmentId, f64> {
759        frags.iter().map(|f| (f.id.clone(), score)).collect()
760    }
761
762    fn cost_of(selected: &[Fragment]) -> u32 {
763        selected.iter().map(|f| f.token_count).sum()
764    }
765
766    /// The budget is a hard contract (`cost(C) <= B`); four separate call sites
767    /// gate on it and none of them was asserted. A `pick_smallest_fitting` that
768    /// returns a non-fitting candidate is one `if` away, and the oracle corpus
769    /// can never catch it because its budget is always >=2.5x the whole repo.
770    #[test]
771    fn selection_never_exceeds_the_budget() {
772        let frags: Vec<Fragment> = (0..12)
773            .map(|i| {
774                frag(
775                    "a.rs",
776                    1 + i * 50,
777                    40 + i * 50,
778                    FragmentKind::Function,
779                    30 + i * 17,
780                )
781            })
782            .collect();
783        let core: FxHashSet<FragmentId> = std::iter::once(frags[0].id.clone()).collect();
784        let rel = rel_map(&frags, 0.7);
785
786        for budget in [1u32, 7, 30, 31, 60, 200, 1_000] {
787            let result =
788                lazy_greedy_select(frags.clone(), &core, &rel, &[], budget, 0.12, None, None);
789            assert!(
790                cost_of(&result.selected) <= budget,
791                "budget {budget} overrun: cost {} via {:?}",
792                cost_of(&result.selected),
793                result.reason
794            );
795            assert!(
796                result.used_tokens <= budget,
797                "reported used_tokens {} exceeds budget {budget}",
798                result.used_tokens
799            );
800        }
801    }
802
803    #[test]
804    fn selected_fragments_never_overlap_and_are_never_duplicated() {
805        let frags: Vec<Fragment> = (0..8)
806            .map(|i| frag("a.rs", 1 + i * 10, 12 + i * 10, FragmentKind::Function, 25))
807            .collect();
808        let core: FxHashSet<FragmentId> = std::iter::once(frags[0].id.clone()).collect();
809        let rel = rel_map(&frags, 0.9);
810        let result = lazy_greedy_select(frags, &core, &rel, &[], 400, 0.12, None, None);
811
812        let ids: FxHashSet<FragmentId> = result.selected.iter().map(|f| f.id.clone()).collect();
813        assert_eq!(
814            ids.len(),
815            result.selected.len(),
816            "duplicate fragment selected"
817        );
818    }
819
820    /// `find_best_singleton_full_set` used to return `vec![full]`, discarding
821    /// every core fragment. The core IS the changed code, so a selection that
822    /// drops all of it answers a different question than the one asked.
823    #[test]
824    fn a_winning_singleton_never_evicts_the_core_selection() {
825        // A heavy, highly relevant non-core fragment is the shape that makes the
826        // full-set singleton win.
827        let core_frag = frag("changed.rs", 1, 8, FragmentKind::Function, 20);
828        let heavy = frag("other.rs", 1, 400, FragmentKind::Class, 900);
829        let filler = frag("other.rs", 500, 520, FragmentKind::Function, 40);
830        let frags = vec![core_frag.clone(), heavy.clone(), filler];
831
832        let core: FxHashSet<FragmentId> = std::iter::once(core_frag.id.clone()).collect();
833        let mut rel = FxHashMap::default();
834        rel.insert(core_frag.id.clone(), 0.05);
835        rel.insert(heavy.id.clone(), 1.0);
836        rel.insert(frags[2].id.clone(), 0.1);
837
838        let result = lazy_greedy_select(frags, &core, &rel, &[], 2_000, 0.12, None, None);
839        assert!(
840            result.selected.iter().any(|f| core.contains(&f.id)),
841            "no core fragment survived; reason was {:?}",
842            result.reason
843        );
844        assert!(cost_of(&result.selected) <= 2_000);
845    }
846
847    #[test]
848    fn empty_ground_set_reports_no_candidates() {
849        let result = lazy_greedy_select(
850            Vec::new(),
851            &FxHashSet::default(),
852            &FxHashMap::default(),
853            &[],
854            1_000,
855            0.12,
856            None,
857            None,
858        );
859        assert!(result.selected.is_empty());
860        assert_eq!(result.reason, SelectionReason::NoCandidates);
861        assert_eq!(result.used_tokens, 0);
862    }
863
864    /// Keyed on `(path, start_line)`, this used to be last-write-wins, so a
865    /// small co-located sibling could delete the stub that was the only
866    /// affordable representation of an oversized fragment.
867    #[test]
868    fn drop_redundant_signatures_is_independent_of_candidate_order() {
869        let header = frag("a.rs", 10, 12, FragmentKind::Definition, 40);
870        let whole = frag("a.rs", 10, 300, FragmentKind::Class, 4_000);
871        let stub = frag("a.rs", 10, 11, FragmentKind::ClassSignature, 15);
872
873        let forward =
874            drop_redundant_signatures(&[header.clone(), whole.clone(), stub.clone()], 500);
875        let backward = drop_redundant_signatures(&[whole, header, stub], 500);
876
877        let kinds = |v: &[Fragment]| -> Vec<FragmentKind> { v.iter().map(|f| f.kind).collect() };
878        assert!(
879            kinds(&forward).contains(&FragmentKind::ClassSignature),
880            "the stub for an unaffordable class was dropped: {:?}",
881            kinds(&forward)
882        );
883        let mut a: Vec<FragmentKind> = kinds(&forward);
884        let mut b: Vec<FragmentKind> = kinds(&backward);
885        a.sort_by_key(|k| format!("{k:?}"));
886        b.sort_by_key(|k| format!("{k:?}"));
887        assert_eq!(a, b, "verdict depended on candidate order");
888    }
889
890    #[test]
891    fn drop_redundant_signatures_removes_a_stub_whose_full_fragment_fits() {
892        let whole = frag("a.rs", 10, 40, FragmentKind::Class, 100);
893        let stub = frag("a.rs", 10, 11, FragmentKind::ClassSignature, 15);
894        let kept = drop_redundant_signatures(&[whole, stub], 500);
895        assert!(
896            !kept.iter().any(|f| f.kind.is_signature()),
897            "stub survived even though the full fragment fits the budget"
898        );
899    }
900    /// tau is the adaptive stop: once a candidate's density falls below
901    /// `tau * peak_density` the loop stops instead of spending the rest of the
902    /// budget. The whole oracle corpus runs at tau=0.0, which makes the
903    /// predicate unreachable, so deleting the rule failed no test. This runs at
904    /// the shipped default so the assertion covers the real operating point.
905    #[test]
906    fn tau_stops_the_greedy_loop_before_the_budget_is_spent() {
907        // Descending relevance against escalating cost gives sharply
908        // descending density, which is what the stop rule reacts to.
909        let frags: Vec<Fragment> = (0..6)
910            .map(|i| {
911                frag(
912                    "a.rs",
913                    1 + i * 20,
914                    10 + i * 20,
915                    FragmentKind::Function,
916                    20 + i * i * 120,
917                )
918            })
919            .collect();
920        let core: FxHashSet<FragmentId> = std::iter::once(frags[0].id.clone()).collect();
921        let mut rel = FxHashMap::default();
922        for (i, f) in frags.iter().enumerate() {
923            rel.insert(f.id.clone(), 1.0 / (1.0 + 3.0 * i as f64));
924        }
925
926        let budget = 10_000;
927        let default_tau = lazy_greedy_select(
928            frags.clone(),
929            &core,
930            &rel,
931            &[],
932            budget,
933            crate::config::limits::DEFAULT_STOPPING_THRESHOLD,
934            None,
935            None,
936        );
937        let no_tau = lazy_greedy_select(frags, &core, &rel, &[], budget, 0.0, None, None);
938
939        assert_eq!(
940            default_tau.reason,
941            SelectionReason::StoppedByTau,
942            "the adaptive stop did not fire at the shipped default"
943        );
944        assert!(
945            default_tau.selected.len() < no_tau.selected.len(),
946            "tau={} selected {} fragments, same as tau=0.0 — the rule is inert",
947            crate::config::limits::DEFAULT_STOPPING_THRESHOLD,
948            default_tau.selected.len()
949        );
950        assert!(
951            default_tau.used_tokens < no_tau.used_tokens,
952            "the stop saved no budget: {} vs {}",
953            default_tau.used_tokens,
954            no_tau.used_tokens
955        );
956        assert!(
957            default_tau.stopping_certificate > 0.0,
958            "StoppedByTau must carry a positive certificate"
959        );
960        assert_eq!(
961            no_tau.stopping_certificate, 0.0,
962            "tau=0.0 cannot produce a stopping certificate"
963        );
964    }
965}