1use std::borrow::Cow;
2use std::path::{Path, PathBuf};
3use std::sync::OnceLock;
4
5use rayon::prelude::*;
6use rustc_hash::{FxHashMap, FxHashSet};
7
8use crate::config::bm25::BM25;
9use crate::token_corpus::{DocTokens, TokenCorpus};
10use crate::types::extract_identifier_list;
11
12pub struct DiscoveryContext {
13 pub root_dir: PathBuf,
14 pub changed_files: Vec<PathBuf>,
15 pub all_candidates: Vec<PathBuf>,
16 pub diff_text: String,
17 pub expansion_concepts: FxHashSet<String>,
18 pub file_cache: FxHashMap<PathBuf, String>,
19 pub token_corpus: OnceLock<TokenCorpus>,
20}
21
22impl DiscoveryContext {
23 pub fn read_file(&self, path: &Path) -> Option<Cow<'_, str>> {
24 if let Some(content) = self.file_cache.get(path) {
25 return Some(Cow::Borrowed(content.as_str()));
26 }
27 std::fs::read_to_string(path).ok().map(Cow::Owned)
28 }
29
30 pub fn shared_corpus(&self) -> &TokenCorpus {
31 self.token_corpus.get_or_init(|| TokenCorpus::build(self))
32 }
33}
34
35pub trait DiscoveryStrategy: Send + Sync {
36 fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf>;
37
38 fn name(&self) -> &'static str;
42
43 fn discover_attributed(
52 &self,
53 ctx: &DiscoveryContext,
54 ) -> (Vec<PathBuf>, Vec<(PathBuf, &'static str)>) {
55 let paths = self.discover(ctx);
56 let attribution = paths.iter().map(|p| (p.clone(), self.name())).collect();
57 (paths, attribution)
58 }
59}
60
61pub struct DefaultDiscovery;
62
63impl DiscoveryStrategy for DefaultDiscovery {
64 fn name(&self) -> &'static str {
65 "structural"
66 }
67
68 fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
69 let changed_set: FxHashSet<&Path> = ctx.changed_files.iter().map(|p| p.as_path()).collect();
70
71 let mut discovered = crate::edges::discover_all_related_files(
72 &ctx.changed_files,
73 &ctx.all_candidates,
74 Some(ctx.root_dir.as_path()),
75 Some(&ctx.file_cache),
76 );
77 discovered.retain(|p| !changed_set.contains(p.as_path()));
78
79 let rare_files = expand_by_rare_identifiers(ctx);
80 let existing: FxHashSet<PathBuf> = discovered.iter().cloned().collect();
81 for f in rare_files {
82 if !existing.contains(&f) {
83 discovered.push(f);
84 }
85 }
86
87 discovered
88 }
89}
90
91fn expand_by_rare_identifiers(ctx: &DiscoveryContext) -> Vec<PathBuf> {
92 let rare_threshold = crate::config::limits::LIMITS.rare_identifier_threshold;
93
94 let mut ident_to_files: FxHashMap<String, Vec<PathBuf>> = FxHashMap::default();
95 for (path, doc) in &ctx.shared_corpus().docs {
96 for ident in &ctx.expansion_concepts {
97 if doc.term_counts.contains_key(ident) {
98 ident_to_files
99 .entry(ident.clone())
100 .or_default()
101 .push(path.clone());
102 }
103 }
104 }
105
106 let mut result: Vec<PathBuf> = Vec::new();
107 let mut seen: FxHashSet<PathBuf> = FxHashSet::default();
108 for (_ident, files) in &ident_to_files {
109 if files.len() <= rare_threshold {
110 for f in files {
111 if seen.insert(f.clone()) {
112 result.push(f.clone());
113 }
114 }
115 }
116 }
117 result
118}
119
120pub struct TestFileDiscovery;
121
122const TEST_PREFIXES: &[&str] = &["test_", "spec_"];
123const TEST_SUFFIXES: &[&str] = &["_test", "_spec", ".test", ".spec", "-test", "-spec"];
124
125fn lowercase_stem(path: &Path) -> String {
126 path.file_stem()
127 .map(|s| s.to_string_lossy().to_lowercase())
128 .unwrap_or_default()
129}
130
131impl DiscoveryStrategy for TestFileDiscovery {
132 fn name(&self) -> &'static str {
133 "test_pairing"
134 }
135
136 fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
137 let changed_set: FxHashSet<&Path> = ctx.changed_files.iter().map(|p| p.as_path()).collect();
138 let mut test_stems: FxHashSet<String> = FxHashSet::default();
142 let mut sibling_stems: FxHashMap<PathBuf, FxHashSet<String>> = FxHashMap::default();
150
151 for f in &ctx.changed_files {
152 let stem = lowercase_stem(f);
153 if TEST_PREFIXES.iter().any(|p| stem.starts_with(p)) {
154 continue;
155 }
156 if TEST_SUFFIXES.iter().any(|s| stem.ends_with(s)) {
157 continue;
158 }
159 sibling_stems
160 .entry(f.parent().unwrap_or(Path::new("")).to_path_buf())
161 .or_default()
162 .insert(stem.clone());
163 for prefix in TEST_PREFIXES {
164 test_stems.insert(format!("{prefix}{stem}"));
165 }
166 for suffix in TEST_SUFFIXES {
167 test_stems.insert(format!("{stem}{suffix}"));
168 }
169 }
170
171 let mut discovered: Vec<PathBuf> = Vec::new();
172 for candidate in &ctx.all_candidates {
173 if changed_set.contains(candidate.as_path()) {
174 continue;
175 }
176 let stem = lowercase_stem(candidate);
177 let is_sibling = sibling_stems
178 .get(candidate.parent().unwrap_or(Path::new("")))
179 .is_some_and(|stems| stems.contains(&stem));
180 if test_stems.contains(&stem) || is_sibling {
181 discovered.push(candidate.clone());
182 }
183 }
184 discovered
185 }
186}
187
188pub struct BM25Discovery {
189 pub top_k: usize,
190}
191
192impl BM25Discovery {
193 pub fn new(top_k: usize) -> Self {
194 Self { top_k }
195 }
196
197 fn bm25_score(
198 doc: &DocTokens,
199 query_set: &FxHashSet<String>,
200 idf: &FxHashMap<String, f64>,
201 avgdl: f64,
202 ) -> f64 {
203 let dl = doc.total_len as f64;
204 let mut s = 0.0;
205 for t in query_set {
206 let freq = doc.term_counts.get(t).copied().unwrap_or(0) as f64;
207 if freq == 0.0 {
208 continue;
209 }
210 let idf_val = idf.get(t).copied().unwrap_or(0.0);
211 s += idf_val * (freq * BM25.k1)
212 / (freq + BM25.k1 * (1.0 - BM25.b + BM25.b * dl / avgdl));
213 }
214 s
215 }
216}
217
218impl DiscoveryStrategy for BM25Discovery {
219 fn name(&self) -> &'static str {
220 "lexical_bm25"
221 }
222
223 fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
224 let query_tokens = extract_identifier_list(&ctx.diff_text, BM25.min_query_token_length);
225 if query_tokens.is_empty() {
226 return Vec::new();
227 }
228 let query_set: FxHashSet<String> = query_tokens.into_iter().collect();
229
230 let pairs = &ctx.shared_corpus().docs;
231
232 if pairs.is_empty() {
233 return Vec::new();
234 }
235 let n_docs = pairs.len();
236 if n_docs > 5000 {
237 tracing::warn!(
238 "BM25Discovery: large candidate corpus ({n_docs} docs) — using inverted-index fast path"
239 );
240 }
241
242 let mut df: FxHashMap<String, usize> = FxHashMap::default();
246 let mut postings: FxHashMap<String, Vec<usize>> = FxHashMap::default();
247 let mut total_len: usize = 0;
248 for (doc_id, (_, doc)) in pairs.iter().enumerate() {
249 total_len += doc.total_len as usize;
250 for term in doc.term_counts.keys() {
251 *df.entry(term.clone()).or_insert(0) += 1;
252 if query_set.contains(term.as_str()) {
253 postings.entry(term.clone()).or_default().push(doc_id);
254 }
255 }
256 }
257 let avgdl = total_len as f64 / n_docs as f64;
258
259 let idf: FxHashMap<String, f64> = query_set
260 .iter()
261 .map(|t| {
262 let d = df.get(t).copied().unwrap_or(0) as f64;
263 let val =
264 ((n_docs as f64 - d + BM25.idf_smoothing) / (d + BM25.idf_smoothing)).ln_1p();
265 (t.clone(), val)
266 })
267 .collect();
268
269 let mut candidate_ids: FxHashSet<usize> = FxHashSet::default();
275 for term in &query_set {
276 if let Some(p) = postings.get(term) {
277 candidate_ids.extend(p);
278 }
279 }
280 if candidate_ids.is_empty() {
281 return Vec::new();
282 }
283
284 let candidate_vec: Vec<usize> = candidate_ids.into_iter().collect();
285 let scored: Vec<(usize, f64)> = candidate_vec
286 .par_iter()
287 .map(|&doc_id| {
288 let s = Self::bm25_score(&pairs[doc_id].1, &query_set, &idf, avgdl);
289 (doc_id, s)
290 })
291 .collect();
292
293 let mut ranked: Vec<(usize, f64)> = scored.into_iter().filter(|(_, s)| *s > 0.0).collect();
294 ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
295
296 ranked
297 .into_iter()
298 .take(self.top_k)
299 .map(|(i, _)| pairs[i].0.clone())
300 .collect()
301 }
302}
303
304pub struct EnsembleDiscovery {
305 strategies: Vec<Box<dyn DiscoveryStrategy>>,
306}
307
308impl EnsembleDiscovery {
309 pub fn new(strategies: Vec<Box<dyn DiscoveryStrategy>>) -> Self {
310 Self { strategies }
311 }
312
313 pub fn default_ensemble() -> Self {
314 let bm25_k =
321 crate::config::env_overrides::read_env_usize("DIFFCTX_BM25_DISCOVERY_TOP_K", 1);
322 Self {
323 strategies: vec![
324 Box::new(DefaultDiscovery),
325 Box::new(TestFileDiscovery),
326 Box::new(BM25Discovery::new(bm25_k)),
327 ],
328 }
329 }
330}
331
332impl DiscoveryStrategy for EnsembleDiscovery {
333 fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
334 self.discover_attributed(ctx).0
335 }
336
337 fn name(&self) -> &'static str {
338 "ensemble"
339 }
340
341 fn discover_attributed(
349 &self,
350 ctx: &DiscoveryContext,
351 ) -> (Vec<PathBuf>, Vec<(PathBuf, &'static str)>) {
352 let per_strategy: Vec<(&'static str, Vec<PathBuf>)> = self
353 .strategies
354 .par_iter()
355 .map(|strategy| (strategy.name(), strategy.discover(ctx)))
356 .collect();
357
358 let mut seen: FxHashSet<PathBuf> = FxHashSet::default();
359 let mut result: Vec<PathBuf> = Vec::new();
360 let mut attribution: Vec<(PathBuf, &'static str)> = Vec::new();
361 for (source, paths) in per_strategy {
362 for path in paths {
363 if seen.insert(path.clone()) {
364 attribution.push((path.clone(), source));
365 result.push(path);
366 }
367 }
368 }
369
370 (result, attribution)
371 }
372}
373
374#[cfg(test)]
375mod tests {
376 use super::*;
377
378 fn doc(text: &str) -> DocTokens {
379 let terms = extract_identifier_list(text, 1);
380 let total_len = terms.len() as u32;
381 let mut term_counts: FxHashMap<String, u32> = FxHashMap::default();
382 for t in terms {
383 *term_counts.entry(t).or_insert(0) += 1;
384 }
385 DocTokens {
386 term_counts,
387 total_len,
388 }
389 }
390
391 struct CtxBuilder {
392 changed: Vec<&'static str>,
393 candidates: Vec<&'static str>,
394 diff_text: String,
395 concepts: Vec<&'static str>,
396 docs: Vec<(&'static str, &'static str)>,
397 }
398
399 impl CtxBuilder {
400 fn new() -> Self {
401 Self {
402 changed: Vec::new(),
403 candidates: Vec::new(),
404 diff_text: String::new(),
405 concepts: Vec::new(),
406 docs: Vec::new(),
407 }
408 }
409
410 fn build(self) -> DiscoveryContext {
411 let root = PathBuf::from("/repo");
412 let corpus = TokenCorpus {
413 docs: self
414 .docs
415 .iter()
416 .map(|(p, text)| (root.join(p), doc(text)))
417 .collect(),
418 };
419 let token_corpus = OnceLock::new();
420 token_corpus
421 .set(corpus)
422 .unwrap_or_else(|_| unreachable!("fresh OnceLock"));
423 DiscoveryContext {
424 root_dir: root.clone(),
425 changed_files: self.changed.iter().map(|p| root.join(p)).collect(),
426 all_candidates: self.candidates.iter().map(|p| root.join(p)).collect(),
427 diff_text: self.diff_text,
428 expansion_concepts: self.concepts.iter().map(|s| s.to_string()).collect(),
429 file_cache: FxHashMap::default(),
430 token_corpus,
431 }
432 }
433 }
434
435 fn names(paths: &[PathBuf]) -> Vec<String> {
436 let mut v: Vec<String> = paths
437 .iter()
438 .map(|p| {
439 p.strip_prefix("/repo")
440 .unwrap_or(p)
441 .to_string_lossy()
442 .into_owned()
443 })
444 .collect();
445 v.sort();
446 v
447 }
448
449 #[test]
458 fn bare_stem_pairing_does_not_drag_in_same_named_files_repo_wide() {
459 let ctx = CtxBuilder {
460 changed: vec!["crates/a/src/mod.rs"],
461 candidates: vec![
462 "crates/a/src/mod_test.rs",
463 "crates/b/src/mod.rs",
464 "crates/c/src/mod.rs",
465 "vendor/d/mod.rs",
466 ],
467 ..CtxBuilder::new()
468 }
469 .build();
470
471 assert_eq!(
472 names(&TestFileDiscovery.discover(&ctx)),
473 vec!["crates/a/src/mod_test.rs"],
474 "unrelated same-basename files entered the universe"
475 );
476 }
477
478 #[test]
481 fn bare_stem_pairing_still_finds_a_co_located_counterpart() {
482 let ctx = CtxBuilder {
483 changed: vec!["src/parser.c"],
484 candidates: vec!["src/parser.h", "other/parser.h"],
485 ..CtxBuilder::new()
486 }
487 .build();
488
489 assert_eq!(
490 names(&TestFileDiscovery.discover(&ctx)),
491 vec!["src/parser.h"]
492 );
493 }
494
495 #[test]
496 fn test_file_discovery_pairs_every_supported_naming_convention() {
497 let ctx = CtxBuilder {
498 changed: vec!["src/auth.py", "web/handler.go", "ui/widget.ts"],
499 candidates: vec![
500 "tests/test_auth.py",
501 "web/handler_test.go",
502 "ui/widget.test.ts",
503 "ui/widget.spec.ts",
504 "ui/widget-spec.ts",
505 ],
506 ..CtxBuilder::new()
507 }
508 .build();
509
510 assert_eq!(
511 names(&TestFileDiscovery.discover(&ctx)),
512 vec![
513 "tests/test_auth.py",
514 "ui/widget-spec.ts",
515 "ui/widget.spec.ts",
516 "ui/widget.test.ts",
517 "web/handler_test.go",
518 ]
519 );
520 }
521
522 #[test]
523 fn test_file_discovery_does_not_pair_on_a_prefix_match() {
524 let ctx = CtxBuilder {
525 changed: vec!["src/authenticate.py"],
526 candidates: vec!["tests/test_auth.py", "tests/test_authenticate.py"],
527 ..CtxBuilder::new()
528 }
529 .build();
530 assert_eq!(
531 names(&TestFileDiscovery.discover(&ctx)),
532 vec!["tests/test_authenticate.py"]
533 );
534 }
535
536 #[test]
537 fn test_file_discovery_skips_changed_test_files_and_never_returns_a_changed_file() {
538 let ctx = CtxBuilder {
541 changed: vec!["tests/test_auth.py", "src/auth.py"],
542 candidates: vec!["tests/test_auth.py", "src/auth.py", "tests/test_other.py"],
543 ..CtxBuilder::new()
544 }
545 .build();
546 let found = names(&TestFileDiscovery.discover(&ctx));
547 assert!(!found.contains(&"src/auth.py".to_string()));
548 assert!(!found.contains(&"tests/test_auth.py".to_string()));
549 }
550
551 #[test]
552 fn rare_identifier_expansion_keeps_rare_terms_and_drops_common_ones() {
553 let threshold = crate::config::limits::LIMITS.rare_identifier_threshold;
554 let mut docs: Vec<(&'static str, &'static str)> = vec![
555 ("rare_a.py", "unique_marker"),
556 ("rare_b.py", "unique_marker"),
557 ];
558 let common: [&'static str; 6] = ["c0.py", "c1.py", "c2.py", "c3.py", "c4.py", "c5.py"];
560 for p in common.iter().take(threshold + 2) {
561 docs.push((p, "common_marker"));
562 }
563
564 let ctx = CtxBuilder {
565 concepts: vec!["unique_marker", "common_marker"],
566 docs,
567 ..CtxBuilder::new()
568 }
569 .build();
570
571 let found = names(&expand_by_rare_identifiers(&ctx));
572 assert!(
573 found.contains(&"rare_a.py".to_string()),
574 "rare term did not expand: {found:?}"
575 );
576 assert!(
577 found.contains(&"rare_b.py".to_string()),
578 "rare term did not expand: {found:?}"
579 );
580 assert!(
581 !found.iter().any(|f| f.starts_with("c")),
582 "a term appearing in more than {threshold} files still expanded: {found:?}"
583 );
584 }
585
586 #[test]
587 fn rare_identifier_expansion_is_empty_without_concepts() {
588 let ctx = CtxBuilder {
589 docs: vec![("a.py", "anything")],
590 ..CtxBuilder::new()
591 }
592 .build();
593 assert!(expand_by_rare_identifiers(&ctx).is_empty());
594 }
595
596 #[test]
599 fn bm25_ranks_a_rare_query_term_above_a_ubiquitous_one() {
600 let ctx = CtxBuilder {
601 diff_text: "+ use rare_needle; use ubiquitous_helper;".into(),
602 docs: vec![
603 ("has_rare.py", "rare_needle body body"),
604 ("common_1.py", "ubiquitous_helper body body"),
605 ("common_2.py", "ubiquitous_helper body body"),
606 ("common_3.py", "ubiquitous_helper body body"),
607 ("common_4.py", "ubiquitous_helper body body"),
608 ("common_5.py", "ubiquitous_helper body body"),
609 ],
610 ..CtxBuilder::new()
611 }
612 .build();
613
614 let ranked = BM25Discovery::new(6).discover(&ctx);
615 assert!(!ranked.is_empty(), "BM25 returned nothing");
616 assert_eq!(
617 names(&ranked[..1]),
618 vec!["has_rare.py"],
619 "the rare term did not win: {:?}",
620 names(&ranked)
621 );
622 }
623
624 #[test]
625 fn bm25_returns_nothing_when_no_document_contains_a_query_term() {
626 let ctx = CtxBuilder {
627 diff_text: "+ absent_symbol_xyz".into(),
628 docs: vec![("a.py", "unrelated content here")],
629 ..CtxBuilder::new()
630 }
631 .build();
632 assert!(BM25Discovery::new(5).discover(&ctx).is_empty());
633 }
634
635 #[test]
636 fn bm25_returns_nothing_on_an_empty_query_or_an_empty_corpus() {
637 let empty_query = CtxBuilder {
638 docs: vec![("a.py", "content")],
639 ..CtxBuilder::new()
640 }
641 .build();
642 assert!(BM25Discovery::new(5).discover(&empty_query).is_empty());
643
644 let empty_corpus = CtxBuilder {
645 diff_text: "+ some_symbol".into(),
646 ..CtxBuilder::new()
647 }
648 .build();
649 assert!(BM25Discovery::new(5).discover(&empty_corpus).is_empty());
650 }
651
652 #[test]
653 fn bm25_honours_top_k() {
654 let ctx = CtxBuilder {
655 diff_text: "+ shared_term".into(),
656 docs: vec![
657 ("a.py", "shared_term shared_term a"),
658 ("b.py", "shared_term b b b"),
659 ("c.py", "shared_term c c c c"),
660 ],
661 ..CtxBuilder::new()
662 }
663 .build();
664 assert_eq!(BM25Discovery::new(2).discover(&ctx).len(), 2);
665 }
666
667 #[test]
671 fn ensemble_deduplicates_across_strategies_and_preserves_first_hit_order() {
672 struct Fixed(&'static str, Vec<&'static str>);
673 impl DiscoveryStrategy for Fixed {
674 fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
675 self.1.iter().map(|p| ctx.root_dir.join(p)).collect()
676 }
677 fn name(&self) -> &'static str {
678 self.0
679 }
680 }
681
682 let ctx = CtxBuilder::new().build();
683 let ensemble = EnsembleDiscovery::new(vec![
684 Box::new(Fixed("first", vec!["a.py", "b.py"])),
685 Box::new(Fixed("second", vec!["b.py", "c.py"])),
686 Box::new(Fixed("third", vec![])),
687 ]);
688 let found = ensemble.discover(&ctx);
689 assert_eq!(
690 found
691 .iter()
692 .map(|p| p.file_name().unwrap().to_string_lossy().into_owned())
693 .collect::<Vec<_>>(),
694 vec!["a.py", "b.py", "c.py"]
695 );
696 }
697
698 #[test]
703 fn the_ensemble_records_which_strategy_first_surfaced_each_path() {
704 struct Fixed(&'static str, Vec<&'static str>);
705 impl DiscoveryStrategy for Fixed {
706 fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
707 self.1.iter().map(|p| ctx.root_dir.join(p)).collect()
708 }
709 fn name(&self) -> &'static str {
710 self.0
711 }
712 }
713
714 let ctx = CtxBuilder::new().build();
715 let ensemble = EnsembleDiscovery::new(vec![
716 Box::new(Fixed("structural", vec!["a.py", "b.py"])),
717 Box::new(Fixed("lexical", vec!["b.py", "c.py"])),
718 ]);
719 let (paths, attribution) = ensemble.discover_attributed(&ctx);
720
721 assert_eq!(paths.len(), 3, "dedup must still collapse the shared path");
722 let by_name: Vec<(String, &str)> = attribution
723 .iter()
724 .map(|(p, s)| (p.file_name().unwrap().to_string_lossy().into_owned(), *s))
725 .collect();
726 assert_eq!(
727 by_name,
728 vec![
729 ("a.py".to_string(), "structural"),
730 ("b.py".to_string(), "structural"),
731 ("c.py".to_string(), "lexical"),
732 ]
733 );
734 }
735
736 #[test]
738 fn a_lone_strategy_attributes_everything_to_itself() {
739 struct Fixed;
740 impl DiscoveryStrategy for Fixed {
741 fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
742 vec![ctx.root_dir.join("only.py")]
743 }
744 fn name(&self) -> &'static str {
745 "solo"
746 }
747 }
748 let ctx = CtxBuilder::new().build();
749 let (paths, attribution) = Fixed.discover_attributed(&ctx);
750 assert_eq!(paths.len(), 1);
751 assert_eq!(attribution[0].1, "solo");
752 }
753
754 #[test]
755 fn default_ensemble_wires_three_channels() {
756 assert_eq!(EnsembleDiscovery::default_ensemble().strategies.len(), 3);
759 }
760}