Skip to main content

_diffctx/edges/
base.rs

1use std::path::{Path, PathBuf};
2
3use once_cell::sync::Lazy;
4use rustc_hash::{FxHashMap, FxHashSet};
5
6use crate::config::edge_weights::SEMANTIC_DISCOVERY;
7use crate::config::extensions::CODE_EXTENSIONS;
8use crate::types::{Fragment, FragmentId};
9
10use super::EdgeDict;
11
12pub trait EdgeBuilder: Send + Sync {
13    fn build(&self, fragments: &[Fragment], repo_root: Option<&Path>) -> EdgeDict;
14
15    fn discover_related_files(
16        &self,
17        _changed: &[PathBuf],
18        _candidates: &[PathBuf],
19        _repo_root: Option<&Path>,
20        _file_cache: Option<&FxHashMap<PathBuf, String>>,
21    ) -> Vec<PathBuf> {
22        vec![]
23    }
24
25    fn category_label(&self) -> Option<&str> {
26        None
27    }
28
29    fn is_expensive(&self) -> bool {
30        false
31    }
32}
33
34static INDEX_FILE_STEMS: Lazy<FxHashSet<&str>> =
35    Lazy::new(|| ["__init__", "index", "mod"].iter().copied().collect());
36
37fn strip_source_prefix(parts: &[&str]) -> Vec<String> {
38    for (i, part) in parts.iter().enumerate() {
39        if *part == "src" || *part == "lib" || *part == "packages" {
40            return parts[i + 1..].iter().map(|s| s.to_string()).collect();
41        }
42    }
43    parts.iter().map(|s| s.to_string()).collect()
44}
45
46fn strip_file_extension(stem: &str) -> &str {
47    for ext in CODE_EXTENSIONS.iter() {
48        if let Some(stripped) = stem.strip_suffix(ext) {
49            return stripped;
50        }
51    }
52    stem
53}
54
55pub fn path_to_module(path: &Path, repo_root: Option<&Path>) -> String {
56    let effective = if let Some(root) = repo_root {
57        if path.is_absolute() {
58            path.strip_prefix(root).unwrap_or(path)
59        } else {
60            path
61        }
62    } else {
63        path
64    };
65
66    let parts_raw: Vec<&str> = effective.iter().filter_map(|c| c.to_str()).collect();
67    let mut parts = strip_source_prefix(&parts_raw);
68
69    if let Some(last) = parts.last_mut() {
70        let stripped = strip_file_extension(last).to_string();
71        *last = stripped;
72    }
73
74    if let Some(last) = parts.last() {
75        if INDEX_FILE_STEMS.contains(last.as_str()) {
76            parts.pop();
77        }
78    }
79
80    parts.join(".")
81}
82
83pub struct FragmentIndex {
84    pub by_name: FxHashMap<String, Vec<FragmentId>>,
85    pub by_path: FxHashMap<String, Vec<FragmentId>>,
86}
87
88impl FragmentIndex {
89    pub fn new(fragments: &[Fragment], repo_root: Option<&Path>) -> Self {
90        let mut by_name: FxHashMap<String, Vec<FragmentId>> = FxHashMap::default();
91        let mut by_path: FxHashMap<String, Vec<FragmentId>> = FxHashMap::default();
92
93        for f in fragments {
94            let path = Path::new(f.path());
95            if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
96                by_name
97                    .entry(name.to_lowercase())
98                    .or_default()
99                    .push(f.id.clone());
100            }
101            by_path
102                .entry(f.path().to_string())
103                .or_default()
104                .push(f.id.clone());
105
106            if let Some(root) = repo_root {
107                if let Ok(rel) = Path::new(f.path()).strip_prefix(root) {
108                    let rel_str = rel.to_string_lossy().to_string();
109                    by_path
110                        .entry(rel_str.clone())
111                        .or_default()
112                        .push(f.id.clone());
113                    let posix = rel_str.replace('\\', "/");
114                    if posix != rel_str {
115                        by_path.entry(posix).or_default().push(f.id.clone());
116                    }
117                }
118            }
119        }
120
121        Self { by_name, by_path }
122    }
123}
124
125pub fn add_edge(
126    edges: &mut EdgeDict,
127    src: &FragmentId,
128    dst: &FragmentId,
129    weight: f64,
130    reverse_factor: f64,
131) {
132    let key_fwd = (src.clone(), dst.clone());
133    let existing_fwd = edges.get(&key_fwd).copied().unwrap_or(0.0);
134    if weight > existing_fwd {
135        edges.insert(key_fwd, weight);
136    }
137    let rev_w = weight * reverse_factor;
138    let key_rev = (dst.clone(), src.clone());
139    let existing_rev = edges.get(&key_rev).copied().unwrap_or(0.0);
140    if rev_w > existing_rev {
141        edges.insert(key_rev, rev_w);
142    }
143}
144
145pub fn add_edge_unidirectional(
146    edges: &mut EdgeDict,
147    src: &FragmentId,
148    dst: &FragmentId,
149    weight: f64,
150) {
151    let key = (src.clone(), dst.clone());
152    let existing = edges.get(&key).copied().unwrap_or(0.0);
153    if weight > existing {
154        edges.insert(key, weight);
155    }
156}
157
158pub fn add_edges_from_ids(
159    edges: &mut EdgeDict,
160    src: &FragmentId,
161    targets: &[FragmentId],
162    weight: f64,
163    reverse_factor: f64,
164) {
165    for target in targets {
166        if target != src {
167            add_edge(edges, src, target, weight, reverse_factor);
168        }
169    }
170}
171
172pub fn link_by_name(
173    src_id: &FragmentId,
174    name: &str,
175    idx: &FragmentIndex,
176    edges: &mut EdgeDict,
177    weight: f64,
178    reverse_factor: f64,
179) {
180    let target = name.split('/').next_back().unwrap_or(name).to_lowercase();
181    if let Some(frag_ids) = idx.by_name.get(&target) {
182        for fid in frag_ids {
183            if fid != src_id {
184                add_edge(edges, src_id, fid, weight, reverse_factor);
185                return;
186            }
187        }
188    }
189    link_by_path_match(src_id, name, idx, edges, weight, reverse_factor);
190}
191
192pub fn link_by_path_match(
193    src_id: &FragmentId,
194    ref_str: &str,
195    idx: &FragmentIndex,
196    edges: &mut EdgeDict,
197    weight: f64,
198    reverse_factor: f64,
199) {
200    let ref_lower = ref_str.to_lowercase();
201    for (path_str, frag_ids) in &idx.by_path {
202        if path_str.contains(ref_str) || path_str.to_lowercase().contains(&ref_lower) {
203            for fid in frag_ids {
204                if fid != src_id {
205                    add_edge(edges, src_id, fid, weight, reverse_factor);
206                }
207            }
208        }
209    }
210}
211
212pub fn read_file_cached<'a>(
213    path: &Path,
214    cache: Option<&'a FxHashMap<PathBuf, String>>,
215) -> Option<String> {
216    if let Some(c) = cache {
217        if let Some(content) = c.get(path) {
218            return Some(content.clone());
219        }
220    }
221    std::fs::read_to_string(path).ok()
222}
223
224fn candidate_rel_path(candidate: &Path, repo_root: Option<&Path>) -> String {
225    if let Some(root) = repo_root {
226        if let Ok(rel) = candidate.strip_prefix(root) {
227            return rel.to_string_lossy().to_lowercase();
228        }
229    }
230    candidate
231        .file_name()
232        .map(|n| n.to_string_lossy().to_lowercase())
233        .unwrap_or_default()
234}
235
236fn matches_any_ref(candidate_name: &str, candidate_rel: &str, refs: &FxHashSet<String>) -> bool {
237    for r in refs {
238        let ref_name = r.split('/').next_back().unwrap_or(r).to_lowercase();
239        if candidate_name == ref_name {
240            return true;
241        }
242        let ref_lower = r.to_lowercase();
243        if ref_lower.len() >= SEMANTIC_DISCOVERY.min_ref_length_for_path_match {
244            if let Some(idx) = candidate_rel.find(&ref_lower) {
245                let end_idx = idx + ref_lower.len();
246                let start_ok = idx == 0
247                    || candidate_rel.as_bytes().get(idx - 1) == Some(&b'/')
248                    || candidate_rel.as_bytes().get(idx - 1) == Some(&b'\\');
249                let end_ok = end_idx == candidate_rel.len()
250                    || matches!(
251                        candidate_rel.as_bytes().get(end_idx),
252                        Some(b'/') | Some(b'\\') | Some(b'.')
253                    );
254                if start_ok && end_ok {
255                    return true;
256                }
257            }
258        }
259    }
260    false
261}
262
263pub fn discover_files_by_refs(
264    refs: &FxHashSet<String>,
265    changed_files: &[PathBuf],
266    all_candidates: &[PathBuf],
267    repo_root: Option<&Path>,
268) -> Vec<PathBuf> {
269    if refs.is_empty() {
270        return vec![];
271    }
272    let changed_set: FxHashSet<&PathBuf> = changed_files.iter().collect();
273    let mut discovered = Vec::new();
274    for candidate in all_candidates {
275        if changed_set.contains(candidate) {
276            continue;
277        }
278        let candidate_name = candidate
279            .file_name()
280            .map(|n| n.to_string_lossy().to_lowercase())
281            .unwrap_or_default();
282        let candidate_rel = candidate_rel_path(candidate, repo_root);
283        if matches_any_ref(&candidate_name, &candidate_rel, refs) {
284            discovered.push(candidate.clone());
285        }
286    }
287    discovered
288}
289
290pub fn file_ext(path: &Path) -> String {
291    path.extension()
292        .map(|e| format!(".{}", e.to_string_lossy().to_lowercase()))
293        .unwrap_or_default()
294}