WIP code

ghstack-source-id: 686c0e7911e549984aaeb4a7d9fdb38b7a56cda8
Pull Request resolved: #536

Cargo.toml

@@ -67,6 +67,7 @@ getset = "0.1.2"
 pyo3 = "0.19.0"
 pyo3-log = "0.8.1"
 glob = "0.3.1"
+ast-grep-core = "0.7.2"
 
 [features]
 extension-module = ["pyo3/extension-module"]

src/models/capture_group_patterns.rs

@@ -14,18 +14,23 @@ Copyright (c) 2023 Uber Technologies, Inc.
 use crate::{
  models::Validator,
  utilities::{
+ ast_grep_utilities::get_all_matches_for_ast_grep_pattern,
  regex_utilities::get_all_matches_for_regex,
  tree_sitter_utilities::{get_all_matches_for_query, get_ts_query_parser, number_of_errors},
  Instantiate,
  },
 };
+use ast_grep_core::{language::TSLanguage, Pattern, StrDoc};
 use pyo3::prelude::pyclass;
 use regex::Regex;
 use serde_derive::Deserialize;
 use std::collections::HashMap;
 use tree_sitter::{Node, Query};
 
-use super::{default_configs::REGEX_QUERY_PREFIX, matches::Match};
+use super::{
+ default_configs::{AST_GREP_PREFIX, REGEX_QUERY_PREFIX},
+ matches::Match,
+};
 
 #[pyclass]
 #[derive(Deserialize, Debug, Clone, Default, PartialEq, Hash, Eq)]
@@ -54,6 +59,9 @@ impl Validator for CGPattern {
  .map(|_| Ok(()))
  .unwrap_or(Err(format!("Cannot parse the regex - {}", self.pattern())));
  }
+ if self.pattern().starts_with(AST_GREP_PREFIX) {
+ return Ok(());
+ }
  let mut parser = get_ts_query_parser();
  parser
  .parse(self.pattern(), None)
@@ -78,6 +86,7 @@ impl Instantiate for CGPattern {
 
 #[derive(Debug)]
 pub(crate) enum CompiledCGPattern {
+ P(Pattern<StrDoc<TSLanguage>>),
  Q(Query),
  R(Regex),
 }
@@ -104,6 +113,7 @@ impl CompiledCGPattern {
  replace_node_idx: Option<u8>,
  ) -> Vec<Match> {
  match self {
+ CompiledCGPattern::P(_pattern) => panic!("ast-grep pattern is not supported"),
  CompiledCGPattern::Q(query) => get_all_matches_for_query(
  node,
  source_code,
@@ -112,9 +122,7 @@ impl CompiledCGPattern {
  replace_node,
  replace_node_idx,
  ),
- CompiledCGPattern::R(regex) => {
- get_all_matches_for_regex(node, source_code, regex, recursive, replace_node)
- }
+ CompiledCGPattern::R(regex) => get_all_matches_for_ast_grep_pattern(),
  }
  }
 }

src/models/default_configs.rs

@@ -32,6 +32,7 @@ pub const STRINGS: &str = "strings";
 pub const TS_SCHEME: &str = "scm"; // We support scheme files that contain tree-sitter query
 
 pub const REGEX_QUERY_PREFIX: &str = "rgx ";
+pub const AST_GREP_PREFIX: &str = "sg ";
 
 #[cfg(test)]
 //FIXME: Remove this hack by not passing PiranhaArguments to SourceCodeUnit

src/models/matches.rs

@@ -13,6 +13,7 @@ Copyright (c) 2023 Uber Technologies, Inc.
 
 use std::collections::HashMap;
 
+use ast_grep_core::{language::TSLanguage, StrDoc};
 use getset::{Getters, MutGetters};
 use itertools::Itertools;
 use log::trace;
@@ -67,6 +68,19 @@ impl Match {
  }
  }
 
+ pub(crate) fn from_ast_grep_captures(
+ captures: &ast_grep_core::Node<'_, StrDoc<TSLanguage>>, matches: HashMap<String, String>,
+ source_code: &str,
+ ) -> Self {
+ Match {
+ matched_string: captures.text().to_string(),
+ range: Range::from_range(&captures.range(), source_code),
+ matches,
+ associated_comma: None,
+ associated_comments: Vec::new(),
+ }
+ }
+
  pub(crate) fn new(
  matched_string: String, range: tree_sitter::Range, matches: HashMap<String, String>,
  ) -> Self {
@@ -281,6 +295,15 @@ impl Range {
  end_point: position_for_offset(source_code.as_bytes(), mtch.end()),
  }
  }
