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
314impl DiscoveryStrategy for EnsembleDiscovery {
315 fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
316 self.discover_attributed(ctx).0
317 }
318
319 fn name(&self) -> &'static str {
320 "ensemble"
321 }
322
323 fn discover_attributed(
331 &self,
332 ctx: &DiscoveryContext,
333 ) -> (Vec<PathBuf>, Vec<(PathBuf, &'static str)>) {
334 let per_strategy: Vec<(&'static str, Vec<PathBuf>)> = self
335 .strategies
336 .par_iter()
337 .map(|strategy| (strategy.name(), strategy.discover(ctx)))
338 .collect();
339
340 let mut seen: FxHashSet<PathBuf> = FxHashSet::default();
341 let mut result: Vec<PathBuf> = Vec::new();
342 let mut attribution: Vec<(PathBuf, &'static str)> = Vec::new();
343 for (source, paths) in per_strategy {
344 for path in paths {
345 if seen.insert(path.clone()) {
346 attribution.push((path.clone(), source));
347 result.push(path);
348 }
349 }
350 }
351
352 (result, attribution)
353 }
354}
355
356#[cfg(test)]
357mod tests {
358 use super::*;
359
360 fn doc(text: &str) -> DocTokens {
361 let terms = extract_identifier_list(text, 1);
362 let total_len = terms.len() as u32;
363 let mut term_counts: FxHashMap<String, u32> = FxHashMap::default();
364 for t in terms {
365 *term_counts.entry(t).or_insert(0) += 1;
366 }
367 DocTokens {
368 term_counts,
369 total_len,
370 }
371 }
372
373 struct CtxBuilder {
374 changed: Vec<&'static str>,
375 candidates: Vec<&'static str>,
376 diff_text: String,
377 concepts: Vec<&'static str>,
378 docs: Vec<(&'static str, &'static str)>,
379 }
380
381 impl CtxBuilder {
382 fn new() -> Self {
383 Self {
384 changed: Vec::new(),
385 candidates: Vec::new(),
386 diff_text: String::new(),
387 concepts: Vec::new(),
388 docs: Vec::new(),
389 }
390 }
391
392 fn build(self) -> DiscoveryContext {
393 let root = PathBuf::from("/repo");
394 let corpus = TokenCorpus {
395 docs: self
396 .docs
397 .iter()
398 .map(|(p, text)| (root.join(p), doc(text)))
399 .collect(),
400 };
401 let token_corpus = OnceLock::new();
402 token_corpus
403 .set(corpus)
404 .unwrap_or_else(|_| unreachable!("fresh OnceLock"));
405 DiscoveryContext {
406 root_dir: root.clone(),
407 changed_files: self.changed.iter().map(|p| root.join(p)).collect(),
408 all_candidates: self.candidates.iter().map(|p| root.join(p)).collect(),
409 diff_text: self.diff_text,
410 expansion_concepts: self.concepts.iter().map(|s| s.to_string()).collect(),
411 file_cache: FxHashMap::default(),
412 token_corpus,
413 }
414 }
415 }
416
417 fn names(paths: &[PathBuf]) -> Vec<String> {
418 let mut v: Vec<String> = paths
419 .iter()
420 .map(|p| {
421 p.strip_prefix("/repo")
422 .unwrap_or(p)
423 .to_string_lossy()
424 .into_owned()
425 })
426 .collect();
427 v.sort();
428 v
429 }
430
431 #[test]
440 fn bare_stem_pairing_does_not_drag_in_same_named_files_repo_wide() {
441 let ctx = CtxBuilder {
442 changed: vec!["crates/a/src/mod.rs"],
443 candidates: vec![
444 "crates/a/src/mod_test.rs",
445 "crates/b/src/mod.rs",
446 "crates/c/src/mod.rs",
447 "vendor/d/mod.rs",
448 ],
449 ..CtxBuilder::new()
450 }
451 .build();
452
453 assert_eq!(
454 names(&TestFileDiscovery.discover(&ctx)),
455 vec!["crates/a/src/mod_test.rs"],
456 "unrelated same-basename files entered the universe"
457 );
458 }
459
460 #[test]
463 fn bare_stem_pairing_still_finds_a_co_located_counterpart() {
464 let ctx = CtxBuilder {
465 changed: vec!["src/parser.c"],
466 candidates: vec!["src/parser.h", "other/parser.h"],
467 ..CtxBuilder::new()
468 }
469 .build();
470
471 assert_eq!(
472 names(&TestFileDiscovery.discover(&ctx)),
473 vec!["src/parser.h"]
474 );
475 }
476
477 #[test]
478 fn test_file_discovery_pairs_every_supported_naming_convention() {
479 let ctx = CtxBuilder {
480 changed: vec!["src/auth.py", "web/handler.go", "ui/widget.ts"],
481 candidates: vec![
482 "tests/test_auth.py",
483 "web/handler_test.go",
484 "ui/widget.test.ts",
485 "ui/widget.spec.ts",
486 "ui/widget-spec.ts",
487 ],
488 ..CtxBuilder::new()
489 }
490 .build();
491
492 assert_eq!(
493 names(&TestFileDiscovery.discover(&ctx)),
494 vec![
495 "tests/test_auth.py",
496 "ui/widget-spec.ts",
497 "ui/widget.spec.ts",
498 "ui/widget.test.ts",
499 "web/handler_test.go",
500 ]
501 );
502 }
503
504 #[test]
505 fn test_file_discovery_does_not_pair_on_a_prefix_match() {
506 let ctx = CtxBuilder {
507 changed: vec!["src/authenticate.py"],
508 candidates: vec!["tests/test_auth.py", "tests/test_authenticate.py"],
509 ..CtxBuilder::new()
510 }
511 .build();
512 assert_eq!(
513 names(&TestFileDiscovery.discover(&ctx)),
514 vec!["tests/test_authenticate.py"]
515 );
516 }
517
518 #[test]
519 fn test_file_discovery_skips_changed_test_files_and_never_returns_a_changed_file() {
520 let ctx = CtxBuilder {
523 changed: vec!["tests/test_auth.py", "src/auth.py"],
524 candidates: vec!["tests/test_auth.py", "src/auth.py", "tests/test_other.py"],
525 ..CtxBuilder::new()
526 }
527 .build();
528 let found = names(&TestFileDiscovery.discover(&ctx));
529 assert!(!found.contains(&"src/auth.py".to_string()));
530 assert!(!found.contains(&"tests/test_auth.py".to_string()));
531 }
532
533 #[test]
534 fn rare_identifier_expansion_keeps_rare_terms_and_drops_common_ones() {
535 let threshold = crate::config::limits::LIMITS.rare_identifier_threshold;
536 let mut docs: Vec<(&'static str, &'static str)> = vec![
537 ("rare_a.py", "unique_marker"),
538 ("rare_b.py", "unique_marker"),
539 ];
540 let common: [&'static str; 6] = ["c0.py", "c1.py", "c2.py", "c3.py", "c4.py", "c5.py"];
542 for p in common.iter().take(threshold + 2) {
543 docs.push((p, "common_marker"));
544 }
545
546 let ctx = CtxBuilder {
547 concepts: vec!["unique_marker", "common_marker"],
548 docs,
549 ..CtxBuilder::new()
550 }
551 .build();
552
553 let found = names(&expand_by_rare_identifiers(&ctx));
554 assert!(
555 found.contains(&"rare_a.py".to_string()),
556 "rare term did not expand: {found:?}"
557 );
558 assert!(
559 found.contains(&"rare_b.py".to_string()),
560 "rare term did not expand: {found:?}"
561 );
562 assert!(
563 !found.iter().any(|f| f.starts_with("c")),
564 "a term appearing in more than {threshold} files still expanded: {found:?}"
565 );
566 }
567
568 #[test]
569 fn rare_identifier_expansion_is_empty_without_concepts() {
570 let ctx = CtxBuilder {
571 docs: vec![("a.py", "anything")],
572 ..CtxBuilder::new()
573 }
574 .build();
575 assert!(expand_by_rare_identifiers(&ctx).is_empty());
576 }
577
578 #[test]
581 fn bm25_ranks_a_rare_query_term_above_a_ubiquitous_one() {
582 let ctx = CtxBuilder {
583 diff_text: "+ use rare_needle; use ubiquitous_helper;".into(),
584 docs: vec![
585 ("has_rare.py", "rare_needle body body"),
586 ("common_1.py", "ubiquitous_helper body body"),
587 ("common_2.py", "ubiquitous_helper body body"),
588 ("common_3.py", "ubiquitous_helper body body"),
589 ("common_4.py", "ubiquitous_helper body body"),
590 ("common_5.py", "ubiquitous_helper body body"),
591 ],
592 ..CtxBuilder::new()
593 }
594 .build();
595
596 let ranked = BM25Discovery::new(6).discover(&ctx);
597 assert!(!ranked.is_empty(), "BM25 returned nothing");
598 assert_eq!(
599 names(&ranked[..1]),
600 vec!["has_rare.py"],
601 "the rare term did not win: {:?}",
602 names(&ranked)
603 );
604 }
605
606 #[test]
607 fn bm25_returns_nothing_when_no_document_contains_a_query_term() {
608 let ctx = CtxBuilder {
609 diff_text: "+ absent_symbol_xyz".into(),
610 docs: vec![("a.py", "unrelated content here")],
611 ..CtxBuilder::new()
612 }
613 .build();
614 assert!(BM25Discovery::new(5).discover(&ctx).is_empty());
615 }
616
617 #[test]
618 fn bm25_returns_nothing_on_an_empty_query_or_an_empty_corpus() {
619 let empty_query = CtxBuilder {
620 docs: vec![("a.py", "content")],
621 ..CtxBuilder::new()
622 }
623 .build();
624 assert!(BM25Discovery::new(5).discover(&empty_query).is_empty());
625
626 let empty_corpus = CtxBuilder {
627 diff_text: "+ some_symbol".into(),
628 ..CtxBuilder::new()
629 }
630 .build();
631 assert!(BM25Discovery::new(5).discover(&empty_corpus).is_empty());
632 }
633
634 #[test]
635 fn bm25_honours_top_k() {
636 let ctx = CtxBuilder {
637 diff_text: "+ shared_term".into(),
638 docs: vec![
639 ("a.py", "shared_term shared_term a"),
640 ("b.py", "shared_term b b b"),
641 ("c.py", "shared_term c c c c"),
642 ],
643 ..CtxBuilder::new()
644 }
645 .build();
646 assert_eq!(BM25Discovery::new(2).discover(&ctx).len(), 2);
647 }
648
649 #[test]
653 fn ensemble_deduplicates_across_strategies_and_preserves_first_hit_order() {
654 struct Fixed(&'static str, Vec<&'static str>);
655 impl DiscoveryStrategy for Fixed {
656 fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
657 self.1.iter().map(|p| ctx.root_dir.join(p)).collect()
658 }
659 fn name(&self) -> &'static str {
660 self.0
661 }
662 }
663
664 let ctx = CtxBuilder::new().build();
665 let ensemble = EnsembleDiscovery::new(vec![
666 Box::new(Fixed("first", vec!["a.py", "b.py"])),
667 Box::new(Fixed("second", vec!["b.py", "c.py"])),
668 Box::new(Fixed("third", vec![])),
669 ]);
670 let found = ensemble.discover(&ctx);
671 assert_eq!(
672 found
673 .iter()
674 .map(|p| p.file_name().unwrap().to_string_lossy().into_owned())
675 .collect::<Vec<_>>(),
676 vec!["a.py", "b.py", "c.py"]
677 );
678 }
679
680 #[test]
685 fn the_ensemble_records_which_strategy_first_surfaced_each_path() {
686 struct Fixed(&'static str, Vec<&'static str>);
687 impl DiscoveryStrategy for Fixed {
688 fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
689 self.1.iter().map(|p| ctx.root_dir.join(p)).collect()
690 }
691 fn name(&self) -> &'static str {
692 self.0
693 }
694 }
695
696 let ctx = CtxBuilder::new().build();
697 let ensemble = EnsembleDiscovery::new(vec![
698 Box::new(Fixed("structural", vec!["a.py", "b.py"])),
699 Box::new(Fixed("lexical", vec!["b.py", "c.py"])),
700 ]);
701 let (paths, attribution) = ensemble.discover_attributed(&ctx);
702
703 assert_eq!(paths.len(), 3, "dedup must still collapse the shared path");
704 let by_name: Vec<(String, &str)> = attribution
705 .iter()
706 .map(|(p, s)| (p.file_name().unwrap().to_string_lossy().into_owned(), *s))
707 .collect();
708 assert_eq!(
709 by_name,
710 vec![
711 ("a.py".to_string(), "structural"),
712 ("b.py".to_string(), "structural"),
713 ("c.py".to_string(), "lexical"),
714 ]
715 );
716 }
717
718 #[test]
720 fn a_lone_strategy_attributes_everything_to_itself() {
721 struct Fixed;
722 impl DiscoveryStrategy for Fixed {
723 fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
724 vec![ctx.root_dir.join("only.py")]
725 }
726 fn name(&self) -> &'static str {
727 "solo"
728 }
729 }
730 let ctx = CtxBuilder::new().build();
731 let (paths, attribution) = Fixed.discover_attributed(&ctx);
732 assert_eq!(paths.len(), 1);
733 assert_eq!(attribution[0].1, "solo");
734 }
735}