1use std::path::PathBuf;
2use std::sync::Arc;
3
4use rustc_hash::{FxHashMap, FxHashSet};
5
6use crate::config::limits::UTILITY;
7use crate::config::needs::NEEDS;
8use crate::types::Fragment;
9use crate::utility::needs::{InformationNeed, match_strength_typed};
10
11pub struct UtilityState {
12 pub max_rel: FxHashMap<(String, String), f64>,
13 pub priorities: FxHashMap<(String, String), f64>,
14 pub structural_sum: f64,
15 pub eta: f64,
16 pub structural_bonus_weight: f64,
17 pub r_cap: f64,
18 pub changed_dirs: FxHashSet<PathBuf>,
19 pub proximity_decay: f64,
20 pub file_importance: FxHashMap<Arc<str>, f64>,
21}
22
23impl Default for UtilityState {
24 fn default() -> Self {
25 Self {
26 max_rel: FxHashMap::default(),
27 priorities: FxHashMap::default(),
28 structural_sum: 0.0,
29 eta: UTILITY.eta,
30 structural_bonus_weight: UTILITY.structural_bonus_weight,
31 r_cap: 1.0,
32 changed_dirs: FxHashSet::default(),
33 proximity_decay: UTILITY.proximity_decay,
34 file_importance: FxHashMap::default(),
35 }
36 }
37}
38
39impl UtilityState {
40 pub fn copy(&self) -> Self {
41 Self {
42 max_rel: self.max_rel.clone(),
43 priorities: self.priorities.clone(),
44 structural_sum: self.structural_sum,
45 eta: self.eta,
46 structural_bonus_weight: self.structural_bonus_weight,
47 r_cap: self.r_cap,
48 changed_dirs: self.changed_dirs.clone(),
49 proximity_decay: self.proximity_decay,
50 file_importance: self.file_importance.clone(),
51 }
52 }
53}
54
55fn phi(x: f64) -> f64 {
56 if x > 0.0 { x.sqrt() } else { 0.0 }
57}
58
59fn augmented_score(m: f64, rel_score: f64, state: &UtilityState) -> f64 {
60 let r_norm = if state.r_cap > 0.0 {
61 (rel_score / state.r_cap).min(1.0)
62 } else {
63 0.0
64 };
65 m + state.eta * r_norm
66}
67
68fn needs_from_identifiers(frag: &Fragment) -> Vec<InformationNeed> {
69 frag.identifiers
70 .iter()
71 .map(|c| InformationNeed {
72 need_type: "definition".to_string(),
73 symbol: c.clone(),
74 scope: None,
75 priority: NEEDS.identifier_default_priority,
76 })
77 .collect()
78}
79
80struct GainResult {
81 gain: f64,
82 has_match: bool,
83 need_updates: Vec<((String, String), f64, f64)>,
84 diversity_bonus: f64,
85 structural_bonus: f64,
86}
87
88fn diversity_bonus(
89 needs: &[InformationNeed],
90 rel_score: f64,
91 gain: f64,
92 state: &UtilityState,
93) -> f64 {
94 if needs.is_empty() || rel_score < NEEDS.min_rel_for_bonus {
95 return 0.0;
96 }
97 if gain <= 0.0 {
98 return 0.0;
99 }
100 let total_covered: f64 = needs
101 .iter()
102 .map(|n| {
103 state
104 .max_rel
105 .get(&(n.need_type.clone(), n.symbol.clone()))
106 .copied()
107 .unwrap_or(0.0)
108 .min(1.0)
109 })
110 .sum();
111 let unsatisfied = (1.0 - total_covered / needs.len().max(1) as f64).max(0.0);
112 rel_score * NEEDS.relatedness_bonus * unsatisfied
113}
114
115fn compute_gain_core(
116 frag: &Fragment,
117 rel_score: f64,
118 needs: &[InformationNeed],
119 state: &UtilityState,
120 use_state_priorities: bool,
121) -> GainResult {
122 let effective: Vec<InformationNeed>;
123 let needs_slice = if !needs.is_empty() {
124 needs
125 } else {
126 effective = needs_from_identifiers(frag);
127 &effective
128 };
129
130 let mut result = GainResult {
131 gain: 0.0,
132 has_match: false,
133 need_updates: Vec::new(),
134 diversity_bonus: 0.0,
135 structural_bonus: 0.0,
136 };
137
138 if needs_slice.is_empty() {
139 return result;
140 }
141
142 for need in needs_slice {
143 let m = match_strength_typed(frag, need);
144 if m <= 0.0 {
145 continue;
146 }
147 let mut m_eff = m;
148 if need.need_type == "impact" && !state.file_importance.is_empty() {
149 let path_arc: Arc<str> = Arc::from(frag.path());
150 m_eff *= state.file_importance.get(&path_arc).copied().unwrap_or(1.0);
151 }
152 result.has_match = true;
153 let a_fz = augmented_score(m_eff, rel_score, state);
154 let nkey = (need.need_type.clone(), need.symbol.clone());
155 let old_max = state.max_rel.get(&nkey).copied().unwrap_or(0.0);
156 let new_max = old_max.max(a_fz);
157 let priority = if use_state_priorities {
158 state
159 .priorities
160 .get(&nkey)
161 .copied()
162 .unwrap_or(need.priority)
163 } else {
164 need.priority
165 };
166 result.gain += priority * (phi(new_max) - phi(old_max));
167 result.need_updates.push((nkey, new_max, need.priority));
168 }
169
170 result.diversity_bonus = diversity_bonus(needs, rel_score, result.gain, state);
171
172 if result.has_match {
173 let r_norm = if state.r_cap > 0.0 {
174 (rel_score / state.r_cap).min(1.0)
175 } else {
176 0.0
177 };
178 result.structural_bonus = state.structural_bonus_weight * r_norm;
179 }
180
181 result
182}
183
184pub fn marginal_gain(
185 frag: &Fragment,
186 rel_score: f64,
187 needs: &[InformationNeed],
188 state: &UtilityState,
189) -> f64 {
190 let result = compute_gain_core(frag, rel_score, needs, state, false);
191 result.gain + result.diversity_bonus + result.structural_bonus
192}
193
194pub fn apply_fragment(
195 frag: &Fragment,
196 rel_score: f64,
197 needs: &[InformationNeed],
198 state: &mut UtilityState,
199) {
200 let result = compute_gain_core(frag, rel_score, needs, state, true);
201 for (nkey, new_max, priority) in result.need_updates {
202 state.max_rel.insert(nkey.clone(), new_max);
203 let current = state.priorities.get(&nkey).copied().unwrap_or(0.0);
204 state.priorities.insert(nkey, current.max(priority));
205 }
206 state.structural_sum += result.diversity_bonus + result.structural_bonus;
207}
208
209fn dir_distance(d1: &std::path::Path, d2: &std::path::Path) -> usize {
210 let p1: Vec<_> = d1.components().collect();
211 let p2: Vec<_> = d2.components().collect();
212 let mut common = 0;
213 for (a, b) in p1.iter().zip(p2.iter()) {
214 if a == b {
215 common += 1;
216 } else {
217 break;
218 }
219 }
220 (p1.len() - common) + (p2.len() - common)
221}
222
223fn proximity_factor(frag_path: &str, changed_dirs: &FxHashSet<PathBuf>, alpha: f64) -> f64 {
224 if changed_dirs.is_empty() {
225 return 1.0;
226 }
227 let frag_dir = std::path::Path::new(frag_path)
228 .parent()
229 .unwrap_or_else(|| std::path::Path::new(""));
230 let min_dist = changed_dirs
231 .iter()
232 .map(|d| dir_distance(frag_dir, d))
233 .min()
234 .unwrap_or(0);
235 if min_dist == 0 {
236 return 1.0;
237 }
238 1.0 / (1.0 + alpha * min_dist as f64)
239}
240
241pub fn compute_density(
242 frag: &Fragment,
243 rel_score: f64,
244 needs: &[InformationNeed],
245 state: &UtilityState,
246) -> f64 {
247 if frag.token_count == 0 {
248 return 0.0;
249 }
250 let gain = marginal_gain(frag, rel_score, needs, state);
251 let pf = proximity_factor(frag.path(), &state.changed_dirs, state.proximity_decay);
252 gain * pf / frag.token_count as f64
253}
254
255pub fn utility_value(state: &UtilityState) -> f64 {
256 let u1: f64 = state
257 .max_rel
258 .iter()
259 .map(|(sym, v)| {
260 let p = state.priorities.get(sym).copied().unwrap_or(1.0);
261 p * phi(*v)
262 })
263 .sum();
264 u1 + state.structural_sum
265}
266
267#[cfg(test)]
268mod paper_claim_tests {
269 use super::*;
276 use crate::types::{FragmentId, FragmentKind};
277
278 fn frag(symbol: &str, mentions: &[&str], tokens: u32) -> Fragment {
279 Fragment {
280 id: FragmentId::new(Arc::from(format!("syn/{symbol}.rs")), 1, 10),
281 kind: FragmentKind::Function,
282 content: Arc::from(""),
283 identifiers: mentions.iter().map(|s| s.to_lowercase()).collect(),
284 token_count: tokens,
285 symbol_name: Some(symbol.to_lowercase()),
286 }
287 }
288
289 fn need(symbol: &str, priority: f64) -> InformationNeed {
290 InformationNeed {
291 need_type: "definition".to_string(),
292 symbol: symbol.to_lowercase(),
293 scope: None,
294 priority,
295 }
296 }
297
298 fn xorshift(state: &mut u64) -> u64 {
299 *state ^= *state << 13;
300 *state ^= *state >> 7;
301 *state ^= *state << 17;
302 *state
303 }
304
305 fn random_subset_indices(n: usize, fraction: f64, rng: &mut u64) -> Vec<usize> {
306 (0..n)
307 .filter(|_| (xorshift(rng) % 1000) as f64 / 1000.0 < fraction)
308 .collect()
309 }
310
311 fn build_state_from(
312 fragments: &[Fragment],
313 indices: &[usize],
314 rels: &[f64],
315 needs: &[InformationNeed],
316 ) -> UtilityState {
317 let mut state = UtilityState::default();
318 for &i in indices {
319 apply_fragment(&fragments[i], rels[i], needs, &mut state);
320 }
321 state
322 }
323
324 #[test]
325 fn claim_4_cost_is_modular() {
326 let fragments = vec![
327 frag("alpha", &["beta"], 100),
328 frag("beta", &["gamma"], 250),
329 frag("gamma", &["alpha"], 75),
330 frag("delta", &["alpha", "beta"], 500),
331 ];
332
333 let cost_of =
334 |idx: &[usize]| -> u32 { idx.iter().map(|&i| fragments[i].token_count).sum() };
335
336 let mut rng = 0xCAFEBABE_u64;
337 for _ in 0..1000 {
338 let a: Vec<usize> = random_subset_indices(fragments.len(), 0.5, &mut rng);
339 let b: Vec<usize> = random_subset_indices(fragments.len(), 0.5, &mut rng);
340 let union: Vec<usize> = {
341 let s: FxHashSet<usize> = a.iter().chain(b.iter()).copied().collect();
342 let mut v: Vec<usize> = s.into_iter().collect();
343 v.sort();
344 v
345 };
346 let intersection: Vec<usize> = {
347 let sa: FxHashSet<usize> = a.iter().copied().collect();
348 let mut v: Vec<usize> = b.iter().copied().filter(|i| sa.contains(i)).collect();
349 v.sort();
350 v
351 };
352 let lhs = cost_of(&union);
353 let rhs = cost_of(&a) + cost_of(&b) - cost_of(&intersection);
354 assert_eq!(
355 lhs, rhs,
356 "cost not modular: cost(A∪B)={lhs} ≠ cost(A)+cost(B)-cost(A∩B)={rhs}"
357 );
358 }
359 }
360
361 #[test]
362 fn claim_1a_submodularity_holds_on_random_instances() {
363 let symbols = [
364 "foo", "bar", "baz", "qux", "alpha", "beta", "gamma", "delta",
365 ];
366 let fragments: Vec<Fragment> = symbols
367 .iter()
368 .enumerate()
369 .map(|(i, s)| {
370 let mentions: Vec<&str> = symbols.iter().take(i).copied().collect();
371 frag(s, &mentions, 100 + (i as u32) * 50)
372 })
373 .collect();
374 let needs: Vec<InformationNeed> = symbols.iter().map(|s| need(s, 1.0)).collect();
375 let rels: Vec<f64> = (0..fragments.len()).map(|i| 0.3 + 0.1 * i as f64).collect();
376
377 let mut rng = 0xDEADBEEF_u64;
378 let mut violations = 0;
379 let trials = 500;
380
381 for _ in 0..trials {
382 let s_idx: Vec<usize> = random_subset_indices(fragments.len(), 0.3, &mut rng);
383 let mut t_idx: Vec<usize> = s_idx.clone();
384 for j in 0..fragments.len() {
385 if !t_idx.contains(&j) && (xorshift(&mut rng) % 100) < 40 {
386 t_idx.push(j);
387 }
388 }
389 let candidates: Vec<usize> = (0..fragments.len())
390 .filter(|i| !t_idx.contains(i))
391 .collect();
392 if candidates.is_empty() {
393 continue;
394 }
395 let x = candidates[(xorshift(&mut rng) as usize) % candidates.len()];
396
397 let state_s = build_state_from(&fragments, &s_idx, &rels, &needs);
398 let state_t = build_state_from(&fragments, &t_idx, &rels, &needs);
399
400 let marg_s = marginal_gain(&fragments[x], rels[x], &needs, &state_s);
401 let marg_t = marginal_gain(&fragments[x], rels[x], &needs, &state_t);
402
403 if marg_s + 1e-9 < marg_t {
404 violations += 1;
405 eprintln!("Submodularity violation: marg_S={marg_s:.6}, marg_T={marg_t:.6}, x={x}");
406 }
407 }
408
409 assert_eq!(
410 violations, 0,
411 "Submodularity (Theorem 1) violated in {violations}/{trials} trials"
412 );
413 }
414
415 #[test]
416 fn claim_1b_saturation_zero_marginal_for_duplicate_definition() {
417 let f1 = frag("foo", &[], 100);
418 let f2 = frag("foo", &[], 100);
419 let needs = vec![need("foo", 1.0)];
420
421 let mut state = UtilityState::default();
422 let m1 = marginal_gain(&f1, 1.0, &needs, &state);
423 apply_fragment(&f1, 1.0, &needs, &mut state);
424 let m2 = marginal_gain(&f2, 1.0, &needs, &state);
425
426 assert!(
427 m1 > 0.0,
428 "first fragment must have positive marginal, got {m1}"
429 );
430 let bonus_only = state.structural_bonus_weight;
431 assert!(
432 m2 <= bonus_only + 1e-9,
433 "duplicate definition must have ≈zero core gain marginal (only structural_bonus={bonus_only}); got {m2}"
434 );
435 }
436
437 #[test]
438 fn claim_2a_selected_fragments_never_overlap_within_a_path() {
439 let mut fragments: Vec<Fragment> = Vec::new();
440 for path_idx in 0..3 {
441 for chunk_idx in 0..4 {
442 let start = (chunk_idx * 20 + 1) as u32;
443 let end = (chunk_idx * 20 + 15) as u32;
444 let f = Fragment {
445 id: FragmentId::new(Arc::from(format!("p{path_idx}.rs")), start, end),
446 kind: FragmentKind::Function,
447 content: Arc::from(""),
448 identifiers: ["alpha", "beta"].iter().map(|s| s.to_string()).collect(),
449 token_count: 80 + (path_idx * 30 + chunk_idx * 10) as u32,
450 symbol_name: Some(format!("p{path_idx}_chunk_{chunk_idx}")),
451 };
452 fragments.push(f);
453 }
454 for chunk_idx in 0..3 {
455 let start = (chunk_idx * 20 + 5) as u32;
456 let end = (chunk_idx * 20 + 25) as u32;
457 let f = Fragment {
458 id: FragmentId::new(Arc::from(format!("p{path_idx}.rs")), start, end),
459 kind: FragmentKind::Function,
460 content: Arc::from(""),
461 identifiers: ["alpha", "beta"].iter().map(|s| s.to_string()).collect(),
462 token_count: 90,
463 symbol_name: Some(format!("p{path_idx}_overlap_{chunk_idx}")),
464 };
465 fragments.push(f);
466 }
467 }
468 let needs = vec![need("alpha", 1.0), need("beta", 1.0)];
469 let mut rels: FxHashMap<FragmentId, f64> = FxHashMap::default();
470 for (i, f) in fragments.iter().enumerate() {
471 rels.insert(f.id.clone(), 0.3 + 0.02 * i as f64);
472 }
473 let core_ids = FxHashSet::default();
474 let budget = 2048;
475
476 let result = crate::select::lazy_greedy_select(
477 fragments.clone(),
478 &core_ids,
479 &rels,
480 &needs,
481 budget,
482 0.08,
483 None,
484 None,
485 None,
486 );
487
488 for (i, fi) in result.selected.iter().enumerate() {
489 for fj in result.selected.iter().skip(i + 1) {
490 if fi.id.path != fj.id.path {
491 continue;
492 }
493 let overlap =
494 fi.id.start_line <= fj.id.end_line && fj.id.start_line <= fi.id.end_line;
495 assert!(
496 !overlap,
497 "Matroid (interval) constraint violated: \
498 {}:{}-{} overlaps with {}:{}-{}",
499 fi.id.path,
500 fi.id.start_line,
501 fi.id.end_line,
502 fj.id.path,
503 fj.id.start_line,
504 fj.id.end_line
505 );
506 }
507 }
508 }
509
510 #[test]
511 fn claim_5_greedy_meets_khuller_bound_against_brute_force_optimal() {
512 let fragments: Vec<Fragment> = [
513 ("foo", &[][..], 100u32),
514 ("bar", &["foo"][..], 80),
515 ("baz", &["bar"][..], 250),
516 ("qux", &["foo", "bar"][..], 60),
517 ("alpha", &["baz"][..], 400),
518 ("beta", &["qux"][..], 90),
519 ]
520 .iter()
521 .map(|(name, mentions, tokens)| frag(name, mentions, *tokens))
522 .collect();
523
524 let needs: Vec<InformationNeed> = ["foo", "bar", "baz", "qux", "alpha", "beta"]
525 .iter()
526 .map(|s| need(s, 1.0))
527 .collect();
528
529 let mut rels: FxHashMap<FragmentId, f64> = FxHashMap::default();
530 for (i, f) in fragments.iter().enumerate() {
531 rels.insert(f.id.clone(), 0.4 + 0.1 * i as f64);
532 }
533 let core_ids: FxHashSet<FragmentId> = FxHashSet::default();
534 let budget: u32 = 400;
535
536 let n = fragments.len();
537 let mut optimal: f64 = 0.0;
538 for mask in 0u32..(1 << n) {
539 let cost: u32 = (0..n)
540 .filter(|i| mask & (1 << i) != 0)
541 .map(|i| fragments[i].token_count)
542 .sum();
543 if cost > budget {
544 continue;
545 }
546 let mut state = UtilityState::default();
547 for i in 0..n {
548 if mask & (1 << i) != 0 {
549 apply_fragment(&fragments[i], rels[&fragments[i].id], &needs, &mut state);
550 }
551 }
552 optimal = optimal.max(utility_value(&state));
553 }
554
555 let result = crate::select::lazy_greedy_select(
556 fragments.clone(),
557 &core_ids,
558 &rels,
559 &needs,
560 budget,
561 0.08,
562 None,
563 None,
564 None,
565 );
566 let ratio = result.utility / optimal;
567 let bound = 0.5 * (1.0 - 1.0_f64.exp().recip());
568 assert!(
569 ratio >= bound - 1e-6,
570 "Khuller bound violated: greedy/optimal = {ratio:.4} < {bound:.4} (greedy={}, optimal={})",
571 result.utility,
572 optimal
573 );
574 assert!(
575 ratio >= 0.5,
576 "Realistic-data expectation violated: ratio={ratio:.4} should be ≥ 0.5 on this instance"
577 );
578 }
579
580 #[test]
581 fn claim_9_importance_prior_preserves_submodularity_for_impact_needs() {
582 let mut fragments = vec![
583 frag("hub", &["service"], 200),
584 frag("client_a", &["hub"], 150),
585 frag("client_b", &["hub"], 150),
586 frag("client_c", &["hub"], 150),
587 ];
588 for f in &mut fragments {
589 f.identifiers.insert("service".to_string());
590 }
591 let needs = vec![InformationNeed {
592 need_type: "impact".to_string(),
593 symbol: "service".to_string(),
594 scope: None,
595 priority: 1.0,
596 }];
597 let rels = vec![0.9, 0.7, 0.5, 0.3];
598
599 let mut file_importance = FxHashMap::default();
600 file_importance.insert(Arc::from("syn/hub.rs"), 0.5);
601 file_importance.insert(Arc::from("syn/client_a.rs"), 1.0);
602 file_importance.insert(Arc::from("syn/client_b.rs"), 0.8);
603 file_importance.insert(Arc::from("syn/client_c.rs"), 0.6);
604
605 let mut rng = 0x12345678_u64;
606 for _ in 0..200 {
607 let s_idx: Vec<usize> = random_subset_indices(fragments.len(), 0.3, &mut rng);
608 let mut t_idx: Vec<usize> = s_idx.clone();
609 for j in 0..fragments.len() {
610 if !t_idx.contains(&j) && (xorshift(&mut rng) % 100) < 50 {
611 t_idx.push(j);
612 }
613 }
614 let candidates: Vec<usize> = (0..fragments.len())
615 .filter(|i| !t_idx.contains(i))
616 .collect();
617 if candidates.is_empty() {
618 continue;
619 }
620 let x = candidates[(xorshift(&mut rng) as usize) % candidates.len()];
621
622 let mut state_s = UtilityState {
623 file_importance: file_importance.clone(),
624 ..UtilityState::default()
625 };
626 for &i in &s_idx {
627 apply_fragment(&fragments[i], rels[i], &needs, &mut state_s);
628 }
629 let mut state_t = UtilityState {
630 file_importance: file_importance.clone(),
631 ..UtilityState::default()
632 };
633 for &i in &t_idx {
634 apply_fragment(&fragments[i], rels[i], &needs, &mut state_t);
635 }
636
637 let marg_s = marginal_gain(&fragments[x], rels[x], &needs, &state_s);
638 let marg_t = marginal_gain(&fragments[x], rels[x], &needs, &state_t);
639 assert!(
640 marg_s + 1e-9 >= marg_t,
641 "Importance prior breaks submodularity: marg_S={marg_s}, marg_T={marg_t}"
642 );
643 }
644 }
645}