_diffctx/edges/semantic/

go.rs

use std::path::{Path, PathBuf};

/// Same ambiguity bar as `CFamilySemanticWeights::max_files_per_name`.
const MAX_FILES_PER_NAME: usize = 8;

use once_cell::sync::Lazy;
use regex::Regex;
use rustc_hash::{FxHashMap, FxHashSet};

use crate::config::edge_weights::{GO_SEMANTIC, SEMANTIC_DISCOVERY};
use crate::config::extensions::GO_EXTENSIONS;
use crate::config::weights::EDGE_WEIGHTS;
use crate::types::{Fragment, FragmentId};

use super::super::EdgeDict;
use super::super::base::{self, EdgeBuilder, add_edge, add_edges_from_ids};

fn is_go_file(path: &Path) -> bool {
    let ext = base::file_ext(path);
    GO_EXTENSIONS.contains(ext.as_str())
}

static IMPORT_RE: Lazy<Regex> =
    Lazy::new(|| Regex::new(r#"(?m)^\s*(?:import\s+)?"([^"]+)""#).unwrap());
static PACKAGE_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?m)^\s*package\s+(\w+)").unwrap());
static TYPE_DEF_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?m)^\s*type\s+([A-Z]\w*)").unwrap());
static FUNC_DEF_RE: Lazy<Regex> =
    Lazy::new(|| Regex::new(r"(?m)^\s*func\s+(?:\([^)]*\)\s+)?([A-Z]\w*)\s*\(").unwrap());
static FUNC_CALL_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\b([A-Z]\w*)\s*\(").unwrap());
static TYPE_REF_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\b([A-Z]\w*)\b").unwrap());
static PKG_CALL_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\b([a-z]\w+)\.([A-Z]\w*)").unwrap());
static INIT_FUNC_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"(?m)^\s*func\s+init\s*\(").unwrap());

fn extract_imports(content: &str) -> FxHashSet<String> {
    IMPORT_RE
        .captures_iter(content)
        .map(|c| c[1].to_string())
        .collect()
}

/// None when the fragment carries no `package` line — i.e. every
/// function-body fragment. The old `"main"` fallback funneled all of them
/// into one bucket, and the same-package loop cross-linked essentially the
/// whole Go universe through it: 103.8M edges on a gitpod-scale monorepo
/// (#116).
fn get_package_name(content: &str) -> Option<String> {
    PACKAGE_RE.captures(content).map(|c| c[1].to_string())
}

fn extract_definitions(content: &str) -> (FxHashSet<String>, FxHashSet<String>) {
    let funcs: FxHashSet<String> = FUNC_DEF_RE
        .captures_iter(content)
        .map(|c| c[1].to_string())
        .collect();
    let types: FxHashSet<String> = TYPE_DEF_RE
        .captures_iter(content)
        .map(|c| c[1].to_string())
        .collect();
    (funcs, types)
}

fn extract_references(
    content: &str,
) -> (
    FxHashSet<String>,
    FxHashSet<String>,
    FxHashSet<(String, String)>,
) {
    let func_calls: FxHashSet<String> = FUNC_CALL_RE
        .captures_iter(content)
        .map(|c| c[1].to_string())
        .collect();
    let type_refs: FxHashSet<String> = TYPE_REF_RE
        .captures_iter(content)
        .map(|c| c[1].to_string())
        .collect();
    let pkg_calls: FxHashSet<(String, String)> = PKG_CALL_RE
        .captures_iter(content)
        .map(|c| (c[1].to_string(), c[2].to_string()))
        .collect();
    (func_calls, type_refs, pkg_calls)
}

fn has_init_func(content: &str) -> bool {
    INIT_FUNC_RE.is_match(content)
}

pub struct GoEdgeBuilder;

