Repository: ast-grep/ast-grep (language: rust)

Issue:
[feature] improve HTML language extraction for ts/scss



Discussion / hints:

Fix patch:
diff --git a/crates/cli/src/lang.rs b/crates/cli/src/lang.rs
index c85a1718ae..99b6a9bc35 100644
--- a/crates/cli/src/lang.rs
+++ b/crates/cli/src/lang.rs
@@ -2,7 +2,10 @@ mod custom_lang;
 mod lang_globs;
 
 use anyhow::Result;
-use ast_grep_core::language::TSLanguage;
+use ast_grep_core::{
+  language::{TSLanguage, TSRange},
+  Doc, Node,
+};
 use ast_grep_dynamic::DynamicLang;
 use ast_grep_language::{Language, SupportLang};
 use ignore::types::Types;
@@ -53,6 +56,14 @@ impl SgLang {
     let customs = DynamicLang::all_langs().into_iter().map(Self::Custom);
     builtin.chain(customs).collect()
   }
+
+  pub fn injectable_sg_langs(&self) -> Option<impl Iterator<Item = Self>> {
+    let langs = self.injectable_languages()?;
+    // TODO: handle injected languages not found
+    // e.g vue can inject scss which is not supported by sg
+    let iter = langs.iter().filter_map(|s| SgLang::from_str(s).ok());
+    Some(iter)
+  }
 }
 
 impl Display for SgLang {
@@ -107,6 +118,11 @@ impl From<SupportLang> for SgLang {
     Self::Builtin(value)
   }
 }
+impl From<DynamicLang> for SgLang {
+  fn from(value: DynamicLang) -> Self {
+    Self::Custom(value)
+  }
+}
 
 use SgLang::*;
 impl Language for SgLang {
@@ -148,6 +164,24 @@ impl Language for SgLang {
       Custom(c) => c.expando_char(),
     }
   }
+
+  fn injectable_languages(&self) -> Option<&'static [&'static str]> {
+    match self {
+      Builtin(b) => b.injectable_languages(),
+      Custom(c) => c.injectable_languages(),
+    }
+  }
+
+  fn extract_injections<D: Doc>(
+    &self,
+    root: Node<D>,
+    conv: impl Fn(Self) -> D::Lang,
+  ) -> HashMap<String, Vec<TSRange>> {
+    match self {
+      Builtin(b) => b.extract_injections(root, |b| conv(Builtin(b))),
+      Custom(c) => c.extract_injections(root, |c| conv(Custom(c))),
+    }
+  }
 }
 
 #[cfg(test)]
diff --git a/crates/cli/src/run.rs b/crates/cli/src/run.rs
index b80ee4c6be..c11f5d0351 100644
--- a/crates/cli/src/run.rs
+++ b/crates/cli/src/run.rs
@@ -12,7 +12,9 @@ use crate::debug::DebugFormat;
 use crate::error::ErrorContext as EC;
 use crate::lang::SgLang;
 use crate::print::{ColoredPrinter, Diff, Heading, InteractivePrinter, JSONPrinter, Printer};
-use crate::utils::{filter_file_pattern, InputArgs, MatchUnit, OutputArgs};
+use crate::utils::{
+  filter_file_pattern, filter_file_pattern_injection, InputArgs, MatchUnit, OutputArgs,
+};
 use crate::utils::{Items, PathWorker, StdInWorker, Worker};
 
 // NOTE: have to register custom lang before clap read arg
@@ -228,8 +230,16 @@ impl<P: Printer> PathWorker for RunWithInferredLang<P> {
   fn produce_item(&self, path: &Path) -> Option<Self::Item> {
     let lang = SgLang::from_path(path)?;
     let matcher = self.arg.build_pattern(lang).ok()?;
-    let match_unit = filter_file_pattern(path, lang, matcher)?;
-    Some((match_unit, lang))
+    if let Some(match_unit) = filter_file_pattern(path, lang, matcher) {
+      // match root
+      return Some((match_unit, lang));
+    }
+    // match sub region
+    let matchers = lang.injectable_sg_langs()?.filter_map(|l| {
+      let pattern = self.arg.build_pattern(l).ok()?;
+      Some((l, pattern))
+    });
+    filter_file_pattern_injection(path, lang, matchers)
   }
 }
 
