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 pub greedy_iters: usize,
51 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 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 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 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 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
372fn 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 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 #[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 #[test]
824 fn a_winning_singleton_never_evicts_the_core_selection() {
825 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 #[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 #[test]
906 fn tau_stops_the_greedy_loop_before_the_budget_is_spent() {
907 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}