ast_grep_core/replacer/
template.rs

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

ast_grep_core/replacer/
template.rs
