_diffctx/edges/semantic/
protobuf.rs

1use std::path::{Path, PathBuf};
2
3use once_cell::sync::Lazy;
4use regex::Regex;
5use rustc_hash::{FxHashMap, FxHashSet};
6
7use crate::config::weights::EDGE_WEIGHTS;
8use crate::types::Fragment;
9
10use super::super::EdgeDict;
11use super::super::base::{self, EdgeBuilder, add_edge, discover_files_by_refs};
12
13fn is_proto_file(path: &Path) -> bool {
14    base::file_ext(path) == ".proto"
15}
16
17static IMPORT_RE: Lazy<Regex> =
18    Lazy::new(|| Regex::new(r#"(?m)^\s*import\s+(?:public\s+)?["']([^"']+)["']"#).unwrap());
19static MSG_RE: Lazy<Regex> =
20    Lazy::new(|| Regex::new(r"(?m)^\s*(?:message|enum|service)\s+(\w+)").unwrap());
21static FIELD_TYPE_RE: Lazy<Regex> = Lazy::new(|| {
22    Regex::new(r"(?m)^\s*(?:repeated\s+|optional\s+|required\s+|map<\w+,\s*)?([A-Z]\w+)").unwrap()
23});
24static RPC_TYPE_RE: Lazy<Regex> = Lazy::new(|| {
25    Regex::new(r"(?:returns\s*\(|rpc\s+\w+\s*\()\s*(?:stream\s+)?([A-Z]\w+)").unwrap()
26});
27
28fn extract_imports(content: &str) -> FxHashSet<String> {
29    IMPORT_RE
30        .captures_iter(content)
31        .map(|c| c[1].to_string())
32        .collect()
33}
34
35fn extract_defs(content: &str) -> FxHashSet<String> {
36    MSG_RE
37        .captures_iter(content)
38        .map(|c| c[1].to_string())
39        .collect()
40}
41
42fn extract_type_refs(content: &str) -> FxHashSet<String> {
43    let mut refs: FxHashSet<String> = FIELD_TYPE_RE
44        .captures_iter(content)
45        .map(|c| c[1].to_string())
46        .collect();
47    refs.extend(RPC_TYPE_RE.captures_iter(content).map(|c| c[1].to_string()));
48    refs
49}
50
51pub struct ProtobufEdgeBuilder;
52
53impl EdgeBuilder for ProtobufEdgeBuilder {
54    fn build(&self, fragments: &[Fragment], repo_root: Option<&Path>) -> EdgeDict {
55        let frags: Vec<&Fragment> = fragments
56            .iter()
57            .filter(|f| is_proto_file(Path::new(f.path())))
58            .collect();
59        if frags.is_empty() {
60            return FxHashMap::default();
61        }
62
63        let import_w = EDGE_WEIGHTS["proto_import"].forward;
64        let msg_w = EDGE_WEIGHTS["proto_message_ref"].forward;
65        let rpc_w = EDGE_WEIGHTS["proto_service_rpc"].forward;
66        let reverse_factor = EDGE_WEIGHTS["proto_import"].reverse_factor;
67
68        let idx = base::FragmentIndex::new(fragments, repo_root);
69        let mut name_to_defs: FxHashMap<String, Vec<_>> = FxHashMap::default();
70        for f in &frags {
71            for name in extract_defs(&f.content) {
72                name_to_defs
73                    .entry(name.to_lowercase())
74                    .or_default()
75                    .push(f.id.clone());
76            }
77        }
78
79        let mut edges: EdgeDict = FxHashMap::default();
80
81        for f in &frags {
82            let self_defs = extract_defs(&f.content);
83            for imp in extract_imports(&f.content) {
84                base::link_by_name(&f.id, &imp, &idx, &mut edges, import_w, reverse_factor);
85            }
86            for tref in extract_type_refs(&f.content) {
87                if self_defs.contains(&tref) {
88                    continue;
89                }
90                let w = if RPC_TYPE_RE.captures_iter(&f.content).any(|c| c[1] == tref) {
91                    rpc_w
92                } else {
93                    msg_w
94                };
95                if let Some(targets) = name_to_defs.get(&tref.to_lowercase()) {
96                    for t in targets {
97                        if t != &f.id {
98                            add_edge(&mut edges, &f.id, t, w, reverse_factor);
99                        }
100                    }
101                }
102            }
103        }
104        edges
105    }
106
107    fn discover_related_files(
108        &self,
109        changed: &[PathBuf],
110        candidates: &[PathBuf],
111        repo_root: Option<&Path>,
112        file_cache: Option<&FxHashMap<PathBuf, String>>,
113    ) -> Vec<PathBuf> {
114        let proto_changed: Vec<&PathBuf> = changed.iter().filter(|f| is_proto_file(f)).collect();
115        if proto_changed.is_empty() {
116            return vec![];
117        }
118        let mut refs = FxHashSet::default();
119        for f in &proto_changed {
120            if let Some(content) = base::read_file_cached(f, file_cache) {
121                refs.extend(extract_imports(&content));
122            }
123        }
124        discover_files_by_refs(&refs, changed, candidates, repo_root)
125    }
126}
_diffctx/edges/semantic/protobuf.rs

_diffctx/edges/semantic/
protobuf.rs