diff --git a/crates/cli/src/utils.rs b/crates/cli/src/utils.rs
index 98b4458e09..a9b2abe304 100644
--- a/crates/cli/src/utils.rs
+++ b/crates/cli/src/utils.rs
@@ -19,6 +19,7 @@ use std::fs::read_to_string;
 use std::io::stdout;
 use std::io::Write;
 use std::path::{Path, PathBuf};
+use std::str::FromStr;
 use std::sync::{mpsc, Arc};
 
 fn read_char() -> Result<char> {
@@ -232,6 +233,37 @@ pub fn filter_file_pattern(
   })
 }
 
+pub fn filter_file_pattern_injection(
+  path: &Path,
+  lang: SgLang,
+  mut matchers: impl Iterator<Item = (SgLang, Pattern<SgLang>)>,
+) -> Option<(MatchUnit<Pattern<SgLang>>, SgLang)> {
+  let file_content = read_file(path)?;
+  let grep = lang.ast_grep(&file_content);
+  let injections = grep.inner.get_injections(|s| SgLang::from_str(s).ok())?;
+  matchers.find_map(|(i_lang, matcher)| {
+    let fixed = matcher.fixed_string();
+    if !fixed.is_empty() && !file_content.contains(&*fixed) {
+      return None;
+    }
+    let injection = injections.iter().find(|i| *i.lang() == i_lang)?;
+    let has_match = injection.root().find(&matcher).is_some();
+    let injected = AstGrep {
+      inner: injection.clone(),
+    };
+    has_match.then(|| {
+      (
+        MatchUnit {
+          grep: injected,
+          path: path.to_path_buf(),
+          matcher,
+        },
+        i_lang,
+      )
+    })
+  })
+}
+
 const MAX_FILE_SIZE: usize = 3_000_000;
 const MAX_LINE_COUNT: usize = 200_000;
 
diff --git a/crates/core/src/language.rs b/crates/core/src/language.rs
index 03f27d1ec7..92a05fdd1d 100644
--- a/crates/core/src/language.rs
+++ b/crates/core/src/language.rs
@@ -1,9 +1,10 @@
 use crate::meta_var::{extract_meta_var, MetaVariable};
-use crate::AstGrep;
-use crate::StrDoc;
+use crate::{AstGrep, Doc, Node, StrDoc};
 use std::borrow::Cow;
+use std::collections::HashMap;
 use std::path::Path;
 pub use tree_sitter::Language as TSLanguage;
+pub use tree_sitter::{Point as TSPoint, Range as TSRange};
 
 /// Trait to abstract ts-language usage in ast-grep, which includes:
 /// * which character is used for meta variable.