impl GoEdgeBuilder {
    fn build_indices<'a>(
        &self,
        go_frags: &'a [&'a Fragment],
        repo_root: Option<&Path>,
    ) -> (
        FxHashMap<String, Vec<FragmentId>>,
        FxHashMap<String, Vec<FragmentId>>,
        FxHashMap<String, Vec<FragmentId>>,
        FxHashMap<String, Vec<FragmentId>>,
        FxHashMap<String, FxHashSet<&'a str>>,
        FxHashMap<String, FxHashSet<&'a str>>,
    ) {
        let mut pkg_to_frags: FxHashMap<String, Vec<FragmentId>> = FxHashMap::default();
        let mut path_to_frags: FxHashMap<String, Vec<FragmentId>> = FxHashMap::default();
        let mut type_defs: FxHashMap<String, Vec<FragmentId>> = FxHashMap::default();
        let mut func_defs: FxHashMap<String, Vec<FragmentId>> = FxHashMap::default();
        let mut def_files: FxHashMap<String, FxHashSet<&str>> = FxHashMap::default();
        let mut pkg_files: FxHashMap<String, FxHashSet<&str>> = FxHashMap::default();

        for f in go_frags {
            if let Some(pkg) = get_package_name(&f.content) {
                let lower = pkg.to_lowercase();
                pkg_files.entry(lower.clone()).or_default().insert(f.path());
                pkg_to_frags.entry(lower).or_default().push(f.id.clone());
            }

            if let Some(root) = repo_root {
                if let Ok(rel) = Path::new(f.path()).strip_prefix(root) {
                    if let Some(parent) = rel.parent() {
                        path_to_frags
                            .entry(parent.to_string_lossy().to_string())
                            .or_default()
                            .push(f.id.clone());
                    }
                }
            }

            let (funcs, types) = extract_definitions(&f.content);
            for t in types {
                let lower = t.to_lowercase();
                def_files.entry(lower.clone()).or_default().insert(f.path());
                type_defs.entry(lower).or_default().push(f.id.clone());
            }
            for func in funcs {
                let lower = func.to_lowercase();
                def_files.entry(lower.clone()).or_default().insert(f.path());
                func_defs.entry(lower).or_default().push(f.id.clone());
            }
        }

        (
            pkg_to_frags,
            path_to_frags,
            type_defs,
            func_defs,
            def_files,
            pkg_files,
        )
    }
}

impl EdgeBuilder for GoEdgeBuilder {
    fn build(&self, fragments: &[Fragment], repo_root: Option<&Path>) -> EdgeDict {
        let go_frags: Vec<&Fragment> = fragments
            .iter()
            .filter(|f| is_go_file(Path::new(f.path())))
            .collect();
        if go_frags.is_empty() {
            return FxHashMap::default();
        }

        let import_weight = EDGE_WEIGHTS["go_import"].forward;
        let type_weight = EDGE_WEIGHTS["go_type"].forward;
        let func_weight = EDGE_WEIGHTS["go_func"].forward;
        let same_package_weight = EDGE_WEIGHTS["go_same_package"].forward;
        let reverse_factor = EDGE_WEIGHTS["go_import"].reverse_factor;
        let init_same_package_weight = GO_SEMANTIC.init_same_package_weight;

        let (pkg_to_frags, path_to_frags, type_defs, func_defs, def_files, pkg_files) =
            self.build_indices(&go_frags, repo_root);

        let name_capped = |name: &str| {
            def_files
                .get(name)
                .is_some_and(|s| s.len() > MAX_FILES_PER_NAME)
        };
        let pkg_capped = |name: &str| {
            pkg_files
                .get(name)
                .is_some_and(|s| s.len() > MAX_FILES_PER_NAME)
        };

        let mut edges: EdgeDict = FxHashMap::default();

        let mut path_last_comp: FxHashMap<&str, Vec<&String>> = FxHashMap::default();
        for path_str in path_to_frags.keys() {
            if let Some(last) = path_str.rsplit('/').next() {
                path_last_comp.entry(last).or_default().push(path_str);
            }
        }

        for gf in &go_frags {
            let imports = extract_imports(&gf.content);
            let (func_calls, type_refs, pkg_calls) = extract_references(&gf.content);

            for imp in &imports {
                let imp_pkg = imp.split('/').next_back().unwrap_or(imp).to_lowercase();
                // Direct lookup: the previous full-map scan was
                // O(imports x packages) and cost ~40s alone on a
                // kubernetes-scale commit (#196).
                if !pkg_capped(&imp_pkg) {
                    if let Some(frag_ids) = pkg_to_frags.get(&imp_pkg) {
                        add_edges_from_ids(
                            &mut edges,
                            &gf.id,
                            frag_ids,
                            import_weight,
                            reverse_factor,
                        );
                    }
                }
                // Last-component index instead of a full scan: the old form
                // was O(imports x dirs) with two String allocations per probe
                // and stood at ~39s alone on a kubernetes commit (#196). A
                // dir can only match if its last component appears as a
                // component of the import, so only that posting list is
                // verified against the full predicate.
                for part in imp.split('/') {
                    let Some(dirs) = path_last_comp.get(part) else {
                        continue;
                    };
                    for path_str in dirs {
                        if *imp == **path_str
                            || imp.ends_with(&format!("/{}", path_str))
                            || imp.contains(&format!("/{}/", path_str))
                        {
                            if let Some(frag_ids) = path_to_frags.get(*path_str) {
                                add_edges_from_ids(
                                    &mut edges,
                                    &gf.id,
                                    frag_ids,
                                    import_weight,
                                    reverse_factor,
                                );
                            }
                        }
                    }
                }
            }

            for type_ref in &type_refs {
                let lower = type_ref.to_lowercase();
                if name_capped(&lower) {
                    continue;
                }
                for fid in type_defs.get(&lower).unwrap_or(&vec![]) {
                    if fid != &gf.id {
                        add_edge(&mut edges, &gf.id, fid, type_weight, reverse_factor);
                    }
                }
            }

            for func_call in &func_calls {
                let lower = func_call.to_lowercase();
                if name_capped(&lower) {
                    continue;
                }
                for fid in func_defs.get(&lower).unwrap_or(&vec![]) {
                    if fid != &gf.id {
                        add_edge(&mut edges, &gf.id, fid, func_weight, reverse_factor);
                    }
                }
            }

            for (pkg_name, _symbol) in &pkg_calls {
                let lower = pkg_name.to_lowercase();
                if pkg_capped(&lower) {
                    continue;
                }
                for fid in pkg_to_frags.get(&lower).unwrap_or(&vec![]) {
                    if fid != &gf.id {
                        add_edge(&mut edges, &gf.id, fid, func_weight, reverse_factor);
                    }
                }
            }

            let has_init = has_init_func(&gf.content);
            let sp_weight = if has_init {
                init_same_package_weight
            } else {
                same_package_weight
            };
            if let Some(current_pkg) = get_package_name(&gf.content) {
                for fid in pkg_to_frags
                    .get(&current_pkg.to_lowercase())
                    .unwrap_or(&vec![])
                {
                    if fid != &gf.id {
                        add_edge(&mut edges, &gf.id, fid, sp_weight, reverse_factor);
                    }
                }
            }
        }

        edges
    }

