1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

pub mod language;
pub mod matcher;
pub mod meta_var;
pub mod ops;
pub mod source;
pub mod traversal;

mod match_tree;
mod node;
mod replacer;
mod ts_parser;

pub use language::Language;
pub use matcher::{Matcher, NodeMatch, Pattern, PatternError};
pub use node::Node;
pub use replacer::replace_meta_var_in_string;

use crate::replacer::Replacer;
use node::Root;
use source::Content;
use ts_parser::{Edit, TSParseError};

#[derive(Clone)]
pub struct AstGrep<L: Language> {
  inner: Root<L>,
}

impl<L: Language> AstGrep<L> {
  pub fn new<S: AsRef<str>>(src: S, lang: L) -> Self {
    Self {
      inner: Root::new(src.as_ref(), lang),
    }
  }

  pub fn customized<C: Content>(content: C, lang: L) -> Result<Self, TSParseError> {
    Ok(Self {
      inner: Root::customized(content, lang)?,
    })
  }

  pub fn source(&self) -> &str {
    &self.inner.source
  }

  pub fn root(&self) -> Node<L> {
    self.inner.root()
  }

  pub fn edit(&mut self, edit: Edit) -> Result<&mut Self, TSParseError> {
    self.inner.do_edit(edit)?;
    Ok(self)
  }

  pub fn replace<M: Matcher<L>, R: Replacer<L>>(
    &mut self,
    pattern: M,
    replacer: R,
  ) -> Result<bool, TSParseError> {
    if let Some(edit) = self.root().replace(pattern, replacer) {
      self.edit(edit)?;
      Ok(true)
    } else {
      Ok(false)
    }
  }

  pub fn lang(&self) -> &L {
    &self.inner.lang
  }

  pub fn generate(self) -> String {
    self.inner.source.to_string()
  }
}

#[cfg(test)]
mod test {
  use super::*;
  use language::Tsx;
  use ops::Op;

  #[test]
  fn test_replace() {
    let mut ast_grep = Tsx.ast_grep("var a = 1; let b = 2;");
    ast_grep.replace("var $A = $B", "let $A = $B").unwrap();
    let source = ast_grep.generate();
    assert_eq!(source, "let a = 1; let b = 2;"); // note the semicolon
  }

  #[test]
  fn test_replace_by_rule() {
    let rule = Op::either("let a = 123").or("let b = 456");
    let mut ast_grep = Tsx.ast_grep("let a = 123");
    let replaced = ast_grep.replace(rule, "console.log('it works!')").unwrap();
    assert!(replaced);
    let source = ast_grep.generate();
    assert_eq!(source, "console.log('it works!')");
  }

  #[test]
  fn test_replace_unnamed_node() {
    // ++ and -- is unnamed node in tree-sitter javascript
    let mut ast_grep = Tsx.ast_grep("c++");
    ast_grep.replace("$A++", "$A--").unwrap();
    let source = ast_grep.generate();
    assert_eq!(source, "c--");
  }

  #[test]
  fn test_replace_trivia() {
    let mut ast_grep = Tsx.ast_grep("var a = 1 /*haha*/;");
    ast_grep.replace("var $A = $B", "let $A = $B").unwrap();
    let source = ast_grep.generate();
    assert_eq!(source, "let a = 1 /*haha*/;"); // semicolon

    let mut ast_grep = Tsx.ast_grep("var a = 1; /*haha*/");
    ast_grep.replace("var $A = $B", "let $A = $B").unwrap();
    let source = ast_grep.generate();
    assert_eq!(source, "let a = 1; /*haha*/");
  }
}