pub mod base;
pub mod config_edges;
pub mod document;
pub mod history;
pub mod semantic;
pub mod similarity;
pub mod structural;
use std::cmp::Reverse;
use std::path::{Path, PathBuf};
use rayon::prelude::*;
use rustc_hash::{FxHashMap, FxHashSet};
use tracing::debug;
use crate::graph::{
CappedEdges, CompactEdge, EdgeCapStats, EdgeCategory, RankedCandidate, SourceTopK,
SuppressionFactors, cap_out_edges_per_source, dedup_compact_edges, intern_fragment_nodes,
push_bounded_top_k, read_max_out_edges_per_node,
};
use crate::types::FragmentId;
pub type EdgeDict = FxHashMap<(FragmentId, FragmentId), f64>;
pub type EdgeCategories = FxHashMap<(FragmentId, FragmentId), EdgeCategory>;
use crate::types::Fragment;
use self::base::EdgeBuilder;
pub const NAMING_WEIGHT_FLOOR: f64 = 0.30;
const EXPENSIVE_CATEGORIES: &[&str] = &["similarity", "history"];
struct BuilderCategory {
name: &'static str,
builders: fn() -> Vec<Box<dyn EdgeBuilder>>,
}
fn builder_categories() -> Vec<BuilderCategory> {
vec![
BuilderCategory {
name: "semantic",
builders: || semantic::get_semantic_builders(),
},
BuilderCategory {
name: "structural",
builders: || structural::get_structural_builders(),
},
BuilderCategory {
name: "config",
builders: || config_edges::get_config_builders(),
},
BuilderCategory {
name: "document",
builders: || document::get_document_builders(),
},
BuilderCategory {
name: "similarity",
builders: || similarity::get_similarity_builders(),
},
BuilderCategory {
name: "history",
builders: || history::get_history_builders(),
},
]
}
pub fn get_all_builders() -> Vec<Box<dyn EdgeBuilder>> {
let mut all = Vec::new();
for cat in builder_categories() {
all.extend((cat.builders)());
}
all
}
fn pack_pair(src: u32, dst: u32) -> u64 {
((src as u64) << 32) | dst as u64
}
struct LoggedEmission {
src: u32,
dst: u32,
weight: f64,
}
pub fn collect_capped_edges(
fragments: &[Fragment],
repo_root: Option<&Path>,
skip_expensive: bool,
deadline: crate::deadline::Deadline,
) -> CappedEdges {
let mut all_builders: Vec<(&str, Box<dyn EdgeBuilder>)> = Vec::new();
for cat in builder_categories() {
if skip_expensive && EXPENSIVE_CATEGORIES.contains(&cat.name) {
debug!("skipping {} edge builders (skip_expensive=true)", cat.name);
continue;
}
for builder in (cat.builders)() {
all_builders.push((cat.name, builder));
}
}
let (node_to_idx, idx_to_node) = intern_fragment_nodes(fragments);
let category_weights = *crate::config::category_weights::CATEGORY_WEIGHTS;
let builder_meta: Vec<(EdgeCategory, f64)> = all_builders
.iter()
.map(|(cat_name, builder)| {
let category = EdgeCategory::from_str(builder.category_label().unwrap_or(cat_name));
(category, category_weights.multiplier(category))
})
.collect();
let fallback_flags: Vec<bool> = all_builders
.iter()
.map(|(_, builder)| builder.is_fallback())
.collect();
let per_builder_log: Vec<Vec<LoggedEmission>> = all_builders
.par_iter()
.enumerate()
.map(|(builder_idx, (name, builder))| {
deadline.check("edge construction");
let _in_builder = deadline.enter();
let t = std::time::Instant::now();
let edges = builder.build(fragments, repo_root);
if std::env::var_os("DIFFCTX_TRACE_BUILDERS").is_some() {
eprintln!(
"builder {name}[{builder_idx}]: {:.1}s, {} edges",
t.elapsed().as_secs_f64(),
edges.len()
);
}
let mut log = Vec::with_capacity(edges.len());
for ((src, dst), weight) in edges {
let (Some(&s), Some(&d)) = (node_to_idx.get(&src), node_to_idx.get(&dst)) else {
continue;
};
log.push(LoggedEmission {
src: s,
dst: d,
weight,
});
}
log
})
.collect();
drop(all_builders);
let mut per_builder_log = per_builder_log;
if fallback_flags.iter().any(|&f| f) {
let mut dedicated_files: FxHashSet<&str> = FxHashSet::default();
for (builder_idx, log) in per_builder_log.iter().enumerate() {
if fallback_flags[builder_idx] || builder_meta[builder_idx].0 != EdgeCategory::Semantic
{
continue;
}
for e in log {
dedicated_files.insert(idx_to_node[e.src as usize].path.as_ref());
dedicated_files.insert(idx_to_node[e.dst as usize].path.as_ref());
}
}
for (builder_idx, log) in per_builder_log.iter_mut().enumerate() {
if !fallback_flags[builder_idx] {
continue;
}
log.retain(|e| {
!dedicated_files.contains(idx_to_node[e.src as usize].path.as_ref())
|| !dedicated_files.contains(idx_to_node[e.dst as usize].path.as_ref())
});
}
}
let n_nodes = idx_to_node.len();
let mut in_degree = vec![0u32; n_nodes];
let mut out_degree = vec![0u32; n_nodes];
let mut category_entries: Vec<(u32, u32, EdgeCategory)> = Vec::new();
let mut sem_out_files: FxHashMap<u32, FxHashSet<&str>> = FxHashMap::default();
let mut raw_by_category: FxHashMap<EdgeCategory, u64> = FxHashMap::default();
let mut deduped_by_category: FxHashMap<EdgeCategory, u64> = FxHashMap::default();
let mut seen: FxHashSet<u64> = FxHashSet::default();
for (builder_idx, log) in per_builder_log.iter().enumerate() {
let (category, _) = builder_meta[builder_idx];
*raw_by_category.entry(category).or_default() += log.len() as u64;
for e in log {
if !seen.insert(pack_pair(e.src, e.dst)) {
continue;
}
*deduped_by_category.entry(category).or_default() += 1;
in_degree[e.dst as usize] += 1;
out_degree[e.src as usize] += 1;
category_entries.push((e.src, e.dst, category));
if category == EdgeCategory::Semantic {
sem_out_files
.entry(e.src)
.or_default()
.insert(idx_to_node[e.dst as usize].path.as_ref());
}
}
}
drop(seen);
let mut emissions_by_category: Vec<(EdgeCategory, u64, u64)> = raw_by_category
.iter()
.map(|(&category, &raw)| {
let deduped = deduped_by_category.get(&category).copied().unwrap_or(0);
(category, raw, deduped)
})
.collect();
emissions_by_category.sort_unstable_by_key(|e| e.0.as_str());
category_entries.sort_unstable_by_key(|e| (e.0, e.1));
let mut sem_file_deg = vec![0u32; n_nodes];
for (&src, files) in &sem_out_files {
sem_file_deg[src as usize] = files.len() as u32;
}
drop(sem_out_files);
let deduped_edge_count = category_entries.len();
let factors = SuppressionFactors::from_counters(in_degree, sem_file_deg);
let max_per_node = read_max_out_edges_per_node();
let capped_per_builder: Vec<Vec<CompactEdge>> = per_builder_log
.into_par_iter()
.enumerate()
.map(|(builder_idx, log)| {
let (builder_category, multiplier) = builder_meta[builder_idx];
let mut per_source: FxHashMap<u32, SourceTopK> = FxHashMap::default();
for e in log {
let category = category_entries
.binary_search_by_key(&(e.src, e.dst), |c| (c.0, c.1))
.map(|k| category_entries[k].2)
.unwrap_or(builder_category);
let damped = factors.damp(e.weight * multiplier, category, e.src, e.dst);
let naming = (category == EdgeCategory::Semantic
|| category == EdgeCategory::Config)
&& e.weight > NAMING_WEIGHT_FLOOR;
push_bounded_top_k(
per_source.entry(e.src).or_default(),
RankedCandidate {
weight: damped,
dst: e.dst,
category,
naming,
},
max_per_node,
);
}
let mut survivors =
Vec::with_capacity(per_source.values().map(|h| h.len()).sum::<usize>());
for (src, heap) in per_source {
for Reverse(c) in heap {
survivors.push(CompactEdge {
src,
dst: c.dst,
weight: c.weight,
category: c.category,
naming: c.naming,
});
}
}
survivors
})
.collect();
let total: usize = capped_per_builder.iter().map(|v| v.len()).sum();
let mut edges: Vec<CompactEdge> = Vec::with_capacity(total);
for v in capped_per_builder {
edges.extend(v);
}
dedup_compact_edges(&mut edges);
cap_out_edges_per_source(&mut edges, max_per_node);
let nodes_capped = out_degree
.iter()
.filter(|&&d| d as usize > max_per_node)
.count();
let cap_stats = EdgeCapStats {
edges_before_cap: deduped_edge_count,
edges_after_cap: edges.len(),
edges_dropped_by_cap: deduped_edge_count - edges.len(),
nodes_capped,
max_out_edges_per_node: max_per_node,
emissions_by_category,
};
CappedEdges {
node_to_idx,
idx_to_node,
edges,
cap_stats,
}
}
pub fn discover_all_related_files(
changed_files: &[PathBuf],
all_candidates: &[PathBuf],
repo_root: Option<&Path>,
file_cache: Option<&FxHashMap<PathBuf, String>>,
) -> Vec<PathBuf> {
let mut discovered: FxHashMap<PathBuf, ()> = FxHashMap::default();
for builder in get_all_builders() {
for f in
builder.discover_related_files(changed_files, all_candidates, repo_root, file_cache)
{
discovered.entry(f).or_insert(());
}
}
let mut result: Vec<PathBuf> = discovered.into_keys().collect();
result.sort();
result
}
pub fn naming_reachable_files(
capped: &CappedEdges,
core_ids: &rustc_hash::FxHashSet<crate::types::FragmentId>,
max_depth: usize,
) -> rustc_hash::FxHashSet<std::sync::Arc<str>> {
let n = capped.idx_to_node.len();
let mut adj: Vec<Vec<u32>> = vec![Vec::new(); n];
for e in &capped.edges {
if e.naming {
adj[e.src as usize].push(e.dst);
adj[e.dst as usize].push(e.src);
}
}
let mut seen = vec![false; n];
let mut frontier: Vec<u32> = Vec::new();
for (i, id) in capped.idx_to_node.iter().enumerate() {
if core_ids.contains(id) {
seen[i] = true;
frontier.push(i as u32);
}
}
let mut files: rustc_hash::FxHashSet<std::sync::Arc<str>> = frontier
.iter()
.map(|&i| capped.idx_to_node[i as usize].path.clone())
.collect();
for _ in 0..max_depth {
let mut next = Vec::new();
for &u in &frontier {
for &v in &adj[u as usize] {
if !seen[v as usize] {
seen[v as usize] = true;
files.insert(capped.idx_to_node[v as usize].path.clone());
next.push(v);
}
}
}
if next.is_empty() {
break;
}
frontier = next;
}
files
}
#[cfg(test)]
mod fallback_gate_tests {
use super::*;
use rustc_hash::FxHashSet as Set;
use std::sync::Arc;
fn frag(path: &str, content: &str, idents: &[&str]) -> Fragment {
Fragment {
id: crate::types::FragmentId::new(Arc::from(path), 1, 10),
kind: crate::types::FragmentKind::Function,
content: Arc::from(content),
identifiers: idents.iter().map(|s| s.to_string()).collect::<Set<_>>(),
token_count: 10,
symbol_name: None,
}
}
#[test]
fn tags_edges_survive_only_where_dedicated_builders_came_back_empty() {
let fragments = vec![
frag(
"proj/a.c",
"#include \"bdep.h\"\nint zzcommonzz;\n",
&["zzcommonzz"],
),
frag("proj/bdep.h", "int bdecl(void);\n", &["bdecl"]),
frag(
"proj/c.c",
"#include \"ddep.h\"\nint zzcommonzz;\n",
&["zzcommonzz"],
),
frag("proj/ddep.h", "int ddecl(void);\n", &["ddecl"]),
frag("proj/u1.xyz", "zzcommonzz here\n", &["zzcommonzz"]),
frag("proj/u2.xyz", "zzcommonzz there\n", &["zzcommonzz"]),
];
let capped =
collect_capped_edges(&fragments, None, false, crate::deadline::Deadline::none());
let node_path = |idx: u32| capped.idx_to_node[idx as usize].path.clone();
let has = |a: &str, b: &str| {
capped.edges.iter().any(|e| {
if e.category != EdgeCategory::Semantic {
return false;
}
let s = node_path(e.src);
let d = node_path(e.dst);
(s.ends_with(a) && d.ends_with(b)) || (s.ends_with(b) && d.ends_with(a))
})
};
assert!(
has("u1.xyz", "u2.xyz"),
"fallback must still connect files no dedicated builder covers"
);
assert!(
!has("a.c", "c.c"),
"a tags-only link between two dedicated-covered files is the measured noise class (#131)"
);
}
}