    fn discover_related_files(
        &self,
        changed: &[PathBuf],
        candidates: &[PathBuf],
        _repo_root: Option<&Path>,
        file_cache: Option<&FxHashMap<PathBuf, String>>,
    ) -> Vec<PathBuf> {
        let go_changed: Vec<&PathBuf> = changed.iter().filter(|f| is_go_file(f)).collect();
        if go_changed.is_empty() {
            return vec![];
        }

        let changed_set: FxHashSet<PathBuf> = changed.iter().cloned().collect();
        let go_candidates: Vec<PathBuf> = candidates
            .iter()
            .filter(|c| !changed_set.contains(*c) && is_go_file(c))
            .cloned()
            .collect();

        let mut discovered: FxHashSet<PathBuf> = FxHashSet::default();

        let pkg_dirs: FxHashSet<PathBuf> = go_changed
            .iter()
            .filter_map(|f| f.parent().map(|p| p.to_path_buf()))
            .collect();
        for c in &go_candidates {
            if let Some(parent) = c.parent() {
                if pkg_dirs.contains(&parent.to_path_buf()) {
                    discovered.insert(c.clone());
                }
            }
        }

        let mut candidate_index: FxHashMap<PathBuf, (String, FxHashSet<String>)> =
            FxHashMap::default();
        for c in &go_candidates {
            let content = base::read_file_cached(c, file_cache);
            if let Some(content) = content {
                let Some(pkg) = get_package_name(&content).map(|p| p.to_lowercase()) else {
                    continue;
                };
                let imports = extract_imports(&content);
                candidate_index.insert(c.clone(), (pkg, imports));
            }
        }

        let mut frontier: FxHashSet<PathBuf> = go_changed.iter().map(|f| (*f).clone()).collect();

        for _ in 0..SEMANTIC_DISCOVERY.max_depth {
            let mut next_frontier: FxHashSet<PathBuf> = FxHashSet::default();
            for f in &frontier {
                let content = base::read_file_cached(f, file_cache);
                if let Some(content) = content {
                    let f_imports = extract_imports(&content);
                    let Some(f_pkg) = get_package_name(&content).map(|p| p.to_lowercase()) else {
                        continue;
                    };

                    for c in &go_candidates {
                        if changed_set.contains(c) || discovered.contains(c) {
                            continue;
                        }
                        if let Some((c_pkg, c_imports)) = candidate_index.get(c) {
                            let forward_match = f_imports.iter().any(|imp| {
                                imp.split('/').next_back().unwrap_or(imp).to_lowercase() == *c_pkg
                            });
                            let reverse_match = c_imports.iter().any(|imp| {
                                imp.split('/').next_back().unwrap_or(imp).to_lowercase() == f_pkg
                            });
                            if forward_match || reverse_match {
                                discovered.insert(c.clone());
                                next_frontier.insert(c.clone());
                            }
                        }
                    }
                }
            }
            if next_frontier.is_empty() {
                break;
            }
            frontier = next_frontier;
        }

        let mut result: Vec<PathBuf> = discovered.into_iter().collect();
        result.sort();
        result
    }
}