Add TsQuery::all_captures()

Author: lionel-

crates/ark/src/lsp/inputs/package_namespace.rs

@@ -51,32 +51,23 @@ impl Namespace {
         "#;
         let mut ts_query = TsQuery::new(query_str)?;
 
-        let mut exports = Vec::new();
-        for capture in ts_query.captures_for(root_node, "exported", contents.as_bytes()) {
-            let symbol = capture
-                .utf8_text(contents.as_bytes())
-                .unwrap_or("")
-                .to_string();
-            exports.push(symbol);
-        }
-
-        let mut imports = Vec::new();
-        for capture in ts_query.captures_for(root_node, "imported", contents.as_bytes()) {
-            let symbol = capture
-                .utf8_text(contents.as_bytes())
-                .unwrap_or("")
-                .to_string();
-            imports.push(symbol);
-        }
-
-        let mut package_imports = Vec::new();
-        for capture in ts_query.captures_for(root_node, "bulk_imported", contents.as_bytes()) {
-            let symbol = capture
-                .utf8_text(contents.as_bytes())
-                .unwrap_or("")
-                .to_string();
-            package_imports.push(symbol);
-        }
+        let all_captures = ts_query.all_captures(root_node, contents.as_bytes());
+
+        let filter_captures = |capture_name: &str| -> Vec<String> {
+            all_captures
+                .iter()
+                .filter(|(name, _)| name == capture_name)
+                .map(|(_, node)| {
+                    node.utf8_text(contents.as_bytes())
+                        .unwrap_or("")
+                        .to_string()
+                })
+                .collect()
+        };
+
+        let mut exports = filter_captures("exported");
+        let mut imports = filter_captures("imported");
+        let mut package_imports = filter_captures("bulk_imported");
 
         // Take unique values of imports and exports. In the future we'll lint
         // this but for now just be defensive.

crates/ark/src/treesitter.rs

@@ -659,26 +659,40 @@ impl TsQuery {
         Self { query, cursor }
     }
 
-    /// Match query against `contents` and collect all nodes captured with the
-    /// given capture name
+    /// Run the query on `contents` and collect all captures as (capture_name, node) pairs
+    pub(crate) fn all_captures<'tree>(
+        &mut self,
+        node: tree_sitter::Node<'tree>,
+        contents: &[u8],
+    ) -> Vec<(String, tree_sitter::Node<'tree>)> {
+        self.cursor
+            .matches(&self.query, node, contents)
+            .flat_map(|m| {
+                m.captures.iter().map(|cap| {
+                    let cap_name = self.query.capture_names()[cap.index as usize].to_string();
+                    (cap_name, cap.node)
+                })
+            })
+            .collect()
+    }
+
+    /// Run the query on `contents` and filter captures that match `capture_name`
     pub(crate) fn captures_for<'tree>(
         &mut self,
         node: tree_sitter::Node<'tree>,
         capture_name: &str,
         contents: &[u8],
     ) -> Vec<tree_sitter::Node<'tree>> {
-        let mut result = Vec::new();
-
-        for m in self.cursor.matches(&self.query, node, contents) {
-            for cap in m.captures.iter() {
-                let cap_name = &self.query.capture_names()[cap.index as usize];
-                if *cap_name == capture_name {
-                    result.push(cap.node);
+        self.all_captures(node, contents)
+            .into_iter()
+            .filter_map(|(name, node)| {
+                if name == capture_name {
+                    Some(node)
+                } else {
+                    None
                 }
-            }
-        }
-
-        result
+            })
+            .collect()
     }
 }