Skip to main content

_diffctx/edges/structural/
testing.rs

1use std::path::{Path, PathBuf};
2
3use once_cell::sync::Lazy;
4use regex::Regex;
5use rustc_hash::{FxHashMap, FxHashSet};
6
7use crate::config::weights::EDGE_WEIGHTS;
8use crate::types::Fragment;
9
10use super::super::EdgeDict;
11use super::super::base::{self, EdgeBuilder, add_edge_unidirectional, path_to_module};
12
13static IMPORT_RE: Lazy<Regex> =
14    Lazy::new(|| Regex::new(r"(?m)^\s*(?:from\s+([\w.]+)\s+import|import\s+([\w.]+))").unwrap());
15
16fn is_test_file(path: &Path) -> bool {
17    crate::testfiles::is_test_path(path)
18}
19
20fn extract_imports(content: &str) -> FxHashSet<String> {
21    let mut imports = FxHashSet::default();
22    for cap in IMPORT_RE.captures_iter(content) {
23        if let Some(m) = cap.get(1) {
24            imports.insert(m.as_str().to_string());
25        }
26        if let Some(m) = cap.get(2) {
27            imports.insert(m.as_str().to_string());
28        }
29    }
30    imports
31}
32
33fn has_direct_import(test_imports: &FxHashSet<String>, src_module: &str) -> bool {
34    if src_module.is_empty() {
35        return false;
36    }
37    let suffix = format!(".{}", src_module);
38    test_imports
39        .iter()
40        .any(|imp| imp == src_module || imp.ends_with(&suffix))
41}
42
43fn extract_target_name_from_test(test_name: &str) -> Option<String> {
44    let lower = test_name.to_lowercase();
45    if lower.starts_with("test_") {
46        return Some(lower[5..].to_string());
47    }
48    if lower.ends_with("_test") {
49        return Some(lower[..lower.len() - 5].to_string());
50    }
51    if lower.contains(".test") {
52        return Some(lower.split(".test").next()?.to_string());
53    }
54    if lower.contains(".spec") {
55        return Some(lower.split(".spec").next()?.to_string());
56    }
57    if test_name.starts_with("Test")
58        && test_name.len() > 4
59        && test_name.as_bytes()[4].is_ascii_uppercase()
60    {
61        return Some(test_name[4..].to_lowercase());
62    }
63    if test_name.ends_with("Tests") && test_name.len() > 5 {
64        return Some(test_name[..test_name.len() - 5].to_lowercase());
65    }
66    if test_name.ends_with("Test") && test_name.len() > 4 {
67        return Some(test_name[..test_name.len() - 4].to_lowercase());
68    }
69    None
70}
71
72pub struct TestEdgeBuilder;
73
74/// Same ambiguity bar as the c-family and path-reference caps: a stem carried
75/// by more files than this cannot name a target.
76const MAX_GLOBAL_TARGET_FILES: usize = 8;
77
78impl EdgeBuilder for TestEdgeBuilder {
79    fn build(&self, fragments: &[Fragment], repo_root: Option<&Path>) -> EdgeDict {
80        let weight_direct = EDGE_WEIGHTS["test_direct"].forward;
81        let weight_naming = EDGE_WEIGHTS["test_naming"].forward;
82        let test_reverse_weight = EDGE_WEIGHTS["test_reverse"].forward;
83
84        let mut test_frags: Vec<&Fragment> = Vec::new();
85        let mut by_base: FxHashMap<String, Vec<&Fragment>> = FxHashMap::default();
86
87        for f in fragments {
88            let path = Path::new(f.path());
89            if is_test_file(path) {
90                test_frags.push(f);
91            } else {
92                let stem = path
93                    .file_stem()
94                    .map(|s| s.to_string_lossy().to_lowercase())
95                    .unwrap_or_default();
96                by_base.entry(stem).or_default().push(f);
97            }
98        }
99
100        let mut module_cache: FxHashMap<String, String> = FxHashMap::default();
101        for src_list in by_base.values() {
102            for sf in src_list {
103                let path_str = sf.path().to_string();
104                module_cache
105                    .entry(path_str)
106                    .or_insert_with_key(|_| path_to_module(Path::new(sf.path()), repo_root));
107            }
108        }
109
110        let mut import_cache: FxHashMap<String, FxHashSet<String>> = FxHashMap::default();
111        for tf in &test_frags {
112            let path_str = tf.path().to_string();
113            import_cache
114                .entry(path_str)
115                .or_insert_with(|| extract_imports(&tf.content));
116        }
117
118        let reps = base::file_representatives(fragments);
119
120        let mut edges: EdgeDict = FxHashMap::default();
121
122        // The naming convention names a FILE, so the relation is carried by
123        // the two files' representatives (base::file_representatives) and the
124        // containment star spreads it within each file. One pair per
125        // (test file, source file), not per fragment pair.
126        let mut linked_file_pairs: FxHashSet<(&str, &str)> = FxHashSet::default();
127
128        // One pass per test FILE, not per test fragment: everything below is
129        // file-level (stem, imports, representatives), and a home-assistant
130        // test suite holds thousands of fragments per stem — re-resolving the
131        // same candidate set per fragment burned 37s to emit 2.5k edges
132        // (#196).
133        let mut seen_test_files: FxHashSet<&str> = FxHashSet::default();
134        for test_frag in &test_frags {
135            if !seen_test_files.insert(test_frag.path()) {
136                continue;
137            }
138            let test_stem = Path::new(test_frag.path())
139                .file_stem()
140                .map(|s| s.to_string_lossy().to_string())
141                .unwrap_or_default();
142            let target_name = match extract_target_name_from_test(&test_stem) {
143                Some(name) => name,
144                None => continue,
145            };
146
147            let test_imports = import_cache
148                .get(test_frag.path())
149                .cloned()
150                .unwrap_or_default();
151
152            // `config_test` names `config` — but which one? envoy holds 520
153            // files with that stem, and pairing every test fragment with every
154            // fragment of every one of them put 180M edges out of this builder
155            // alone (half of dcbench hung in the graph build). Co-location
156            // resolves it the way the naming convention means it: the test's
157            // own directory first, then a global fallback only while the stem
158            // stays rare enough to identify a file.
159            let candidates = by_base.get(&target_name).map(Vec::as_slice).unwrap_or(&[]);
160            let test_dir = Path::new(test_frag.path()).parent();
161            let same_dir: Vec<&&Fragment> = candidates
162                .iter()
163                .filter(|sf| Path::new(sf.path()).parent() == test_dir)
164                .collect();
165            let chosen: Vec<&&Fragment> = if !same_dir.is_empty() {
166                same_dir
167            } else {
168                let distinct_files: FxHashSet<&str> =
169                    candidates.iter().map(|sf| sf.path()).collect();
170                if distinct_files.len() > MAX_GLOBAL_TARGET_FILES {
171                    continue;
172                }
173                candidates.iter().collect()
174            };
175
176            let Some(test_rep) = reps.get(test_frag.path()) else {
177                continue;
178            };
179            for src_frag in chosen {
180                if !linked_file_pairs.insert((test_frag.path(), src_frag.path())) {
181                    continue;
182                }
183                let Some(src_rep) = reps.get(src_frag.path()) else {
184                    continue;
185                };
186                let src_module = module_cache
187                    .get(src_frag.path())
188                    .map(|s| s.as_str())
189                    .unwrap_or("");
190                let weight = if has_direct_import(&test_imports, src_module) {
191                    weight_direct
192                } else {
193                    weight_naming
194                };
195
196                add_edge_unidirectional(&mut edges, test_rep, src_rep, weight);
197                add_edge_unidirectional(&mut edges, src_rep, test_rep, test_reverse_weight);
198            }
199        }
200
201        edges
202    }
203
204    fn discover_related_files(
205        &self,
206        changed: &[PathBuf],
207        candidates: &[PathBuf],
208        _repo_root: Option<&Path>,
209        _file_cache: Option<&FxHashMap<PathBuf, String>>,
210    ) -> Vec<PathBuf> {
211        let changed_set: FxHashSet<PathBuf> = changed.iter().cloned().collect();
212        let mut candidate_by_stem: FxHashMap<String, Vec<PathBuf>> = FxHashMap::default();
213        for c in candidates {
214            if !changed_set.contains(c) {
215                let stem = c
216                    .file_stem()
217                    .map(|s| s.to_string_lossy().to_lowercase())
218                    .unwrap_or_default();
219                candidate_by_stem.entry(stem).or_default().push(c.clone());
220            }
221        }
222
223        let mut discovered: Vec<PathBuf> = Vec::new();
224
225        for changed_file in changed {
226            let ext = base::file_ext(changed_file);
227            let stem = changed_file
228                .file_stem()
229                .map(|s| s.to_string_lossy().to_string())
230                .unwrap_or_default();
231
232            if is_test_file(changed_file) {
233                if let Some(target) = extract_target_name_from_test(&stem) {
234                    for c in candidate_by_stem.get(&target).unwrap_or(&vec![]) {
235                        if base::file_ext(c) == ext {
236                            discovered.push(c.clone());
237                        }
238                    }
239                }
240            } else {
241                let stem_lower = stem.to_lowercase();
242                for test_stem in [
243                    format!("test_{}", stem_lower),
244                    format!("{}_test", stem_lower),
245                ] {
246                    for c in candidate_by_stem.get(&test_stem).unwrap_or(&vec![]) {
247                        if base::file_ext(c) == ext && is_test_file(c) {
248                            discovered.push(c.clone());
249                        }
250                    }
251                }
252
253                if matches!(
254                    ext.as_str(),
255                    ".js" | ".ts" | ".jsx" | ".tsx" | ".mjs" | ".cjs"
256                ) {
257                    let stem_test = format!("{stem_lower}.test");
258                    let stem_spec = format!("{stem_lower}.spec");
259                    for c in candidate_by_stem.get(&stem_test).unwrap_or(&vec![]) {
260                        if base::file_ext(c) == ext && is_test_file(c) {
261                            discovered.push(c.clone());
262                        }
263                    }
264                    for c in candidate_by_stem.get(&stem_spec).unwrap_or(&vec![]) {
265                        if base::file_ext(c) == ext && is_test_file(c) {
266                            discovered.push(c.clone());
267                        }
268                    }
269                }
270            }
271        }
272
273        discovered
274    }
275
276    fn category_label(&self) -> Option<&str> {
277        Some("test_edge")
278    }
279}