1 files changed, 27 insertions, 25 deletions

diff --git a/helix-core/src/syntax.rs b/helix-core/src/syntax.rs
index b20d9092..1cfa04e7 100644
--- a/helix-core/src/syntax.rs
+++ b/helix-core/src/syntax.rs
@@ -234,6 +234,7 @@ pub struct TextObjectQuery {
     pub query: Query,
 }
 
+#[derive(Debug)]
 pub enum CapturedNode<'a> {
     Single(Node<'a>),
     /// Guaranteed to be not empty
@@ -268,12 +269,12 @@ impl TextObjectQuery {
     /// and support for this is partial and could use improvement.
     ///
     /// ```query
-    /// ;; supported:
     /// (comment)+ @capture
     ///
-    /// ;; unsupported:
+    /// ; OR
     /// (
-    ///   (comment)+
+    ///   (comment)*
+    ///   .
     ///   (function)
     /// ) @capture
     /// ```
@@ -299,28 +300,29 @@ impl TextObjectQuery {
         let capture_idx = capture_names
             .iter()
             .find_map(|cap| self.query.capture_index_for_name(cap))?;
-        let captures = cursor.matches(&self.query, node, RopeProvider(slice));
-
-        let nodes = captures.flat_map(move |mat| {
-            let captures = mat.captures.iter().filter(move |c| c.index == capture_idx);
-            let nodes = captures.map(|c| c.node);
-            let pattern_idx = mat.pattern_index;
-            let quantifier = self.query.capture_quantifiers(pattern_idx)[capture_idx as usize];
-
-            let iter: Box<dyn Iterator<Item = CapturedNode>> = match quantifier {
-                CaptureQuantifier::OneOrMore | CaptureQuantifier::ZeroOrMore => {
-                    let nodes: Vec<Node> = nodes.collect();
-                    if nodes.is_empty() {
-                        Box::new(std::iter::empty())
-                    } else {
-                        Box::new(std::iter::once(CapturedNode::Grouped(nodes)))
-                    }
+
+        let nodes = cursor
+            .captures(&self.query, node, RopeProvider(slice))
+            .filter_map(move |(mat, _)| {
+                let nodes: Vec<_> = mat
+                    .captures
+                    .iter()
+                    .filter_map(|x| {
+                        if x.index == capture_idx {
+                            Some(x.node)
+                        } else {
+                            None
+                        }
+                    })
+                    .collect();
+
+                if nodes.len() > 1 {
+                    Some(CapturedNode::Grouped(nodes))
+                } else {
+                    nodes.into_iter().map(CapturedNode::Single).next()
                 }
-                _ => Box::new(nodes.map(CapturedNode::Single)),
-            };
+            });
 
-            iter
-        });
         Some(nodes)
     }
 }
@@ -1122,8 +1124,8 @@ pub(crate) fn generate_edits(
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::{iter, mem, ops, str, usize};
 use tree_sitter::{
-    CaptureQuantifier, Language as Grammar, Node, Parser, Point, Query, QueryCaptures, QueryCursor,
-    QueryError, QueryMatch, Range, TextProvider, Tree,
+    Language as Grammar, Node, Parser, Point, Query, QueryCaptures, QueryCursor, QueryError,
+    QueryMatch, Range, TextProvider, Tree,
 };
 
 const CANCELLATION_CHECK_INTERVAL: usize = 100;