ast_grep_core/
ops.rs

1use crate::matcher::{MatchAll, MatchNone, Matcher};
2use crate::meta_var::MetaVarEnv;
3use crate::{Doc, Node};
4use bit_set::BitSet;
5use std::borrow::Cow;
6
7pub struct And<P1: Matcher, P2: Matcher> {
8  pattern1: P1,
9  pattern2: P2,
10}
11
12impl<P1, P2> Matcher for And<P1, P2>
13where
14  P1: Matcher,
15  P2: Matcher,
16{
17  fn match_node_with_env<'tree, D: Doc>(
18    &self,
19    node: Node<'tree, D>,
20    env: &mut Cow<MetaVarEnv<'tree, D>>,
21  ) -> Option<Node<'tree, D>> {
22    // keep the original env intact until both arms match
23    let mut new_env = Cow::Borrowed(env.as_ref());
24    let node = self.pattern1.match_node_with_env(node, &mut new_env)?;
25    let ret = self.pattern2.match_node_with_env(node, &mut new_env)?;
26    // both succeed – commit the combined env
27    *env = Cow::Owned(new_env.into_owned());
28    Some(ret)
29  }
30
31  fn potential_kinds(&self) -> Option<BitSet> {
32    let set1 = self.pattern1.potential_kinds();
33    let set2 = self.pattern2.potential_kinds();
34    // if both constituent have Some(bitset), intersect them
35    // otherwise returns either of the non-null set
36    match (&set1, &set2) {
37      (Some(s1), Some(s2)) => Some(s1.intersection(s2).collect()),
38      _ => set1.xor(set2),
39    }
40  }
41}
42
43// we pre-compute and cache potential_kinds. So patterns should not be mutated.
44// Box<[P]> is used here for immutability so that kinds will never be invalidated.
45pub struct All<P: Matcher> {
46  patterns: Box<[P]>,
47  kinds: Option<BitSet>,
48}
49
50impl<P: Matcher> All<P> {
51  pub fn new<PS: IntoIterator<Item = P>>(patterns: PS) -> Self {
52    let patterns: Box<[P]> = patterns.into_iter().collect();
53    let kinds = Self::compute_kinds(&patterns);
54    Self { patterns, kinds }
55  }
56
57  fn compute_kinds(patterns: &[P]) -> Option<BitSet> {
58    let mut set: Option<BitSet> = None;
59    for pattern in patterns {
60      let Some(n) = pattern.potential_kinds() else {
61        continue;
62      };
63      if let Some(set) = set.as_mut() {
64        set.intersect_with(&n);
65      } else {
66        set = Some(n);
67      }
68    }
69    set
70  }
71
72  pub fn inner(&self) -> &[P] {
73    &self.patterns
74  }
75}
76
77impl<P: Matcher> Matcher for All<P> {
78  fn match_node_with_env<'tree, D: Doc>(
79    &self,
80    node: Node<'tree, D>,
81    env: &mut Cow<MetaVarEnv<'tree, D>>,
82  ) -> Option<Node<'tree, D>> {
83    if let Some(kinds) = &self.kinds {
84      if !kinds.contains(node.kind_id().into()) {
85        return None;
86      }
87    }
88    let mut new_env = Cow::Borrowed(env.as_ref());
89    let all_satisfied = self
90      .patterns
91      .iter()
92      .all(|p| p.match_node_with_env(node.clone(), &mut new_env).is_some());
93    if all_satisfied {
94      *env = Cow::Owned(new_env.into_owned());
95      Some(node)
96    } else {
97      None
98    }
99  }
100
101  fn potential_kinds(&self) -> Option<BitSet> {
102    self.kinds.clone()
103  }
104}
105
106// Box<[P]> for immutability and potential_kinds cache correctness
107pub struct Any<P> {
108  patterns: Box<[P]>,
109  kinds: Option<BitSet>,
110}
111
112impl<P: Matcher> Any<P> {
113  pub fn new<PS: IntoIterator<Item = P>>(patterns: PS) -> Self {
114    let patterns: Box<[P]> = patterns.into_iter().collect();
115    let kinds = Self::compute_kinds(&patterns);
116    Self { patterns, kinds }
117  }
118
119  fn compute_kinds(patterns: &[P]) -> Option<BitSet> {
120    let mut set = BitSet::new();
121    for pattern in patterns {
122      let n = pattern.potential_kinds()?;
123      set.union_with(&n);
124    }
125    Some(set)
126  }
127
128  pub fn inner(&self) -> &[P] {
129    &self.patterns
130  }
131}
132
133impl<M: Matcher> Matcher for Any<M> {
134  fn match_node_with_env<'tree, D: Doc>(
135    &self,
136    node: Node<'tree, D>,
137    env: &mut Cow<MetaVarEnv<'tree, D>>,
138  ) -> Option<Node<'tree, D>> {
139    if let Some(kinds) = &self.kinds {
140      if !kinds.contains(node.kind_id().into()) {
141        return None;
142      }
143    }
144    let mut new_env = Cow::Borrowed(env.as_ref());
145    let found = self.patterns.iter().find_map(|p| {
146      new_env = Cow::Borrowed(env.as_ref());
147      p.match_node_with_env(node.clone(), &mut new_env)
148    });
149    if found.is_some() {
150      *env = Cow::Owned(new_env.into_owned());
151      Some(node)
152    } else {
153      None
154    }
155  }
156
157  fn potential_kinds(&self) -> Option<BitSet> {
158    self.kinds.clone()
159  }
160}
161
162pub struct Or<P1: Matcher, P2: Matcher> {
163  pattern1: P1,
164  pattern2: P2,
165}
166
167impl<P1, P2> Matcher for Or<P1, P2>
168where
169  P1: Matcher,
170  P2: Matcher,
171{
172  fn match_node_with_env<'tree, D: Doc>(
173    &self,
174    node: Node<'tree, D>,
175    env: &mut Cow<MetaVarEnv<'tree, D>>,
176  ) -> Option<Node<'tree, D>> {
177    let mut new_env = Cow::Borrowed(env.as_ref());
178    if let Some(ret) = self
179      .pattern1
180      .match_node_with_env(node.clone(), &mut new_env)
181    {
182      *env = Cow::Owned(new_env.into_owned());
183      Some(ret)
184    } else {
185      self.pattern2.match_node_with_env(node, env)
186    }
187  }
188
189  fn potential_kinds(&self) -> Option<BitSet> {
190    let mut set1 = self.pattern1.potential_kinds()?;
191    let set2 = self.pattern2.potential_kinds()?;
192    set1.union_with(&set2);
193    Some(set1)
194  }
195}
196
197pub struct Not<M: Matcher> {
198  not: M,
199}
200
201impl<M: Matcher> Not<M> {
202  pub fn new(not: M) -> Self {
203    Self { not }
204  }
205
206  pub fn inner(&self) -> &M {
207    &self.not
208  }
209}
210impl<P> Matcher for Not<P>
211where
212  P: Matcher,
213{
214  fn match_node_with_env<'tree, D: Doc>(
215    &self,
216    node: Node<'tree, D>,
217    env: &mut Cow<MetaVarEnv<'tree, D>>,
218  ) -> Option<Node<'tree, D>> {
219    self
220      .not
221      .match_node_with_env(node.clone(), env)
222      .xor(Some(node))
223  }
224}
225
226#[derive(Clone)]
227pub struct Op<M: Matcher> {
228  inner: M,
229}
230
231impl<M> Matcher for Op<M>
232where
233  M: Matcher,
234{
235  fn match_node_with_env<'tree, D: Doc>(
236    &self,
237    node: Node<'tree, D>,
238    env: &mut Cow<MetaVarEnv<'tree, D>>,
239  ) -> Option<Node<'tree, D>> {
240    self.inner.match_node_with_env(node, env)
241  }
242
243  fn potential_kinds(&self) -> Option<BitSet> {
244    self.inner.potential_kinds()
245  }
246}
247
248/*
249pub struct Predicate<F> {
250  func: F,
251}
252
253impl<L, F> Matcher for Predicate<F>
254where
255  L: Language,
256  F: for<'tree> Fn(&Node<'tree, StrDoc<L>>) -> bool,
257{
258  fn match_node_with_env<'tree, D: Doc<Lang=L>>(
259    &self,
260    node: Node<'tree, D>,
261    env: &mut MetaVarEnv<'tree, D>,
262  ) -> Option<Node<'tree, D>> {
263    (self.func)(&node).then_some(node)
264  }
265}
266*/
267
268/*
269// we don't need specify M for static method
270impl<L: Language> Op<L, MatchNone> {
271  pub fn func<F>(func: F) -> Predicate<F>
272  where
273    F: for<'tree> Fn(&Node<'tree, StrDoc<L>>) -> bool,
274  {
275    Predicate { func }
276  }
277}
278*/
279
280impl<M: Matcher> Op<M> {
281  pub fn not(pattern: M) -> Not<M> {
282    Not { not: pattern }
283  }
284}
285
286impl<M: Matcher> Op<M> {
287  pub fn every(pattern: M) -> Op<And<M, MatchAll>> {
288    Op {
289      inner: And {
290        pattern1: pattern,
291        pattern2: MatchAll,
292      },
293    }
294  }
295  pub fn either(pattern: M) -> Op<Or<M, MatchNone>> {
296    Op {
297      inner: Or {
298        pattern1: pattern,
299        pattern2: MatchNone,
300      },
301    }
302  }
303
304  pub fn all<MS: IntoIterator<Item = M>>(patterns: MS) -> All<M> {
305    All::new(patterns)
306  }
307
308  pub fn any<MS: IntoIterator<Item = M>>(patterns: MS) -> Any<M> {
309    Any::new(patterns)
310  }
311
312  pub fn new(matcher: M) -> Op<M> {
313    Self { inner: matcher }
314  }
315}
316
317type NestedAnd<M, N, O> = And<And<M, N>, O>;
318impl<M: Matcher, N: Matcher> Op<And<M, N>> {
319  pub fn and<O: Matcher>(self, other: O) -> Op<NestedAnd<M, N, O>> {
320    Op {
321      inner: And {
322        pattern1: self.inner,
323        pattern2: other,
324      },
325    }
326  }
327}
328
329type NestedOr<M, N, O> = Or<Or<M, N>, O>;
330impl<M: Matcher, N: Matcher> Op<Or<M, N>> {
331  pub fn or<O: Matcher>(self, other: O) -> Op<NestedOr<M, N, O>> {
332    Op {
333      inner: Or {
334        pattern1: self.inner,
335        pattern2: other,
336      },
337    }
338  }
339}
340
341#[cfg(test)]
342mod test {
343  use super::*;
344  use crate::language::Tsx;
345  use crate::matcher::MatcherExt;
346  use crate::meta_var::MetaVarEnv;
347  use crate::Root;
348
349  fn test_find(matcher: &impl Matcher, code: &str) {
350    let node = Root::str(code, Tsx);
351    assert!(matcher.find_node(node.root()).is_some());
352  }
353  fn test_not_find(matcher: &impl Matcher, code: &str) {
354    let node = Root::str(code, Tsx);
355    assert!(matcher.find_node(node.root()).is_none());
356  }
357  fn find_all(matcher: impl Matcher, code: &str) -> Vec<String> {
358    let node = Root::str(code, Tsx);
359    node
360      .root()
361      .find_all(matcher)
362      .map(|n| n.text().to_string())
363      .collect()
364  }
365
366  #[test]
367  fn test_or() {
368    let matcher = Or {
369      pattern1: "let a = 1",
370      pattern2: "const b = 2",
371    };
372    test_find(&matcher, "let a = 1");
373    test_find(&matcher, "const b = 2");
374    test_not_find(&matcher, "let a = 2");
375    test_not_find(&matcher, "const a = 1");
376    test_not_find(&matcher, "let b = 2");
377    test_not_find(&matcher, "const b = 1");
378  }
379
380  #[test]
381  fn test_not() {
382    let matcher = Not { not: "let a = 1" };
383    test_find(&matcher, "const b = 2");
384  }
385
386  #[test]
387  fn test_and() {
388    let matcher = And {
389      pattern1: "let a = $_",
390      pattern2: Not { not: "let a = 123" },
391    };
392    test_find(&matcher, "let a = 233");
393    test_find(&matcher, "let a = 456");
394    test_not_find(&matcher, "let a = 123");
395  }
396
397  #[test]
398  fn test_api_and() {
399    let matcher = Op::every("let a = $_").and(Op::not("let a = 123"));
400    test_find(&matcher, "let a = 233");
401    test_find(&matcher, "let a = 456");
402    test_not_find(&matcher, "let a = 123");
403  }
404
405  #[test]
406  fn test_api_or() {
407    let matcher = Op::either("let a = 1").or("const b = 2");
408    test_find(&matcher, "let a = 1");
409    test_find(&matcher, "const b = 2");
410    test_not_find(&matcher, "let a = 2");
411    test_not_find(&matcher, "const a = 1");
412    test_not_find(&matcher, "let b = 2");
413    test_not_find(&matcher, "const b = 1");
414  }
415  #[test]
416  fn test_multiple_match() {
417    let sequential = find_all("$A + b", "let f = () => a + b; let ff = () => c + b");
418    assert_eq!(sequential.len(), 2);
419    let nested = find_all(
420      "function $A() { $$$ }",
421      "function a() { function b() { b } }",
422    );
423    assert_eq!(nested.len(), 2);
424  }
425
426  #[test]
427  fn test_multiple_match_order() {
428    let ret = find_all(
429      "$A + b",
430      "let f = () => () => () => a + b; let ff = () => c + b",
431    );
432    assert_eq!(ret, ["a + b", "c + b"], "should match source code order");
433  }
434
435  /*
436  #[test]
437  fn test_api_func() {
438    let matcher = Op::func(|n| n.text().contains("114514"));
439    test_find(&matcher, "let a = 114514");
440    test_not_find(&matcher, "let a = 1919810");
441  }
442  */
443  use crate::Pattern;
444  trait TsxMatcher {
445    fn t(self) -> Pattern;
446  }
447  impl TsxMatcher for &str {
448    fn t(self) -> Pattern {
449      Pattern::new(self, Tsx)
450    }
451  }
452
453  #[test]
454  fn test_and_kinds() {
455    // intersect None kinds
456    let matcher = Op::every("let a = $_".t()).and(Op::not("let a = 123".t()));
457    assert_eq!(matcher.potential_kinds().map(|v| v.len()), Some(1));
458    let matcher = Op::every(Op::not("let a = $_".t())).and("let a = 123".t());
459    assert_eq!(matcher.potential_kinds().map(|v| v.len()), Some(1));
460    // intersect Same kinds
461    let matcher = Op::every("let a = $_".t()).and("let b = 123".t());
462    assert_eq!(matcher.potential_kinds().map(|v| v.len()), Some(1));
463    // intersect different kinds
464    let matcher = Op::every("let a = 1".t()).and("console.log(1)".t());
465    assert_eq!(matcher.potential_kinds().map(|v| v.len()), Some(0));
466    // two None kinds
467    let matcher = Op::every(Op::not("let a = $_".t())).and(Op::not("let a = 123".t()));
468    assert_eq!(matcher.potential_kinds(), None);
469  }
470
471  #[test]
472  fn test_or_kinds() {
473    // union None kinds
474    let matcher = Op::either("let a = $_".t()).or(Op::not("let a = 123".t()));
475    assert_eq!(matcher.potential_kinds(), None);
476    let matcher = Op::either(Op::not("let a = $_".t())).or("let a = 123".t());
477    assert_eq!(matcher.potential_kinds(), None);
478    // union Same kinds
479    let matcher = Op::either("let a = $_".t()).or("let b = 123".t());
480    assert_eq!(matcher.potential_kinds().map(|v| v.len()), Some(1));
481    // union different kinds
482    let matcher = Op::either("let a = 1".t()).or("console.log(1)".t());
483    assert_eq!(matcher.potential_kinds().map(|v| v.len()), Some(2));
484    // two None kinds
485    let matcher = Op::either(Op::not("let a = $_".t())).or(Op::not("let a = 123".t()));
486    assert_eq!(matcher.potential_kinds(), None);
487  }
488
489  #[test]
490  fn test_all_kinds() {
491    // intersect None kinds
492    let matcher = Op::all(["let a = $_".t(), "$A".t()]);
493    assert_eq!(matcher.potential_kinds().map(|v| v.len()), Some(1));
494    let matcher = Op::all(["$A".t(), "let a = $_".t()]);
495    assert_eq!(matcher.potential_kinds().map(|v| v.len()), Some(1));
496    // intersect Same kinds
497    let matcher = Op::all(["let a = $_".t(), "let b = 123".t()]);
498    assert_eq!(matcher.potential_kinds().map(|v| v.len()), Some(1));
499    // intersect different kinds
500    let matcher = Op::all(["let a = 1".t(), "console.log(1)".t()]);
501    assert_eq!(matcher.potential_kinds().map(|v| v.len()), Some(0));
502    // two None kinds
503    let matcher = Op::all(["$A".t(), "$B".t()]);
504    assert_eq!(matcher.potential_kinds(), None);
505  }
506
507  #[test]
508  fn test_any_kinds() {
509    // union None kinds
510    let matcher = Op::any(["let a = $_".t(), "$A".t()]);
511    assert_eq!(matcher.potential_kinds(), None);
512    let matcher = Op::any(["$A".t(), "let a = $_".t()]);
513    assert_eq!(matcher.potential_kinds(), None);
514    // union Same kinds
515    let matcher = Op::any(["let a = $_".t(), "let b = 123".t()]);
516    assert_eq!(matcher.potential_kinds().map(|v| v.len()), Some(1));
517    // union different kinds
518    let matcher = Op::any(["let a = 1".t(), "console.log(1)".t()]);
519    assert_eq!(matcher.potential_kinds().map(|v| v.len()), Some(2));
520    // two None kinds
521    let matcher = Op::any(["$A".t(), "$B".t()]);
522    assert_eq!(matcher.potential_kinds(), None);
523  }
524
525  #[test]
526  fn test_or_revert_env() {
527    let matcher = Op::either(Op::every("foo($A)".t()).and("impossible".t())).or("foo($B)".t());
528    let code = Root::str("foo(123)", Tsx);
529    let matches = code.root().find(matcher).expect("should found");
530    assert!(matches.get_env().get_match("A").is_none());
531    assert_eq!(matches.get_env().get_match("B").unwrap().text(), "123");
532  }
533
534  #[test]
535  fn test_any_revert_env() {
536    let matcher = Op::any([
537      Op::all(["foo($A)".t(), "impossible".t()]),
538      Op::all(["foo($B)".t()]),
539    ]);
540    let code = Root::str("foo(123)", Tsx);
541    let matches = code.root().find(matcher).expect("should found");
542    assert!(matches.get_env().get_match("A").is_none());
543    assert_eq!(matches.get_env().get_match("B").unwrap().text(), "123");
544  }
545
546  // gh #1225
547  #[test]
548  fn test_all_revert_env() {
549    let matcher = Op::all(["$A(123)".t(), "$B(456)".t()]);
550    let code = Root::str("foo(123)", Tsx);
551    let node = code.root().find("foo($C)").expect("should exist");
552    let node = node.get_node().clone();
553    let mut env = Cow::Owned(MetaVarEnv::new());
554    assert!(matcher.match_node_with_env(node, &mut env).is_none());
555    assert!(env.get_match("A").is_none());
556  }
557}
ast_grep_core/ops.rs

ast_grep_core/
ops.rs