+
+ pub(crate) fn from_range(range: &std::ops::Range<usize>, source_code: &str) -> Self {
+ Self {
+ start_byte: range.start,
+ end_byte: range.end,
+ start_point: position_for_offset(source_code.as_bytes(), range.start),
+ end_point: position_for_offset(source_code.as_bytes(), range.end),
+ }
+ }
 }
 
 // Finds the position (col and row number) for a given offset.

src/utilities/ast_grep_utilities.rs

@@ -0,0 +1,81 @@
+use std::{collections::HashMap, hash::Hash};
+
+/*
+Copyright (c) 2023 Uber Technologies, Inc.
+
+<p>Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+except in compliance with the License. You may obtain a copy of the License at
+<p>http://www.apache.org/licenses/LICENSE-2.0
+
+<p>Unless required by applicable law or agreed to in writing, software distributed under the
+License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+express or implied. See the License for the specific language governing permissions and
+limitations under the License.
+*/
+use ast_grep_core::{language::TSLanguage, AstGrep, Matcher, Pattern, StrDoc};
+use tree_sitter::Node;
+
+use crate::models::matches::Match;
+
+/// Applies the query upon the given `node`, and gets all the matches
+/// # Arguments
+/// * `node` - the root node to apply the query upon
+/// * `source_code` - the corresponding source code string for the node.
+/// * `recursive` - if `true` it matches the query to `self` and `self`'s sub-ASTs, else it matches the `query` only to `self`.
+/// * `replace_node` - node to replace
+///
+/// # Returns
+/// List containing all the matches against `node`
+pub(crate) fn get_all_matches_for_ast_grep_pattern(
+ node: &Node, source_code: String, pattern: &Pattern<StrDoc<TSLanguage>>, recursive: bool,
+ replace_node: Option<String>, language: tree_sitter::Language,
+) -> Vec<Match> {
+ let x = AstGrep::new(&source_code, TSLanguage::from(language));
+
+ let all_captures = x.root().find_all(pattern);
+ let mut all_matches = vec![];
+ for captures in all_captures {
+ let range_matches_node =
+ node.start_byte() == captures.range().start && node.end_byte() == captures.range().end;
+ let range_matches_inside_node =
+ node.start_byte() <= captures.range().start && node.end_byte() >= captures.range().end;
+ if (recursive && range_matches_inside_node) || range_matches_node {
+ let replace_node_match = if let Some(ref rn) = replace_node {
+ captures
+ .get_env()
+ .get_match(rn)
+ .unwrap_or_else(|| panic!("The tag {rn} provided in the replace node is not present"))
+ } else {
+ captures.get_node()
+ };
+ let matches = extract_captures(&captures);
+ all_matches.push(Match::from_ast_grep_captures(
+ &replace_node_match,
+ matches,
+ &source_code,
+ ));
+ }
+ }
+ vec![]
+}
+
+fn extract_captures(
+ captures: &ast_grep_core::NodeMatch<'_, StrDoc<TSLanguage>>,
+) -> HashMap<String, String> {
+ let mut map = HashMap::new();
+ for v in captures.get_env().get_matched_variables() {
+ let name = match v {
+ ast_grep_core::meta_var::MetaVariable::Named(name, _) => Some(name),
+ ast_grep_core::meta_var::MetaVariable::Anonymous(_) => None,
+ ast_grep_core::meta_var::MetaVariable::Ellipsis => None,
+ ast_grep_core::meta_var::MetaVariable::NamedEllipsis(name) => Some(name),
+ };
+ if let Some(n) = name {
+ map.insert(
+ n.to_string(),
+ captures.get_env().get_match(&n).unwrap().text().to_string(),
+ );
+ }
+ }
+ return map;
+}

src/utilities/mod.rs

@@ -11,6 +11,7 @@ Copyright (c) 2023 Uber Technologies, Inc.
  limitations under the License.
 */
 
+pub(crate) mod ast_grep_utilities;
 pub(crate) mod regex_utilities;
 pub(crate) mod tree_sitter_utilities;
 use std::collections::HashMap;

utilities/ast_grep_utilities.rs

@@ -0,0 +1,13 @@
+/*
+ Copyright (c) 2023 Uber Technologies, Inc.
+
+ <p>Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ except in compliance with the License. You may obtain a copy of the License at
+ <p>http://www.apache.org/licenses/LICENSE-2.0
+
+ <p>Unless required by applicable law or agreed to in writing, software distributed under the
+ License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ express or implied. See the License for the specific language governing permissions and
+ limitations under the License.
+ */
+