Skip to main content

_diffctx/edges/semantic/
r_lang.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_r_file(path: &Path) -> bool {
14    let ext = base::file_ext(path);
15    matches!(ext.as_str(), ".r" | ".rmd")
16}
17
18static SOURCE_RE: Lazy<Regex> =
19    Lazy::new(|| Regex::new(r##"(?m)source\s*\(\s*['"]([^'"]+)['"]"##).unwrap());
20static FUNC_DEF_RE: Lazy<Regex> =
21    Lazy::new(|| Regex::new(r"(?m)^\s*(\w+)\s*<-\s*function\s*\(").unwrap());
22static S4_CLASS_RE: Lazy<Regex> =
23    Lazy::new(|| Regex::new(r##"(?m)setClass\s*\(\s*['"](\w+)['"]"##).unwrap());
24static S4_METHOD_RE: Lazy<Regex> =
25    Lazy::new(|| Regex::new(r##"(?m)setMethod\s*\(\s*['"](\w+)['"]"##).unwrap());
26static CALL_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\b([a-zA-Z_.]\w*)\s*\(").unwrap());
27
28static R_KEYWORDS: Lazy<FxHashSet<&str>> = Lazy::new(|| {
29    [
30        "if",
31        "else",
32        "for",
33        "while",
34        "repeat",
35        "function",
36        "return",
37        "next",
38        "break",
39        "in",
40        "TRUE",
41        "FALSE",
42        "NULL",
43        "NA",
44        "Inf",
45        "NaN",
46        "library",
47        "require",
48        "source",
49        "print",
50        "cat",
51        "paste",
52        "c",
53        "list",
54        "data.frame",
55        "matrix",
56        "length",
57        "nrow",
58        "ncol",
59        "which",
60        "apply",
61        "sapply",
62        "lapply",
63    ]
64    .iter()
65    .copied()
66    .collect()
67});
68
69fn extract_sources(content: &str) -> FxHashSet<String> {
70    SOURCE_RE
71        .captures_iter(content)
72        .map(|c| c[1].to_string())
73        .collect()
74}
75
76fn extract_defs(content: &str) -> FxHashSet<String> {
77    let mut defs: FxHashSet<String> = FUNC_DEF_RE
78        .captures_iter(content)
79        .map(|c| c[1].to_string())
80        .collect();
81    defs.extend(S4_CLASS_RE.captures_iter(content).map(|c| c[1].to_string()));
82    defs.extend(
83        S4_METHOD_RE
84            .captures_iter(content)
85            .map(|c| c[1].to_string()),
86    );
87    defs
88}
89
90fn extract_calls(content: &str) -> FxHashSet<String> {
91    CALL_RE
92        .captures_iter(content)
93        .map(|c| c[1].to_string())
94        .filter(|n| !R_KEYWORDS.contains(n.as_str()))
95        .collect()
96}
97
98pub struct RLangEdgeBuilder;
99
100impl EdgeBuilder for RLangEdgeBuilder {
101    fn build(&self, fragments: &[Fragment], repo_root: Option<&Path>) -> EdgeDict {
102        let frags: Vec<&Fragment> = fragments
103            .iter()
104            .filter(|f| is_r_file(Path::new(f.path())))
105            .collect();
106        if frags.is_empty() {
107            return FxHashMap::default();
108        }
109
110        let source_w = EDGE_WEIGHTS["r_source"].forward;
111        let fn_w = EDGE_WEIGHTS["r_fn"].forward;
112        let s4_w = EDGE_WEIGHTS["r_s4"].forward;
113        let reverse_factor = EDGE_WEIGHTS["r_source"].reverse_factor;
114
115        let idx = base::FragmentIndex::new(fragments, repo_root);
116        let mut name_to_defs: FxHashMap<String, Vec<_>> = FxHashMap::default();
117        for f in &frags {
118            for name in extract_defs(&f.content) {
119                name_to_defs
120                    .entry(name.to_lowercase())
121                    .or_default()
122                    .push(f.id.clone());
123            }
124        }
125
126        let mut edges: EdgeDict = FxHashMap::default();
127
128        for f in &frags {
129            let self_defs = extract_defs(&f.content);
130            for src in extract_sources(&f.content) {
131                base::link_by_name(&f.id, &src, &idx, &mut edges, source_w, reverse_factor);
132            }
133            for call in extract_calls(&f.content) {
134                if self_defs.contains(&call) {
135                    continue;
136                }
137                let w = if S4_CLASS_RE.is_match(&f.content) || S4_METHOD_RE.is_match(&f.content) {
138                    s4_w
139                } else {
140                    fn_w
141                };
142                if let Some(targets) = name_to_defs.get(&call.to_lowercase()) {
143                    for t in targets {
144                        if t != &f.id {
145                            add_edge(&mut edges, &f.id, t, w, reverse_factor);
146                        }
147                    }
148                }
149            }
150        }
151        edges
152    }
153
154    fn discover_related_files(
155        &self,
156        changed: &[PathBuf],
157        candidates: &[PathBuf],
158        repo_root: Option<&Path>,
159        file_cache: Option<&FxHashMap<PathBuf, String>>,
160    ) -> Vec<PathBuf> {
161        let r_changed: Vec<&PathBuf> = changed.iter().filter(|f| is_r_file(f)).collect();
162        if r_changed.is_empty() {
163            return vec![];
164        }
165        let mut refs = FxHashSet::default();
166        for f in &r_changed {
167            if let Some(content) = base::read_file_cached(f, file_cache) {
168                refs.extend(extract_sources(&content));
169            }
170        }
171        discover_files_by_refs(&refs, changed, candidates, repo_root)
172    }
173}