1pub mod base;
2pub mod config_edges;
3pub mod document;
4pub mod history;
5pub mod semantic;
6pub mod similarity;
7pub mod structural;
8
9use std::cmp::Reverse;
10use std::path::{Path, PathBuf};
11
12use rayon::prelude::*;
13use rustc_hash::{FxHashMap, FxHashSet};
14use tracing::debug;
15
16use crate::graph::{
17 CappedEdges, CompactEdge, EdgeCapStats, EdgeCategory, RankedCandidate, SourceTopK,
18 SuppressionFactors, cap_out_edges_per_source, dedup_compact_edges, intern_fragment_nodes,
19 push_bounded_top_k, read_max_out_edges_per_node,
20};
21use crate::types::FragmentId;
22
23pub type EdgeDict = FxHashMap<(FragmentId, FragmentId), f64>;
24pub type EdgeCategories = FxHashMap<(FragmentId, FragmentId), EdgeCategory>;
25
26use crate::types::Fragment;
27
28use self::base::EdgeBuilder;
29
30pub const NAMING_WEIGHT_FLOOR: f64 = 0.30;
34
35const EXPENSIVE_CATEGORIES: &[&str] = &["similarity", "history"];
36
37struct BuilderCategory {
38 name: &'static str,
39 builders: fn() -> Vec<Box<dyn EdgeBuilder>>,
40}
41
42fn builder_categories() -> Vec<BuilderCategory> {
43 vec![
44 BuilderCategory {
45 name: "semantic",
46 builders: || semantic::get_semantic_builders(),
47 },
48 BuilderCategory {
49 name: "structural",
50 builders: || structural::get_structural_builders(),
51 },
52 BuilderCategory {
53 name: "config",
54 builders: || config_edges::get_config_builders(),
55 },
56 BuilderCategory {
57 name: "document",
58 builders: || document::get_document_builders(),
59 },
60 BuilderCategory {
61 name: "similarity",
62 builders: || similarity::get_similarity_builders(),
63 },
64 BuilderCategory {
65 name: "history",
66 builders: || history::get_history_builders(),
67 },
68 ]
69}
70
71pub fn get_all_builders() -> Vec<Box<dyn EdgeBuilder>> {
72 let mut all = Vec::new();
73 for cat in builder_categories() {
74 all.extend((cat.builders)());
75 }
76 all
77}
78
79fn pack_pair(src: u32, dst: u32) -> u64 {
80 ((src as u64) << 32) | dst as u64
81}
82
83struct LoggedEmission {
84 src: u32,
85 dst: u32,
86 weight: f64,
87}
88
89pub fn collect_capped_edges(
113 fragments: &[Fragment],
114 repo_root: Option<&Path>,
115 skip_expensive: bool,
116 deadline: crate::deadline::Deadline,
117) -> CappedEdges {
118 let mut all_builders: Vec<(&str, Box<dyn EdgeBuilder>)> = Vec::new();
119 for cat in builder_categories() {
120 if skip_expensive && EXPENSIVE_CATEGORIES.contains(&cat.name) {
121 debug!("skipping {} edge builders (skip_expensive=true)", cat.name);
122 continue;
123 }
124 for builder in (cat.builders)() {
125 all_builders.push((cat.name, builder));
126 }
127 }
128
129 let (node_to_idx, idx_to_node) = intern_fragment_nodes(fragments);
130 let category_weights = *crate::config::category_weights::CATEGORY_WEIGHTS;
131 let builder_meta: Vec<(EdgeCategory, f64)> = all_builders
132 .iter()
133 .map(|(cat_name, builder)| {
134 let category = EdgeCategory::from_str(builder.category_label().unwrap_or(cat_name));
135 (category, category_weights.multiplier(category))
136 })
137 .collect();
138 let fallback_flags: Vec<bool> = all_builders
139 .iter()
140 .map(|(_, builder)| builder.is_fallback())
141 .collect();
142
143 let per_builder_log: Vec<Vec<LoggedEmission>> = all_builders
144 .par_iter()
145 .enumerate()
146 .map(|(builder_idx, (name, builder))| {
147 deadline.check("edge construction");
148 let _in_builder = deadline.enter();
149 let t = std::time::Instant::now();
150 let edges = builder.build(fragments, repo_root);
151 if std::env::var_os("DIFFCTX_TRACE_BUILDERS").is_some() {
152 eprintln!(
156 "builder {name}[{builder_idx}]: {:.1}s, {} edges",
157 t.elapsed().as_secs_f64(),
158 edges.len()
159 );
160 }
161 let mut log = Vec::with_capacity(edges.len());
162 for ((src, dst), weight) in edges {
163 let (Some(&s), Some(&d)) = (node_to_idx.get(&src), node_to_idx.get(&dst)) else {
164 continue;
165 };
166 log.push(LoggedEmission {
167 src: s,
168 dst: d,
169 weight,
170 });
171 }
172 log
173 })
174 .collect();
175 drop(all_builders);
176
177 let mut per_builder_log = per_builder_log;
183 if fallback_flags.iter().any(|&f| f) {
184 let mut dedicated_files: FxHashSet<&str> = FxHashSet::default();
185 for (builder_idx, log) in per_builder_log.iter().enumerate() {
186 if fallback_flags[builder_idx] || builder_meta[builder_idx].0 != EdgeCategory::Semantic
187 {
188 continue;
189 }
190 for e in log {
191 dedicated_files.insert(idx_to_node[e.src as usize].path.as_ref());
192 dedicated_files.insert(idx_to_node[e.dst as usize].path.as_ref());
193 }
194 }
195 for (builder_idx, log) in per_builder_log.iter_mut().enumerate() {
196 if !fallback_flags[builder_idx] {
197 continue;
198 }
199 log.retain(|e| {
200 !dedicated_files.contains(idx_to_node[e.src as usize].path.as_ref())
201 || !dedicated_files.contains(idx_to_node[e.dst as usize].path.as_ref())
202 });
203 }
204 }
205
206 let n_nodes = idx_to_node.len();
207 let mut in_degree = vec![0u32; n_nodes];
208 let mut out_degree = vec![0u32; n_nodes];
209 let mut category_entries: Vec<(u32, u32, EdgeCategory)> = Vec::new();
210 let mut sem_out_files: FxHashMap<u32, FxHashSet<&str>> = FxHashMap::default();
211 let mut raw_by_category: FxHashMap<EdgeCategory, u64> = FxHashMap::default();
212 let mut deduped_by_category: FxHashMap<EdgeCategory, u64> = FxHashMap::default();
213 let mut seen: FxHashSet<u64> = FxHashSet::default();
214 for (builder_idx, log) in per_builder_log.iter().enumerate() {
215 let (category, _) = builder_meta[builder_idx];
216 *raw_by_category.entry(category).or_default() += log.len() as u64;
217 for e in log {
218 if !seen.insert(pack_pair(e.src, e.dst)) {
219 continue;
220 }
221 *deduped_by_category.entry(category).or_default() += 1;
222 in_degree[e.dst as usize] += 1;
223 out_degree[e.src as usize] += 1;
224 category_entries.push((e.src, e.dst, category));
225 if category == EdgeCategory::Semantic {
226 sem_out_files
227 .entry(e.src)
228 .or_default()
229 .insert(idx_to_node[e.dst as usize].path.as_ref());
230 }
231 }
232 }
233 drop(seen);
234 let mut emissions_by_category: Vec<(EdgeCategory, u64, u64)> = raw_by_category
235 .iter()
236 .map(|(&category, &raw)| {
237 let deduped = deduped_by_category.get(&category).copied().unwrap_or(0);
238 (category, raw, deduped)
239 })
240 .collect();
241 emissions_by_category.sort_unstable_by_key(|e| e.0.as_str());
242 category_entries.sort_unstable_by_key(|e| (e.0, e.1));
243
244 let mut sem_file_deg = vec![0u32; n_nodes];
245 for (&src, files) in &sem_out_files {
246 sem_file_deg[src as usize] = files.len() as u32;
247 }
248 drop(sem_out_files);
249
250 let deduped_edge_count = category_entries.len();
251 let factors = SuppressionFactors::from_counters(in_degree, sem_file_deg);
252 let max_per_node = read_max_out_edges_per_node();
253
254 let capped_per_builder: Vec<Vec<CompactEdge>> = per_builder_log
255 .into_par_iter()
256 .enumerate()
257 .map(|(builder_idx, log)| {
258 let (builder_category, multiplier) = builder_meta[builder_idx];
259 let mut per_source: FxHashMap<u32, SourceTopK> = FxHashMap::default();
260 for e in log {
261 let category = category_entries
262 .binary_search_by_key(&(e.src, e.dst), |c| (c.0, c.1))
263 .map(|k| category_entries[k].2)
264 .unwrap_or(builder_category);
265 let damped = factors.damp(e.weight * multiplier, category, e.src, e.dst);
266 let naming = (category == EdgeCategory::Semantic
273 || category == EdgeCategory::Config)
274 && e.weight > NAMING_WEIGHT_FLOOR;
275 push_bounded_top_k(
276 per_source.entry(e.src).or_default(),
277 RankedCandidate {
278 weight: damped,
279 dst: e.dst,
280 category,
281 naming,
282 },
283 max_per_node,
284 );
285 }
286 let mut survivors =
287 Vec::with_capacity(per_source.values().map(|h| h.len()).sum::<usize>());
288 for (src, heap) in per_source {
289 for Reverse(c) in heap {
290 survivors.push(CompactEdge {
291 src,
292 dst: c.dst,
293 weight: c.weight,
294 category: c.category,
295 naming: c.naming,
296 });
297 }
298 }
299 survivors
300 })
301 .collect();
302
303 let total: usize = capped_per_builder.iter().map(|v| v.len()).sum();
304 let mut edges: Vec<CompactEdge> = Vec::with_capacity(total);
305 for v in capped_per_builder {
306 edges.extend(v);
307 }
308 dedup_compact_edges(&mut edges);
309 cap_out_edges_per_source(&mut edges, max_per_node);
310
311 let nodes_capped = out_degree
312 .iter()
313 .filter(|&&d| d as usize > max_per_node)
314 .count();
315 let cap_stats = EdgeCapStats {
316 edges_before_cap: deduped_edge_count,
317 edges_after_cap: edges.len(),
318 edges_dropped_by_cap: deduped_edge_count - edges.len(),
319 nodes_capped,
320 max_out_edges_per_node: max_per_node,
321 emissions_by_category,
322 };
323
324 CappedEdges {
325 node_to_idx,
326 idx_to_node,
327 edges,
328 cap_stats,
329 }
330}
331
332pub fn discover_all_related_files(
333 changed_files: &[PathBuf],
334 all_candidates: &[PathBuf],
335 repo_root: Option<&Path>,
336 file_cache: Option<&FxHashMap<PathBuf, String>>,
337) -> Vec<PathBuf> {
338 let mut discovered: FxHashMap<PathBuf, ()> = FxHashMap::default();
339 for builder in get_all_builders() {
340 for f in
341 builder.discover_related_files(changed_files, all_candidates, repo_root, file_cache)
342 {
343 discovered.entry(f).or_insert(());
344 }
345 }
346 let mut result: Vec<PathBuf> = discovered.into_keys().collect();
347 result.sort();
348 result
349}
350
351pub fn naming_reachable_files(
356 capped: &CappedEdges,
357 core_ids: &rustc_hash::FxHashSet<crate::types::FragmentId>,
358 max_depth: usize,
359) -> rustc_hash::FxHashSet<std::sync::Arc<str>> {
360 let n = capped.idx_to_node.len();
361 let mut adj: Vec<Vec<u32>> = vec![Vec::new(); n];
368 for e in &capped.edges {
369 if e.naming {
370 adj[e.src as usize].push(e.dst);
371 adj[e.dst as usize].push(e.src);
372 }
373 }
374 let mut seen = vec![false; n];
375 let mut frontier: Vec<u32> = Vec::new();
376 for (i, id) in capped.idx_to_node.iter().enumerate() {
377 if core_ids.contains(id) {
378 seen[i] = true;
379 frontier.push(i as u32);
380 }
381 }
382 let mut files: rustc_hash::FxHashSet<std::sync::Arc<str>> = frontier
383 .iter()
384 .map(|&i| capped.idx_to_node[i as usize].path.clone())
385 .collect();
386 for _ in 0..max_depth {
387 let mut next = Vec::new();
388 for &u in &frontier {
389 for &v in &adj[u as usize] {
390 if !seen[v as usize] {
391 seen[v as usize] = true;
392 files.insert(capped.idx_to_node[v as usize].path.clone());
393 next.push(v);
394 }
395 }
396 }
397 if next.is_empty() {
398 break;
399 }
400 frontier = next;
401 }
402 files
403}
404
405#[cfg(test)]
406mod fallback_gate_tests {
407 use super::*;
408 use rustc_hash::FxHashSet as Set;
409 use std::sync::Arc;
410
411 fn frag(path: &str, content: &str, idents: &[&str]) -> Fragment {
412 Fragment {
413 id: crate::types::FragmentId::new(Arc::from(path), 1, 10),
414 kind: crate::types::FragmentKind::Function,
415 content: Arc::from(content),
416 identifiers: idents.iter().map(|s| s.to_string()).collect::<Set<_>>(),
417 token_count: 10,
418 symbol_name: None,
419 }
420 }
421
422 #[test]
423 fn tags_edges_survive_only_where_dedicated_builders_came_back_empty() {
424 let fragments = vec![
428 frag(
429 "proj/a.c",
430 "#include \"bdep.h\"\nint zzcommonzz;\n",
431 &["zzcommonzz"],
432 ),
433 frag("proj/bdep.h", "int bdecl(void);\n", &["bdecl"]),
434 frag(
435 "proj/c.c",
436 "#include \"ddep.h\"\nint zzcommonzz;\n",
437 &["zzcommonzz"],
438 ),
439 frag("proj/ddep.h", "int ddecl(void);\n", &["ddecl"]),
440 frag("proj/u1.xyz", "zzcommonzz here\n", &["zzcommonzz"]),
441 frag("proj/u2.xyz", "zzcommonzz there\n", &["zzcommonzz"]),
442 ];
443 let capped =
444 collect_capped_edges(&fragments, None, false, crate::deadline::Deadline::none());
445 let node_path = |idx: u32| capped.idx_to_node[idx as usize].path.clone();
446 let has = |a: &str, b: &str| {
449 capped.edges.iter().any(|e| {
450 if e.category != EdgeCategory::Semantic {
451 return false;
452 }
453 let s = node_path(e.src);
454 let d = node_path(e.dst);
455 (s.ends_with(a) && d.ends_with(b)) || (s.ends_with(b) && d.ends_with(a))
456 })
457 };
458 assert!(
459 has("u1.xyz", "u2.xyz"),
460 "fallback must still connect files no dedicated builder covers"
461 );
462 assert!(
463 !has("a.c", "c.c"),
464 "a tags-only link between two dedicated-covered files is the measured noise class (#131)"
465 );
466 }
467}