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

Issue:
support finding injected code in `run` with specified language



Discussion / hints:

Fix patch:
diff --git a/crates/cli/src/lang.rs b/crates/cli/src/lang.rs
index 10f0eaa2ec..13357875fb 100644
--- a/crates/cli/src/lang.rs
+++ b/crates/cli/src/lang.rs
@@ -68,14 +68,14 @@ impl SgLang {
 
   pub fn augmented_file_type(&self) -> Types {
     let self_type = self.file_types();
-    let injected_from = Self::all_langs().into_iter().filter_map(|lang| {
+    let injector = Self::all_langs().into_iter().filter_map(|lang| {
       lang
         .injectable_sg_langs()?
         .any(|l| l == *self)
         .then_some(lang)
     });
-    let injected_types = injected_from.map(|lang| lang.file_types());
-    let all_types = std::iter::once(self_type).chain(injected_types);
+    let injector_types = injector.map(|lang| lang.file_types());
+    let all_types = std::iter::once(self_type).chain(injector_types);
     lang_globs::merge_types(all_types)
   }
 }
diff --git a/crates/cli/src/run.rs b/crates/cli/src/run.rs
index 6697637f35..8d0b9d3ca1 100644
--- a/crates/cli/src/run.rs
+++ b/crates/cli/src/run.rs
@@ -227,19 +227,22 @@ impl<P: Printer> PathWorker for RunWithInferredLang<P> {
     self.arg.input.walk()
   }
 
-  fn produce_item(&self, path: &Path) -> Option<Self::Item> {
-    let lang = SgLang::from_path(path)?;
-    let matcher = self.arg.build_pattern(lang).ok()?;
-    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)
+  fn produce_item(&self, path: &Path) -> impl Iterator<Item = Self::Item> {
+    let root_region = || {
+      let lang = SgLang::from_path(path)?;
+      let matcher = self.arg.build_pattern(lang).ok()?;
+      filter_file_pattern(path, lang, matcher).map(|unit| (unit, lang))
+    };
+    let option = || {
+      let lang = SgLang::from_path(path)?;
+      // 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)
+    };
+    root_region().into_iter().chain(option())
   }
 }
 
@@ -296,16 +299,20 @@ impl<P: Printer> PathWorker for RunWithSpecificLang<P> {
     let lang = self.arg.lang.expect("must present");
     self.arg.input.walk_lang(lang)
   }
-  fn produce_item(&self, path: &Path) -> Option<Self::Item> {
+  fn produce_item(&self, path: &Path) -> impl Iterator<Item = Self::Item> {
     let arg = &self.arg;
     let pattern = self.pattern.clone();
     let lang = arg.lang.expect("must present");
-    let path_lang = SgLang::from_path(path)?;
-    if path_lang == lang {
-      filter_file_pattern(path, lang, pattern)
-    } else {
-      filter_file_pattern_injection(path, path_lang, std::iter::once((lang, pattern))).map(|n| n.0)
-    }
+    let option = || {
+      let path_lang = SgLang::from_path(path)?;
+      if path_lang == lang {
+        filter_file_pattern(path, lang, pattern)
+      } else {
+        filter_file_pattern_injection(path, path_lang, std::iter::once((lang, pattern)))
+          .map(|n| n.0)
+      }
+    };
+    option().into_iter()
   }
 }
 
diff --git a/crates/cli/src/scan.rs b/crates/cli/src/scan.rs
index 80fd65c1f7..bc77f4fb8a 100644
--- a/crates/cli/src/scan.rs
+++ b/crates/cli/src/scan.rs
@@ -165,19 +165,21 @@ impl<P: Printer> PathWorker for ScanWithConfig<P> {
   fn build_walk(&self) -> WalkParallel {
     self.arg.input.walk()
   }
-  fn produce_item(&self, path: &Path) -> Option<Self::Item> {
+  fn produce_item(&self, path: &Path) -> impl Iterator<Item = Self::Item> {
     let rules = self.configs.for_path(path);
     if rules.is_empty() {
-      return None;
+      return None.into_iter();
     }
     let lang = rules[0].language;
     let combined = CombinedScan::new(rules);
-    let unit = filter_file_interactive(path, lang, ast_grep_core::matcher::MatchAll)?;
+    let Some(unit) = filter_file_interactive(path, lang, ast_grep_core::matcher::MatchAll) else {
+      return None.into_iter();
+    };
     let hit_set = combined.find(&unit.grep);
     if !hit_set.is_empty() {
-      return Some((unit.path, unit.grep, hit_set));
+      return Some((unit.path, unit.grep, hit_set)).into_iter();
     }
-    None
+    None.into_iter()
   }
 }
 
diff --git a/crates/cli/src/utils.rs b/crates/cli/src/utils.rs
index c3270e8876..6ef6cbefe2 100644
--- a/crates/cli/src/utils.rs
+++ b/crates/cli/src/utils.rs
@@ -101,7 +101,7 @@ pub trait PathWorker: Worker {
   /// WalkParallel will determine what files will be processed.
   fn build_walk(&self) -> WalkParallel;
   /// Parse and find_match can be done in `produce_item`.
-  fn produce_item(&self, path: &Path) -> Option<Self::Item>;
+  fn produce_item(&self, path: &Path) -> impl Iterator<Item = Self::Item>;
 
   fn run_path(self) -> Result<()>
   where
@@ -169,16 +169,16 @@ fn run_worker<W: PathWorker + ?Sized + 'static>(worker: Arc<W>) -> Result<()> {
       let tx = tx.clone();
       let w = w.clone();
       Box::new(move |result| {
-        let producer = |path: PathBuf| w.produce_item(&path);
-        let maybe_result = filter_result(result).and_then(producer);
-        let result = match maybe_result {
-          Some(ret) => ret,
-          None => return WalkState::Continue,
+        let Some(p) = filter_result(result) else {
+          return WalkState::Continue;
         };
-        match tx.send(result) {
-          Ok(_) => WalkState::Continue,
-          Err(_) => WalkState::Quit,
+        for result in w.produce_item(&p) {
+          match tx.send(result) {
+            Ok(_) => continue,
+            Err(_) => return WalkState::Quit,
+          }
         }
+        WalkState::Continue
       })
     });
   });


Test patch: