_diffctx/
discovery.rs

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
39pub struct DefaultDiscovery;
40
41impl DiscoveryStrategy for DefaultDiscovery {
42    fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
43        let changed_set: FxHashSet<&Path> = ctx.changed_files.iter().map(|p| p.as_path()).collect();
44
45        let mut discovered = crate::edges::discover_all_related_files(
46            &ctx.changed_files,
47            &ctx.all_candidates,
48            Some(ctx.root_dir.as_path()),
49            Some(&ctx.file_cache),
50        );
51        discovered.retain(|p| !changed_set.contains(p.as_path()));
52
53        let rare_files = expand_by_rare_identifiers(ctx);
54        let existing: FxHashSet<PathBuf> = discovered.iter().cloned().collect();
55        for f in rare_files {
56            if !existing.contains(&f) {
57                discovered.push(f);
58            }
59        }
60
61        discovered
62    }
63}
64
65fn expand_by_rare_identifiers(ctx: &DiscoveryContext) -> Vec<PathBuf> {
66    let rare_threshold = crate::config::limits::LIMITS.rare_identifier_threshold;
67
68    let mut ident_to_files: FxHashMap<String, Vec<PathBuf>> = FxHashMap::default();
69    for (path, doc) in &ctx.shared_corpus().docs {
70        for ident in &ctx.expansion_concepts {
71            if doc.term_counts.contains_key(ident) {
72                ident_to_files
73                    .entry(ident.clone())
74                    .or_default()
75                    .push(path.clone());
76            }
77        }
78    }
79
80    let mut result: Vec<PathBuf> = Vec::new();
81    let mut seen: FxHashSet<PathBuf> = FxHashSet::default();
82    for (_ident, files) in &ident_to_files {
83        if files.len() <= rare_threshold {
84            for f in files {
85                if seen.insert(f.clone()) {
86                    result.push(f.clone());
87                }
88            }
89        }
90    }
91    result
92}
93
94pub struct TestFileDiscovery;
95
96const TEST_PREFIXES: &[&str] = &["test_", "spec_"];
97const TEST_SUFFIXES: &[&str] = &["_test", "_spec", ".test", ".spec", "-test", "-spec"];
98
99impl DiscoveryStrategy for TestFileDiscovery {
100    fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
101        let changed_set: FxHashSet<&Path> = ctx.changed_files.iter().map(|p| p.as_path()).collect();
102        let mut target_stems: FxHashSet<String> = FxHashSet::default();
103
104        for f in &ctx.changed_files {
105            let stem = f
106                .file_stem()
107                .map(|s| s.to_string_lossy().to_lowercase())
108                .unwrap_or_default();
109            if TEST_PREFIXES.iter().any(|p| stem.starts_with(p)) {
110                continue;
111            }
112            if TEST_SUFFIXES.iter().any(|s| stem.ends_with(s)) {
113                continue;
114            }
115            target_stems.insert(stem.clone());
116            for prefix in TEST_PREFIXES {
117                target_stems.insert(format!("{}{}", prefix, stem));
118            }
119            for suffix in TEST_SUFFIXES {
120                target_stems.insert(format!("{}{}", stem, suffix));
121            }
122        }
123
124        let mut discovered: Vec<PathBuf> = Vec::new();
125        for candidate in &ctx.all_candidates {
126            if changed_set.contains(candidate.as_path()) {
127                continue;
128            }
129            let stem = candidate
130                .file_stem()
131                .map(|s| s.to_string_lossy().to_lowercase())
132                .unwrap_or_default();
133            if target_stems.contains(&stem) {
134                discovered.push(candidate.clone());
135            }
136        }
137        discovered
138    }
139}
140
141pub struct BM25Discovery {
142    pub top_k: usize,
143}
144
145impl BM25Discovery {
146    pub fn new(top_k: usize) -> Self {
147        Self { top_k }
148    }
149
150    fn bm25_score(
151        doc: &DocTokens,
152        query_set: &FxHashSet<String>,
153        idf: &FxHashMap<String, f64>,
154        avgdl: f64,
155    ) -> f64 {
156        let dl = doc.total_len as f64;
157        let mut s = 0.0;
158        for t in query_set {
159            let freq = doc.term_counts.get(t).copied().unwrap_or(0) as f64;
160            if freq == 0.0 {
161                continue;
162            }
163            let idf_val = idf.get(t).copied().unwrap_or(0.0);
164            s += idf_val * (freq * BM25.k1)
165                / (freq + BM25.k1 * (1.0 - BM25.b + BM25.b * dl / avgdl));
166        }
167        s
168    }
169}
170
171impl DiscoveryStrategy for BM25Discovery {
172    fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
173        let query_tokens = extract_identifier_list(&ctx.diff_text, BM25.min_query_token_length);
174        if query_tokens.is_empty() {
175            return Vec::new();
176        }
177        let query_set: FxHashSet<String> = query_tokens.into_iter().collect();
178
179        let pairs = &ctx.shared_corpus().docs;
180
181        if pairs.is_empty() {
182            return Vec::new();
183        }
184        let n_docs = pairs.len();
185        if n_docs > 5000 {
186            tracing::warn!(
187                "BM25Discovery: large candidate corpus ({n_docs} docs) — using inverted-index fast path"
188            );
189        }
190
191        // Single pass: compute df globally + inverted-index posting lists
192        // for query terms only (skip indexing terms not in the query — they
193        // are never needed and would balloon memory on large repos).
194        let mut df: FxHashMap<String, usize> = FxHashMap::default();
195        let mut postings: FxHashMap<String, Vec<usize>> = FxHashMap::default();
196        let mut total_len: usize = 0;
197        for (doc_id, (_, doc)) in pairs.iter().enumerate() {
198            total_len += doc.total_len as usize;
199            for term in doc.term_counts.keys() {
200                *df.entry(term.clone()).or_insert(0) += 1;
201                if query_set.contains(term.as_str()) {
202                    postings.entry(term.clone()).or_default().push(doc_id);
203                }
204            }
205        }
206        let avgdl = total_len as f64 / n_docs as f64;
207
208        let idf: FxHashMap<String, f64> = query_set
209            .iter()
210            .map(|t| {
211                let d = df.get(t).copied().unwrap_or(0) as f64;
212                let val =
213                    ((n_docs as f64 - d + BM25.idf_smoothing) / (d + BM25.idf_smoothing)).ln_1p();
214                (t.clone(), val)
215            })
216            .collect();
217
218        // Candidate doc-ids = union of posting lists for query terms. Docs
219        // not in this set contain zero query terms and would score 0 — skip
220        // them. This is the algorithmic win: scoring shrinks from O(N_docs)
221        // to O(|posting-list union|), typically ~10-100× smaller on big
222        // corpora where the query is sparse against the corpus vocabulary.
223        let mut candidate_ids: FxHashSet<usize> = FxHashSet::default();
224        for term in &query_set {
225            if let Some(p) = postings.get(term) {
226                candidate_ids.extend(p);
227            }
228        }
229        if candidate_ids.is_empty() {
230            return Vec::new();
231        }
232
233        let candidate_vec: Vec<usize> = candidate_ids.into_iter().collect();
234        let scored: Vec<(usize, f64)> = candidate_vec
235            .par_iter()
236            .map(|&doc_id| {
237                let s = Self::bm25_score(&pairs[doc_id].1, &query_set, &idf, avgdl);
238                (doc_id, s)
239            })
240            .collect();
241
242        let mut ranked: Vec<(usize, f64)> = scored.into_iter().filter(|(_, s)| *s > 0.0).collect();
243        ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
244
245        ranked
246            .into_iter()
247            .take(self.top_k)
248            .map(|(i, _)| pairs[i].0.clone())
249            .collect()
250    }
251}
252
253pub struct EnsembleDiscovery {
254    strategies: Vec<Box<dyn DiscoveryStrategy>>,
255}
256
257impl EnsembleDiscovery {
258    pub fn new(strategies: Vec<Box<dyn DiscoveryStrategy>>) -> Self {
259        Self { strategies }
260    }
261
262    pub fn default_ensemble() -> Self {
263        Self {
264            strategies: vec![
265                Box::new(DefaultDiscovery),
266                Box::new(TestFileDiscovery),
267                Box::new(BM25Discovery::new(1)),
268            ],
269        }
270    }
271}
272
273impl DiscoveryStrategy for EnsembleDiscovery {
274    fn discover(&self, ctx: &DiscoveryContext) -> Vec<PathBuf> {
275        let per_strategy: Vec<Vec<PathBuf>> = self
276            .strategies
277            .par_iter()
278            .map(|strategy| strategy.discover(ctx))
279            .collect();
280
281        let mut seen: FxHashSet<PathBuf> = FxHashSet::default();
282        let mut result: Vec<PathBuf> = Vec::new();
283        for paths in per_strategy {
284            for path in paths {
285                if seen.insert(path.clone()) {
286                    result.push(path);
287                }
288            }
289        }
290
291        result
292    }
293}
_diffctx/discovery.rs

_diffctx/
discovery.rs
