thread_ast_engine/replacer/
template.rs

1// SPDX-FileCopyrightText: 2022 Herrington Darkholme <2883231+HerringtonDarkholme@users.noreply.github.com>
2// SPDX-FileCopyrightText: 2025 Knitli Inc. <knitli@knit.li>
3// SPDX-FileContributor: Adam Poulemanos <adam@knit.li>
4//
5// SPDX-License-Identifier: AGPL-3.0-or-later AND MIT
6
7use super::indent::{DeindentedExtract, extract_with_deindent, get_indent_at_offset, indent_lines};
8use super::{MetaVarExtract, Replacer, split_first_meta_var};
9use crate::NodeMatch;
10use crate::language::Language;
11use crate::meta_var::{MetaVarEnv, Underlying};
12use crate::source::{Content, Doc};
13
14use thiserror::Error;
15
16use std::borrow::Cow;
17use thread_utilities::{RapidSet, get_set};
18
19#[derive(Debug, Clone)]
20pub enum TemplateFix {
21    // no meta_var, pure text
22    Textual(String),
23    WithMetaVar(Template),
24}
25
26#[derive(Debug, Error)]
27pub enum TemplateFixError {}
28
29impl TemplateFix {
30    pub fn try_new<L: Language>(template: &str, lang: &L) -> Result<Self, TemplateFixError> {
31        Ok(create_template(template, lang.meta_var_char(), &[]))
32    }
33
34    pub fn with_transform<L: Language>(
35        tpl: &str,
36        lang: &L,
37        trans: &[crate::meta_var::MetaVariableID],
38    ) -> Self {
39        create_template(tpl, lang.meta_var_char(), trans)
40    }
41
42    #[must_use]
43    pub fn used_vars(&self) -> RapidSet<&str> {
44        let template = match self {
45            Self::WithMetaVar(t) => t,
46            Self::Textual(_) => return get_set(),
47        };
48        template.vars.iter().map(|v| v.0.used_var()).collect()
49    }
50}
51
52impl<D: Doc> Replacer<D> for TemplateFix {
53    fn generate_replacement(&self, nm: &NodeMatch<'_, D>) -> Underlying<D> {
54        let leading = nm.get_doc().get_source().get_range(0..nm.range().start);
55        let indent = get_indent_at_offset::<D::Source>(leading);
56        let bytes = replace_fixer(self, nm.get_env());
57        let replaced = DeindentedExtract::MultiLine(&bytes, 0);
58        indent_lines::<D::Source>(indent, &replaced).to_vec()
59    }
60}
61
62type Indent = usize;
63
64#[derive(Debug, Clone)]
65pub struct Template {
66    fragments: Vec<String>,
67    vars: Vec<(MetaVarExtract, Indent)>,
68}
69
70fn create_template(
71    tmpl: &str,
72    mv_char: char,
73    transforms: &[crate::meta_var::MetaVariableID],
74) -> TemplateFix {
75    let mut fragments = vec![];
76    let mut vars = vec![];
77    let mut offset = 0;
78    let mut len = 0;
79    while let Some(i) = tmpl[len + offset..].find(mv_char) {
80        if let Some((meta_var, skipped)) =
81            split_first_meta_var(&tmpl[len + offset + i..], mv_char, transforms)
82        {
83            fragments.push(tmpl[len..len + offset + i].to_string());
84            // NB we have to count ident of the full string
85            let indent = get_indent_at_offset::<String>(&tmpl.as_bytes()[..len + offset + i]);
86            vars.push((meta_var, indent));
87            len += skipped + offset + i;
88            offset = 0;
89            continue;
90        }
91        debug_assert!(len + offset + i < tmpl.len());
92        // offset = 0, i = 0,
93        // 0 1 2
94        // $ a $
95        offset = offset + i + 1;
96    }
97    if fragments.is_empty() {
98        TemplateFix::Textual(tmpl[len..].to_string())
99    } else {
100        fragments.push(tmpl[len..].to_string());
101        TemplateFix::WithMetaVar(Template { fragments, vars })
102    }
103}
104
105fn replace_fixer<D: Doc>(fixer: &TemplateFix, env: &MetaVarEnv<'_, D>) -> Underlying<D> {
106    let template = match fixer {
107        TemplateFix::Textual(n) => return D::Source::decode_str(n).to_vec(),
108        TemplateFix::WithMetaVar(t) => t,
109    };
110    let mut ret = vec![];
111    let mut frags = template.fragments.iter();
112    let vars = template.vars.iter();
113    if let Some(frag) = frags.next() {
114        ret.extend_from_slice(&D::Source::decode_str(frag));
115    }
116    for ((var, indent), frag) in vars.zip(frags) {
117        if let Some(bytes) = maybe_get_var(env, var, indent.to_owned()) {
118            ret.extend_from_slice(&bytes);
119        }
120        ret.extend_from_slice(&D::Source::decode_str(frag));
121    }
122    ret
123}
124
125fn maybe_get_var<'e, 't, C, D>(
126    env: &'e MetaVarEnv<'t, D>,
127    var: &MetaVarExtract,
128    indent: usize,
129) -> Option<Cow<'e, [C::Underlying]>>
130where
131    C: Content + 'e,
132    D: Doc<Source = C>,
133{
134    let (source, range) = match var {
135        MetaVarExtract::Transformed(name) => {
136            // transformed source does not have range, directly return bytes
137            let source = env.get_transformed(name)?;
138            let de_intended = DeindentedExtract::MultiLine(source, 0);
139            let bytes = indent_lines::<D::Source>(indent, &de_intended);
140            return Some(Cow::Owned(bytes.into()));
141        }
142        MetaVarExtract::Single(name) => {
143            let replaced = env.get_match(name)?;
144            let source = replaced.get_doc().get_source();
145            let range = replaced.range();
146            (source, range)
147        }
148        MetaVarExtract::Multiple(name) => {
149            let nodes = env.get_multiple_matches(name);
150            if nodes.is_empty() {
151                return None;
152            }
153            // NOTE: start_byte is not always index range of source's slice.
154            // e.g. start_byte is still byte_offset in utf_16 (napi). start_byte
155            // so we need to call source's get_range method
156            let start = nodes[0].range().start;
157            let end = nodes[nodes.len() - 1].range().end;
158            let source = nodes[0].get_doc().get_source();
159            (source, start..end)
160        }
161    };
162    let extracted = extract_with_deindent(source, range);
163    let bytes = indent_lines::<D::Source>(indent, &extracted);
164    Some(Cow::Owned(bytes.into()))
165}
166
167// replace meta_var in template string, e.g. "Hello $NAME" -> "Hello World"
168pub fn gen_replacement<D: Doc>(template: &str, nm: &NodeMatch<'_, D>) -> Underlying<D> {
169    let fixer = create_template(template, nm.lang().meta_var_char(), &[]);
170    fixer.generate_replacement(nm)
171}
172
173#[cfg(test)]
174mod test {
175
176    use super::*;
177    use crate::Pattern;
178    use crate::language::Tsx;
179    use crate::matcher::NodeMatch;
180    use crate::meta_var::{MetaVarEnv, MetaVariable};
181    use crate::tree_sitter::LanguageExt;
182    use std::sync::Arc;
183    use thread_utilities::RapidMap;
184
185    #[test]
186    fn test_example() {
187        let src = r"
188if (true) {
189  a(
190    1
191      + 2
192      + 3
193  )
194}";
195        let pattern = "a($B)";
196        let template = r"c(
197  $B
198)";
199        let mut src = Tsx.ast_grep(src);
200        let pattern = Pattern::new(pattern, &Tsx);
201        let success = src.replace(pattern, template).expect("should replace");
202        assert!(success);
203        let expect = r"if (true) {
204  c(
205    1
206      + 2
207      + 3
208  )
209}";
210        assert_eq!(src.root().text(), expect);
211    }
212
213    fn test_str_replace(replacer: &str, vars: &[(&str, &str)], expected: &str) {
214        let mut env = MetaVarEnv::new();
215        let roots: Vec<_> = vars.iter().map(|(v, p)| (v, Tsx.ast_grep(p))).collect();
216        for (var, root) in &roots {
217            env.insert(var, root.root());
218        }
219        let dummy = Tsx.ast_grep("dummy");
220        let node_match = NodeMatch::new(dummy.root(), env.clone());
221        let replaced = replacer.generate_replacement(&node_match);
222        let replaced = String::from_utf8_lossy(&replaced);
223        assert_eq!(
224            replaced,
225            expected,
226            "wrong replacement {replaced} {expected} {:?}",
227            RapidMap::from(env)
228        );
229    }
230
231    #[test]
232    fn test_no_env() {
233        test_str_replace("let a = 123", &[], "let a = 123");
234        test_str_replace(
235            "console.log('hello world'); let b = 123;",
236            &[],
237            "console.log('hello world'); let b = 123;",
238        );
239    }
240
241    #[test]
242    fn test_single_env() {
243        test_str_replace("let a = $A", &[("A", "123")], "let a = 123");
244        test_str_replace(
245            "console.log($HW); let b = 123;",
246            &[("HW", "'hello world'")],
247            "console.log('hello world'); let b = 123;",
248        );
249    }
250
251    #[test]
252    fn test_multiple_env() {
253        test_str_replace("let $V = $A", &[("A", "123"), ("V", "a")], "let a = 123");
254        test_str_replace(
255            "console.log($HW); let $B = 123;",
256            &[("HW", "'hello world'"), ("B", "b")],
257            "console.log('hello world'); let b = 123;",
258        );
259    }
260
261    #[test]
262    fn test_multiple_occurrences() {
263        test_str_replace("let $A = $A", &[("A", "a")], "let a = a");
264        test_str_replace("var $A = () => $A", &[("A", "a")], "var a = () => a");
265        test_str_replace(
266            "const $A = () => { console.log($B); $A(); };",
267            &[("B", "'hello world'"), ("A", "a")],
268            "const a = () => { console.log('hello world'); a(); };",
269        );
270    }
271
272    fn test_ellipsis_replace(replacer: &str, vars: &[(&str, &str)], expected: &str) {
273        let mut env = MetaVarEnv::new();
274        let roots: Vec<_> = vars.iter().map(|(v, p)| (v, Tsx.ast_grep(p))).collect();
275        for (var, root) in &roots {
276            env.insert_multi(var, root.root().children().collect());
277        }
278        let dummy = Tsx.ast_grep("dummy");
279        let node_match = NodeMatch::new(dummy.root(), env.clone());
280        let replaced = replacer.generate_replacement(&node_match);
281        let replaced = String::from_utf8_lossy(&replaced);
282        assert_eq!(
283            replaced,
284            expected,
285            "wrong replacement {replaced} {expected} {:?}",
286            RapidMap::from(env)
287        );
288    }
289
290    #[test]
291    fn test_ellipsis_meta_var() {
292        test_ellipsis_replace(
293            "let a = () => { $$$B }",
294            &[("B", "alert('works!')")],
295            "let a = () => { alert('works!') }",
296        );
297        test_ellipsis_replace(
298            "let a = () => { $$$B }",
299            &[("B", "alert('works!');console.log(123)")],
300            "let a = () => { alert('works!');console.log(123) }",
301        );
302    }
303
304    #[test]
305    fn test_multi_ellipsis() {
306        test_ellipsis_replace(
307            "import {$$$A, B, $$$C} from 'a'",
308            &[("A", "A"), ("C", "C")],
309            "import {A, B, C} from 'a'",
310        );
311    }
312
313    #[test]
314    fn test_replace_in_string() {
315        test_str_replace("'$A'", &[("A", "123")], "'123'");
316    }
317
318    fn test_template_replace(template: &str, vars: &[(&str, &str)], expected: &str) {
319        let mut env = MetaVarEnv::new();
320        let roots: Vec<_> = vars.iter().map(|(v, p)| (v, Tsx.ast_grep(p))).collect();
321        for (var, root) in &roots {
322            env.insert(var, root.root());
323        }
324        let dummy = Tsx.ast_grep("dummy");
325        let node_match = NodeMatch::new(dummy.root(), env.clone());
326        let bytes = template.generate_replacement(&node_match);
327        let ret = String::from_utf8(bytes).expect("replacement must be valid utf-8");
328        assert_eq!(expected, ret);
329    }
330
331    #[test]
332    fn test_template() {
333        test_template_replace("Hello $A", &[("A", "World")], "Hello World");
334        test_template_replace("$B $A", &[("A", "World"), ("B", "Hello")], "Hello World");
335    }
336
337    #[test]
338    fn test_template_vars() {
339        let tf = TemplateFix::try_new("$A $B $C", &Tsx).expect("ok");
340        assert_eq!(tf.used_vars(), ["A", "B", "C"].into_iter().collect());
341        let tf = TemplateFix::try_new("$a$B$C", &Tsx).expect("ok");
342        assert_eq!(tf.used_vars(), ["B", "C"].into_iter().collect());
343        let tf = TemplateFix::try_new("$a$B$C", &Tsx).expect("ok");
344        assert_eq!(tf.used_vars(), ["B", "C"].into_iter().collect());
345    }
346
347    // GH #641
348    #[test]
349    fn test_multi_row_replace() {
350        test_template_replace(
351            "$A = $B",
352            &[("A", "x"), ("B", "[\n  1\n]")],
353            "x = [\n  1\n]",
354        );
355    }
356
357    #[test]
358    fn test_replace_rewriter() {
359        let tf = TemplateFix::with_transform("if (a)\n  $A", &Tsx, &[Arc::from("A")]);
360        let mut env = MetaVarEnv::new();
361        env.insert_transformation(
362            &MetaVariable::Multiple,
363            "A",
364            "if (b)\n  foo".bytes().collect(),
365        );
366        let dummy = Tsx.ast_grep("dummy");
367        let node_match = NodeMatch::new(dummy.root(), env.clone());
368        let bytes = tf.generate_replacement(&node_match);
369        let ret = String::from_utf8(bytes).expect("replacement must be valid utf-8");
370        assert_eq!("if (a)\n  if (b)\n    foo", ret);
371    }
372
373    #[test]
374    fn test_nested_matching_replace() {
375        // TODO we don't support nested replacement yet
376    }
377}
thread_ast_engine/replacer/template.rs

thread_ast_engine/replacer/
template.rs
