Skip to main content

_diffctx/edges/semantic/
go.rs

1use std::path::{Path, PathBuf};
2
3use once_cell::sync::Lazy;
4use regex::Regex;
5use rustc_hash::{FxHashMap, FxHashSet};
6
7use crate::config::edge_weights::{GO_SEMANTIC, SEMANTIC_DISCOVERY};
8use crate::config::extensions::GO_EXTENSIONS;
9use crate::config::weights::EDGE_WEIGHTS;
10use crate::types::{Fragment, FragmentId};
11
12use super::super::EdgeDict;
13use super::super::base::{self, EdgeBuilder, add_edge, add_edges_from_ids};
14
15fn is_go_file(path: &Path) -> bool {
16    let ext = base::file_ext(path);
17    GO_EXTENSIONS.contains(ext.as_str())
18}
19
20static IMPORT_RE: Lazy<Regex> =
21    Lazy::new(|| Regex::new(r#"(?m)^\s*(?:import\s+)?"([^"]+)""#).unwrap());
22static PACKAGE_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?m)^\s*package\s+(\w+)").unwrap());
23static TYPE_DEF_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?m)^\s*type\s+([A-Z]\w*)").unwrap());
24static FUNC_DEF_RE: Lazy<Regex> =
25    Lazy::new(|| Regex::new(r"(?m)^\s*func\s+(?:\([^)]*\)\s+)?([A-Z]\w*)\s*\(").unwrap());
26static FUNC_CALL_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\b([A-Z]\w*)\s*\(").unwrap());
27static TYPE_REF_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\b([A-Z]\w*)\b").unwrap());
28static PKG_CALL_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\b([a-z]\w+)\.([A-Z]\w*)").unwrap());
29static INIT_FUNC_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?m)^\s*func\s+init\s*\(").unwrap());
30
31fn extract_imports(content: &str) -> FxHashSet<String> {
32    IMPORT_RE
33        .captures_iter(content)
34        .map(|c| c[1].to_string())
35        .collect()
36}
37
38fn get_package_name(content: &str) -> String {
39    PACKAGE_RE
40        .captures(content)
41        .map(|c| c[1].to_string())
42        .unwrap_or_else(|| "main".to_string())
43}
44
45fn extract_definitions(content: &str) -> (FxHashSet<String>, FxHashSet<String>) {
46    let funcs: FxHashSet<String> = FUNC_DEF_RE
47        .captures_iter(content)
48        .map(|c| c[1].to_string())
49        .collect();
50    let types: FxHashSet<String> = TYPE_DEF_RE
51        .captures_iter(content)
52        .map(|c| c[1].to_string())
53        .collect();
54    (funcs, types)
55}
56
57fn extract_references(
58    content: &str,
59) -> (
60    FxHashSet<String>,
61    FxHashSet<String>,
62    FxHashSet<(String, String)>,
63) {
64    let func_calls: FxHashSet<String> = FUNC_CALL_RE
65        .captures_iter(content)
66        .map(|c| c[1].to_string())
67        .collect();
68    let type_refs: FxHashSet<String> = TYPE_REF_RE
69        .captures_iter(content)
70        .map(|c| c[1].to_string())
71        .collect();
72    let pkg_calls: FxHashSet<(String, String)> = PKG_CALL_RE
73        .captures_iter(content)
74        .map(|c| (c[1].to_string(), c[2].to_string()))
75        .collect();
76    (func_calls, type_refs, pkg_calls)
77}
78
79fn has_init_func(content: &str) -> bool {
80    INIT_FUNC_RE.is_match(content)
81}
82
83pub struct GoEdgeBuilder;
84
85impl GoEdgeBuilder {
86    fn build_indices(
87        &self,
88        go_frags: &[&Fragment],
89        repo_root: Option<&Path>,
90    ) -> (
91        FxHashMap<String, Vec<FragmentId>>,
92        FxHashMap<String, Vec<FragmentId>>,
93        FxHashMap<String, Vec<FragmentId>>,
94        FxHashMap<String, Vec<FragmentId>>,
95    ) {
96        let mut pkg_to_frags: FxHashMap<String, Vec<FragmentId>> = FxHashMap::default();
97        let mut path_to_frags: FxHashMap<String, Vec<FragmentId>> = FxHashMap::default();
98        let mut type_defs: FxHashMap<String, Vec<FragmentId>> = FxHashMap::default();
99        let mut func_defs: FxHashMap<String, Vec<FragmentId>> = FxHashMap::default();
100
101        for f in go_frags {
102            let pkg = get_package_name(&f.content).to_lowercase();
103            pkg_to_frags.entry(pkg).or_default().push(f.id.clone());
104
105            if let Some(root) = repo_root {
106                if let Ok(rel) = Path::new(f.path()).strip_prefix(root) {
107                    if let Some(parent) = rel.parent() {
108                        path_to_frags
109                            .entry(parent.to_string_lossy().to_string())
110                            .or_default()
111                            .push(f.id.clone());
112                    }
113                }
114            }
115
116            let (funcs, types) = extract_definitions(&f.content);
117            for t in types {
118                type_defs
119                    .entry(t.to_lowercase())
120                    .or_default()
121                    .push(f.id.clone());
122            }
123            for func in funcs {
124                func_defs
125                    .entry(func.to_lowercase())
126                    .or_default()
127                    .push(f.id.clone());
128            }
129        }
130
131        (pkg_to_frags, path_to_frags, type_defs, func_defs)
132    }
133}
134
135impl EdgeBuilder for GoEdgeBuilder {
136    fn build(&self, fragments: &[Fragment], repo_root: Option<&Path>) -> EdgeDict {
137        let go_frags: Vec<&Fragment> = fragments
138            .iter()
139            .filter(|f| is_go_file(Path::new(f.path())))
140            .collect();
141        if go_frags.is_empty() {
142            return FxHashMap::default();
143        }
144
145        let import_weight = EDGE_WEIGHTS["go_import"].forward;
146        let type_weight = EDGE_WEIGHTS["go_type"].forward;
147        let func_weight = EDGE_WEIGHTS["go_func"].forward;
148        let same_package_weight = EDGE_WEIGHTS["go_same_package"].forward;
149        let reverse_factor = EDGE_WEIGHTS["go_import"].reverse_factor;
150        let init_same_package_weight = GO_SEMANTIC.init_same_package_weight;
151
152        let (pkg_to_frags, path_to_frags, type_defs, func_defs) =
153            self.build_indices(&go_frags, repo_root);
154
155        let mut edges: EdgeDict = FxHashMap::default();
156
157        for gf in &go_frags {
158            let imports = extract_imports(&gf.content);
159            let (func_calls, type_refs, pkg_calls) = extract_references(&gf.content);
160
161            for imp in &imports {
162                let imp_pkg = imp.split('/').next_back().unwrap_or(imp).to_lowercase();
163                for (pkg, frag_ids) in &pkg_to_frags {
164                    if *pkg == imp_pkg {
165                        add_edges_from_ids(
166                            &mut edges,
167                            &gf.id,
168                            frag_ids,
169                            import_weight,
170                            reverse_factor,
171                        );
172                    }
173                }
174                for (path_str, frag_ids) in &path_to_frags {
175                    if *imp == *path_str
176                        || imp.ends_with(&format!("/{}", path_str))
177                        || imp.contains(&format!("/{}/", path_str))
178                    {
179                        add_edges_from_ids(
180                            &mut edges,
181                            &gf.id,
182                            frag_ids,
183                            import_weight,
184                            reverse_factor,
185                        );
186                    }
187                }
188            }
189
190            for type_ref in &type_refs {
191                for fid in type_defs.get(&type_ref.to_lowercase()).unwrap_or(&vec![]) {
192                    if fid != &gf.id {
193                        add_edge(&mut edges, &gf.id, fid, type_weight, reverse_factor);
194                    }
195                }
196            }
197
198            for func_call in &func_calls {
199                for fid in func_defs.get(&func_call.to_lowercase()).unwrap_or(&vec![]) {
200                    if fid != &gf.id {
201                        add_edge(&mut edges, &gf.id, fid, func_weight, reverse_factor);
202                    }
203                }
204            }
205
206            for (pkg_name, _symbol) in &pkg_calls {
207                for fid in pkg_to_frags
208                    .get(&pkg_name.to_lowercase())
209                    .unwrap_or(&vec![])
210                {
211                    if fid != &gf.id {
212                        add_edge(&mut edges, &gf.id, fid, func_weight, reverse_factor);
213                    }
214                }
215            }
216
217            let has_init = has_init_func(&gf.content);
218            let sp_weight = if has_init {
219                init_same_package_weight
220            } else {
221                same_package_weight
222            };
223            let current_pkg = get_package_name(&gf.content).to_lowercase();
224            for fid in pkg_to_frags.get(&current_pkg).unwrap_or(&vec![]) {
225                if fid != &gf.id {
226                    add_edge(&mut edges, &gf.id, fid, sp_weight, reverse_factor);
227                }
228            }
229        }
230
231        edges
232    }
233
234    fn discover_related_files(
235        &self,
236        changed: &[PathBuf],
237        candidates: &[PathBuf],
238        _repo_root: Option<&Path>,
239        file_cache: Option<&FxHashMap<PathBuf, String>>,
240    ) -> Vec<PathBuf> {
241        let go_changed: Vec<&PathBuf> = changed.iter().filter(|f| is_go_file(f)).collect();
242        if go_changed.is_empty() {
243            return vec![];
244        }
245
246        let changed_set: FxHashSet<PathBuf> = changed.iter().cloned().collect();
247        let go_candidates: Vec<PathBuf> = candidates
248            .iter()
249            .filter(|c| !changed_set.contains(*c) && is_go_file(c))
250            .cloned()
251            .collect();
252
253        let mut discovered: FxHashSet<PathBuf> = FxHashSet::default();
254
255        let pkg_dirs: FxHashSet<PathBuf> = go_changed
256            .iter()
257            .filter_map(|f| f.parent().map(|p| p.to_path_buf()))
258            .collect();
259        for c in &go_candidates {
260            if let Some(parent) = c.parent() {
261                if pkg_dirs.contains(&parent.to_path_buf()) {
262                    discovered.insert(c.clone());
263                }
264            }
265        }
266
267        let mut candidate_index: FxHashMap<PathBuf, (String, FxHashSet<String>)> =
268            FxHashMap::default();
269        for c in &go_candidates {
270            let content = base::read_file_cached(c, file_cache);
271            if let Some(content) = content {
272                let pkg = get_package_name(&content).to_lowercase();
273                let imports = extract_imports(&content);
274                candidate_index.insert(c.clone(), (pkg, imports));
275            }
276        }
277
278        let mut frontier: FxHashSet<PathBuf> = go_changed.iter().map(|f| (*f).clone()).collect();
279
280        for _ in 0..SEMANTIC_DISCOVERY.max_depth {
281            let mut next_frontier: FxHashSet<PathBuf> = FxHashSet::default();
282            for f in &frontier {
283                let content = base::read_file_cached(f, file_cache);
284                if let Some(content) = content {
285                    let f_imports = extract_imports(&content);
286                    let f_pkg = get_package_name(&content).to_lowercase();
287
288                    for c in &go_candidates {
289                        if changed_set.contains(c) || discovered.contains(c) {
290                            continue;
291                        }
292                        if let Some((c_pkg, c_imports)) = candidate_index.get(c) {
293                            let forward_match = f_imports.iter().any(|imp| {
294                                imp.split('/').next_back().unwrap_or(imp).to_lowercase() == *c_pkg
295                            });
296                            let reverse_match = c_imports.iter().any(|imp| {
297                                imp.split('/').next_back().unwrap_or(imp).to_lowercase() == f_pkg
298                            });
299                            if forward_match || reverse_match {
300                                discovered.insert(c.clone());
301                                next_frontier.insert(c.clone());
302                            }
303                        }
304                    }
305                }
306            }
307            if next_frontier.is_empty() {
308                break;
309            }
310            frontier = next_frontier;
311        }
312
313        let mut result: Vec<PathBuf> = discovered.into_iter().collect();
314        result.sort();
315        result
316    }
317}