Include macro attributes to impls, structs, enums, functions etc. textobjects (#2494)

2 files changed, 78 insertions, 37 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;
diff --git a/runtime/queries/rust/textobjects.scm b/runtime/queries/rust/textobjects.scm
index 086db67a..99b5eb4d 100644
--- a/runtime/queries/rust/textobjects.scm
+++ b/runtime/queries/rust/textobjects.scm
@@ -1,23 +1,62 @@
-(function_item
-  body: (_) @function.inside) @function.around
+(
+  [
+    (attribute_item)+
+    (line_comment)+
+  ]*
+  .
+  (function_item
+    body: (_) @function.inside)) @function.around
 
-(struct_item
-  body: (_) @class.inside) @class.around
+(
+  [
+    (attribute_item)+
+    (line_comment)+
+  ]*
+  .
+  (struct_item
+    body: (_) @class.inside)) @class.around
 
-(enum_item
-  body: (_) @class.inside) @class.around
+(
+  [
+    (attribute_item)+
+    (line_comment)+
+  ]*
+  .
+  (enum_item
+    body: (_) @class.inside)) @class.around
 
-(union_item
-  body: (_) @class.inside) @class.around
+(
+  [
+    (attribute_item)+
+    (line_comment)+
+  ]*
+  .
+  (union_item
+    body: (_) @class.inside)) @class.around
 
-(trait_item
-  body: (_) @class.inside) @class.around
+(
+  [
+    (attribute_item)+
+    (line_comment)+
+  ]*
+  .
+  (trait_item
+    body: (_) @class.inside)) @class.around
 
-(impl_item
-  body: (_) @class.inside) @class.around
+(
+  [
+    (attribute_item)+
+    (line_comment)+
+  ]*
+  .
+  (impl_item
+    body: (_) @class.inside)) @class.around
 
 (parameters
   (_) @parameter.inside)
+  
+(type_parameters
+  (_) @parameter.inside)
 
 (closure_parameters
   (_) @parameter.inside)