@@ -54,6 +55,23 @@ pub trait Language: Clone {
   fn extract_meta_var(&self, source: &str) -> Option<MetaVariable> {
     extract_meta_var(source, self.expando_char())
   }
+
+  fn injectable_languages(&self) -> Option<&'static [&'static str]> {
+    None
+  }
+
+  /// get injected language regions in the root document. e.g. get JavaScripts in HTML
+  /// it will return a list of tuples of (language, regions).
+  /// The first item is the embedded region language, e.g. javascript
+  /// The second item is a list of regions in tree_sitter.
+  /// also see https://tree-sitter.github.io/tree-sitter/using-parsers#multi-language-documents
+  fn extract_injections<D: Doc>(
+    &self,
+    _root: Node<D>,
+    _conv: impl Fn(Self) -> D::Lang,
+  ) -> HashMap<String, Vec<TSRange>> {
+    HashMap::new()
+  }
 }
 
 impl Language for TSLanguage {
diff --git a/crates/core/src/node.rs b/crates/core/src/node.rs
index 2737c777c4..c72418441b 100644
--- a/crates/core/src/node.rs
+++ b/crates/core/src/node.rs
@@ -86,6 +86,33 @@ impl<D: Doc> Root<D> {
     debug_assert!(self.check_lineage(&node.inner));
     node.root = self;
   }
+
+  pub fn get_injections<F: Fn(&str) -> Option<D::Lang>>(
+    &self,
+    get_lang: F,
+  ) -> Option<Vec<Root<D>>> {
+    let root = self.root();
+    let range = self.lang().extract_injections(root, |l| l);
+    if range.is_empty() {
+      return None;
+    }
+    let roots = range
+      .into_iter()
+      .filter_map(|(lang, ranges)| {
+        let lang = get_lang(&lang)?;
+        let source = self.doc.get_source();
+        let mut parser = tree_sitter::Parser::new().expect("TODO");
+        parser.set_included_ranges(&ranges).expect("TODO");
+        parser.set_language(&lang.get_ts_language()).expect("TODO");
+        let tree = source.parse_tree_sitter(&mut parser, None).expect("TODO");
+        tree.map(|t| Self {
+          inner: t,
+          doc: self.doc.clone_with_lang(lang),
+        })
+      })
+      .collect();
+    Some(roots)
+  }
 }
 
 /// 'r represents root lifetime
diff --git a/crates/core/src/source.rs b/crates/core/src/source.rs
index 0cd3c8f98a..28fc6d2a55 100644
--- a/crates/core/src/source.rs
+++ b/crates/core/src/source.rs
@@ -88,6 +88,7 @@ pub trait Doc: Clone {
     let lang = self.get_lang().get_ts_language();
     parse_lang(|p| source.parse_tree_sitter(p, old_tree), lang)
   }
+  fn clone_with_lang(&self, lang: Self::Lang) -> Self;
   /// TODO: are we paying too much to support str as Pattern/Replacer??
   /// this method converts string to Doc, so that we can support using
   /// string as replacer/searcher. Natively.
@@ -124,6 +125,9 @@ impl<L: Language> Doc for StrDoc<L> {
   fn from_str(src: &str, lang: L) -> Self {
     Self::new(src, lang)
   }
+  fn clone_with_lang(&self, lang: Self::Lang) -> Self {
+    Self::new(&self.src, lang)
+  }
 }
 
 pub trait Content: Sized {
diff --git a/crates/language/src/html.rs b/crates/language/src/html.rs
index 0439a5c41a..d96ab5ae31 100644
--- a/crates/language/src/html.rs
+++ b/crates/language/src/html.rs
@@ -1,40 +1,120 @@
-#![cfg(test)]
-use ast_grep_core::source::TSParseError;
+use super::pre_process_pattern;
+use ast_grep_core::language::{TSPoint, TSRange};
+use ast_grep_core::{matcher::KindMatcher, Doc, Node};
+use std::collections::HashMap;
 
-use super::*;
-
-fn test_match(query: &str, source: &str) {
-  use crate::test::test_match_lang;
-  test_match_lang(query, source, Html);
+// tree-sitter-html uses locale dependent iswalnum for tagName
+// https://github.com/tree-sitter/tree-sitter-html/blob/b5d9758e22b4d3d25704b72526670759a9e4d195/src/scanner.c#L194
+#[derive(Clone, Copy)]
+pub struct Html;
+impl ast_grep_core::language::Language for Html {
+  fn get_ts_language(&self) -> ast_grep_core::language::TSLanguage {
+    crate::parsers::language_html()
+  }
+  fn expando_char(&self) -> char {
+    'z'
+  }
+  fn pre_process_pattern<'q>(&self, query: &'q str) -> std::borrow::Cow<'q, str> {
+    pre_process_pattern(self.expando_char(), query)
+  }
+  fn injectable_languages(&self) -> Option<&'static [&'static str]> {
+    Some(&["css", "js", "ts", "tsx", "scss", "less", "stylus", "coffee"])
+  }
+  fn extract_injections<D: Doc>(
+    &self,
+    root: Node<D>,
+    conv: impl Fn(Self) -> D::Lang,
+  ) -> HashMap<String, Vec<TSRange>> {
+    let lang = conv(*self);
+    let mut map = HashMap::new();
+    let matcher = KindMatcher::new("script_element", lang.clone());
+    for script in root.find_all(matcher) {
+      let injected = find_lang(&script).unwrap_or_else(|| "js".into());
+      let content = script.children().find(|c| c.kind() == "raw_text");
+      if let Some(content) = content {
+        map
+          .entry(injected)
+          .or_insert_with(Vec::new)
+          .push(node_to_range(&content));
+      };
+    }
+    let matcher = KindMatcher::new("style_element", lang);
+    for style in root.find_all(matcher) {
+      let injected = find_lang(&style).unwrap_or_else(|| "css".into());
+      let content = style.children().find(|c| c.kind() == "raw_text");
+      if let Some(content) = content {
+        map
+          .entry(injected)
+          .or_insert_with(Vec::new)
+          .push(node_to_range(&content));
+      };
+    }
+    map
+  }
 }
 
