_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        for test_frag in &test_frags {
129            let test_stem = Path::new(test_frag.path())
130                .file_stem()
131                .map(|s| s.to_string_lossy().to_string())
132                .unwrap_or_default();
133            let target_name = match extract_target_name_from_test(&test_stem) {
134                Some(name) => name,
135                None => continue,
136            };
137
138            let test_imports = import_cache
139                .get(test_frag.path())
140                .cloned()
141                .unwrap_or_default();
142
143            // `config_test` names `config` — but which one? envoy holds 520
144            // files with that stem, and pairing every test fragment with every
145            // fragment of every one of them put 180M edges out of this builder
146            // alone (half of dcbench hung in the graph build). Co-location
147            // resolves it the way the naming convention means it: the test's
148            // own directory first, then a global fallback only while the stem
149            // stays rare enough to identify a file.
150            let candidates = by_base.get(&target_name).map(Vec::as_slice).unwrap_or(&[]);
151            let test_dir = Path::new(test_frag.path()).parent();
152            let same_dir: Vec<&&Fragment> = candidates
153                .iter()
154                .filter(|sf| Path::new(sf.path()).parent() == test_dir)
155                .collect();
156            let chosen: Vec<&&Fragment> = if !same_dir.is_empty() {
157                same_dir
158            } else {
159                let distinct_files: FxHashSet<&str> =
160                    candidates.iter().map(|sf| sf.path()).collect();
161                if distinct_files.len() > MAX_GLOBAL_TARGET_FILES {
162                    continue;
163                }
164                candidates.iter().collect()
165            };
166
167            let Some(test_rep) = reps.get(test_frag.path()) else {
168                continue;
169            };
170            for src_frag in chosen {
171                if !linked_file_pairs.insert((test_frag.path(), src_frag.path())) {
172                    continue;
173                }
174                let Some(src_rep) = reps.get(src_frag.path()) else {
175                    continue;
176                };
177                let src_module = module_cache
178                    .get(src_frag.path())
179                    .map(|s| s.as_str())
180                    .unwrap_or("");
181                let weight = if has_direct_import(&test_imports, src_module) {
182                    weight_direct
183                } else {
184                    weight_naming
185                };
186
187                add_edge_unidirectional(&mut edges, test_rep, src_rep, weight);
188                add_edge_unidirectional(&mut edges, src_rep, test_rep, test_reverse_weight);
189            }
190        }
191
192        edges
193    }
194
195    fn discover_related_files(
196        &self,
197        changed: &[PathBuf],
198        candidates: &[PathBuf],
199        _repo_root: Option<&Path>,
200        _file_cache: Option<&FxHashMap<PathBuf, String>>,
201    ) -> Vec<PathBuf> {
202        let changed_set: FxHashSet<PathBuf> = changed.iter().cloned().collect();
203        let mut candidate_by_stem: FxHashMap<String, Vec<PathBuf>> = FxHashMap::default();
204        for c in candidates {
205            if !changed_set.contains(c) {
206                let stem = c
207                    .file_stem()
208                    .map(|s| s.to_string_lossy().to_lowercase())
209                    .unwrap_or_default();
210                candidate_by_stem.entry(stem).or_default().push(c.clone());
211            }
212        }
213
214        let mut discovered: Vec<PathBuf> = Vec::new();
215
216        for changed_file in changed {
217            let ext = base::file_ext(changed_file);
218            let stem = changed_file
219                .file_stem()
220                .map(|s| s.to_string_lossy().to_string())
221                .unwrap_or_default();
222
223            if is_test_file(changed_file) {
224                if let Some(target) = extract_target_name_from_test(&stem) {
225                    for c in candidate_by_stem.get(&target).unwrap_or(&vec![]) {
226                        if base::file_ext(c) == ext {
227                            discovered.push(c.clone());
228                        }
229                    }
230                }
231            } else {
232                let stem_lower = stem.to_lowercase();
233                for test_stem in [
234                    format!("test_{}", stem_lower),
235                    format!("{}_test", stem_lower),
236                ] {
237                    for c in candidate_by_stem.get(&test_stem).unwrap_or(&vec![]) {
238                        if base::file_ext(c) == ext && is_test_file(c) {
239                            discovered.push(c.clone());
240                        }
241                    }
242                }
243
244                if matches!(
245                    ext.as_str(),
246                    ".js" | ".ts" | ".jsx" | ".tsx" | ".mjs" | ".cjs"
247                ) {
248                    let stem_test = format!("{stem_lower}.test");
249                    let stem_spec = format!("{stem_lower}.spec");
250                    for c in candidate_by_stem.get(&stem_test).unwrap_or(&vec![]) {
251                        if base::file_ext(c) == ext && is_test_file(c) {
252                            discovered.push(c.clone());
253                        }
254                    }
255                    for c in candidate_by_stem.get(&stem_spec).unwrap_or(&vec![]) {
256                        if base::file_ext(c) == ext && is_test_file(c) {
257                            discovered.push(c.clone());
258                        }
259                    }
260                }
261            }
262        }
263
264        discovered
265    }
266
267    fn category_label(&self) -> Option<&str> {
268        Some("test_edge")
269    }
270}
_diffctx/edges/structural/testing.rs

_diffctx/edges/structural/
testing.rs