-fn test_non_match(query: &str, source: &str) {
-  use crate::test::test_non_match_lang;
-  test_non_match_lang(query, source, Html);
+fn find_lang<D: Doc>(node: &Node<D>) -> Option<String> {
+  let html = node.lang();
+  let attr_matcher = KindMatcher::new("attribute", html.clone());
+  let name_matcher = KindMatcher::new("attribute_name", html.clone());
+  let val_matcher = KindMatcher::new("attribute_value", html.clone());
+  node.find_all(attr_matcher).find_map(|attr| {
+    let name = attr.find(&name_matcher)?;
+    if name.text() != "lang" {
+      return None;
+    }
+    let val = attr.find(&val_matcher)?;
+    Some(val.text().to_string())
+  })
 }
 
-#[test]
-fn test_html_match() {
-  test_match("<input>", "<input>");
-  test_match("<$TAG>", "<input>");
-  test_match("<$TAG class='foo'>$$$</$TAG>", "<div class='foo'></div>");
-  test_match("<div>$$$</div>", "<div>123</div>");
-  test_non_match("<$TAG class='foo'>$$$</$TAG>", "<div></div>");
-  test_non_match("<div>$$$</div>", "<div class='foo'>123</div>");
+fn node_to_range<D: Doc>(node: &Node<D>) -> TSRange {
+  let r = node.range();
+  let start = node.start_pos();
+  let sp = TSPoint::new(start.0 as u32, start.1 as u32);
+  let end = node.end_pos();
+  let ep = TSPoint::new(end.0 as u32, end.1 as u32);
+  TSRange::new(r.start as u32, r.end as u32, &sp, &ep)
 }
 
-fn test_replace(src: &str, pattern: &str, replacer: &str) -> Result<String, TSParseError> {
-  use crate::test::test_replace_lang;
-  test_replace_lang(src, pattern, replacer, Html)
-}
+#[cfg(test)]
+mod test {
+  use super::*;
+  use ast_grep_core::source::TSParseError;
+
+  fn test_match(query: &str, source: &str) {
+    use crate::test::test_match_lang;
+    test_match_lang(query, source, Html);
+  }
+
+  fn test_non_match(query: &str, source: &str) {
+    use crate::test::test_non_match_lang;
+    test_non_match_lang(query, source, Html);
+  }
+
+  #[test]
+  fn test_html_match() {
+    test_match("<input>", "<input>");
+    test_match("<$TAG>", "<input>");
+    test_match("<$TAG class='foo'>$$$</$TAG>", "<div class='foo'></div>");
+    test_match("<div>$$$</div>", "<div>123</div>");
+    test_non_match("<$TAG class='foo'>$$$</$TAG>", "<div></div>");
+    test_non_match("<div>$$$</div>", "<div class='foo'>123</div>");
+  }
+
+  fn test_replace(src: &str, pattern: &str, replacer: &str) -> Result<String, TSParseError> {
+    use crate::test::test_replace_lang;
+    test_replace_lang(src, pattern, replacer, Html)
+  }
 
-#[test]
-fn test_html_replace() -> Result<(), TSParseError> {
-  let ret = test_replace(
-    r#"<div class='foo'>bar</div>"#,
-    r#"<$TAG class='foo'>$$$B</$TAG>"#,
-    r#"<$TAG class='$$$B'>foo</$TAG>"#,
-  )?;
-  assert_eq!(ret, r#"<div class='bar'>foo</div>"#);
-  Ok(())
+  #[test]
+  fn test_html_replace() -> Result<(), TSParseError> {
+    let ret = test_replace(
+      r#"<div class='foo'>bar</div>"#,
+      r#"<$TAG class='foo'>$$$B</$TAG>"#,
+      r#"<$TAG class='$$$B'>foo</$TAG>"#,
+    )?;
+    assert_eq!(ret, r#"<div class='bar'>foo</div>"#);
+    Ok(())
+  }
 }
diff --git a/crates/language/src/lib.rs b/crates/language/src/lib.rs
index dec5331ad0..4893a3bd53 100644
--- a/crates/language/src/lib.rs
+++ b/crates/language/src/lib.rs
@@ -26,11 +26,15 @@ mod rust;
 mod scala;
 mod swift;
 
-use ast_grep_core::language::TSLanguage;
+use html::Html;
+
+use ast_grep_core::language::{TSLanguage, TSRange};
 use ast_grep_core::meta_var::MetaVariable;
+use ast_grep_core::{Doc, Node};
 use ignore::types::{Types, TypesBuilder};
 use serde::{de, Deserialize, Deserializer, Serialize};
 use std::borrow::Cow;
+use std::collections::HashMap;
 use std::fmt;
 use std::fmt::{Display, Formatter};
 use std::iter::repeat;
@@ -114,9 +118,6 @@ impl_lang_expando!(Go, language_go, 'µ');
 // https://ghc.gitlab.haskell.org/ghc/doc/users_guide/exts/unicode_syntax.html
 // and the tree-sitter-haskell grammar parses it too.
 impl_lang_expando!(Haskell, language_haskell, 'µ');
-// tree-sitter-html uses locale dependent iswalnum for tagName
-// https://github.com/tree-sitter/tree-sitter-html/blob/b5d9758e22b4d3d25704b72526670759a9e4d195/src/scanner.c#L194
-impl_lang_expando!(Html, language_html, 'z');
 // https://github.com/fwcd/tree-sitter-kotlin/pull/93
 impl_lang_expando!(Kotlin, language_kotlin, 'µ');
 // we can use any char in unicode range [:XID_Start:]
@@ -313,6 +314,18 @@ impl Language for SupportLang {
   impl_lang_method!(meta_var_char, () => char);
   impl_lang_method!(expando_char, () => char);
   impl_lang_method!(extract_meta_var, (source: &str) => Option<MetaVariable>);
+  impl_lang_method!(injectable_languages, () => Option<&'static [&'static str]>);
+
+  fn extract_injections<D: Doc>(
+    &self,
+    root: Node<D>,
+    conv: impl Fn(Self) -> D::Lang,
+  ) -> HashMap<String, Vec<TSRange>> {
+    match self {
+      SupportLang::Html => Html.extract_injections(root, |_| conv(SupportLang::Html)),
+      _ => HashMap::new(),
+    }
+  }
 
   fn pre_process_pattern<'q>(&self, query: &'q str) -> Cow<'q, str> {
     execute_lang_method! { self, pre_process_pattern, query }
diff --git a/crates/napi/src/doc.rs b/crates/napi/src/doc.rs
index 32612c42b7..89922c259a 100644
--- a/crates/napi/src/doc.rs
+++ b/crates/napi/src/doc.rs
@@ -161,6 +161,12 @@ impl Doc for JsDoc {
   fn from_str(src: &str, lang: Self::Lang) -> Self {
     JsDoc::new(src.into(), lang)
   }
+  fn clone_with_lang(&self, lang: Self::Lang) -> Self {
+    JsDoc {
+      source: self.source.clone(),
+      lang,
+    }
+  }
 }
 
 #[cfg(test)]


Test patch: