Reduce the number of allocations needed to find a specific child/sibling

src/api.rs

@@ -94,7 +94,7 @@ impl<L: Language> From<SyntaxToken<L>> for SyntaxElement<L> {
     }
 }
 
-impl<L: Language> SyntaxNode<L> {
+impl<'a, L: Language> SyntaxNode<L> {
     pub fn new_root(green: GreenNode) -> SyntaxNode<L> {
         SyntaxNode::from(cursor::SyntaxNode::new_root(green))
     }
@@ -137,10 +137,18 @@ impl<L: Language> SyntaxNode<L> {
         SyntaxNodeChildren { raw: self.raw.children(), _p: PhantomData }
     }
 
+    pub fn children_matching(&self, matcher: fn(SyntaxKind) -> bool) -> SyntaxNodeChildrenMatching<L> {
+        SyntaxNodeChildrenMatching { raw: self.raw.children_matching(matcher), _p: PhantomData }
+    }
+
     pub fn children_with_tokens(&self) -> SyntaxElementChildren<L> {
         SyntaxElementChildren { raw: self.raw.children_with_tokens(), _p: PhantomData }
     }
 
+    pub fn children_with_tokens_matching(&self, matcher: &'a impl Fn(SyntaxKind) -> bool) -> SyntaxElementChildrenMatching<'a, L> {
+        SyntaxElementChildrenMatching { raw: self.raw.children_with_tokens_matching(matcher), _p: PhantomData }
+    }
+
     pub fn first_child(&self) -> Option<SyntaxNode<L>> {
         self.raw.first_child().map(Self::from)
     }
@@ -388,6 +396,19 @@ impl<L: Language> Iterator for SyntaxNodeChildren<L> {
     }
 }
 
+#[derive(Debug, Clone)]
+pub struct SyntaxNodeChildrenMatching<L: Language> {
+    raw: cursor::SyntaxNodeChildrenMatching<fn(SyntaxKind) -> bool>,
+    _p: PhantomData<L>,
+}
+
+impl<L: Language> Iterator for SyntaxNodeChildrenMatching<L> {
+    type Item = SyntaxNode<L>;
+    fn next(&mut self) -> Option<Self::Item> {
+        self.raw.next().map(SyntaxNode::from)
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct SyntaxElementChildren<L: Language> {
     raw: cursor::SyntaxElementChildren,
@@ -401,6 +422,18 @@ impl<L: Language> Iterator for SyntaxElementChildren<L> {
     }
 }
 
+pub struct SyntaxElementChildrenMatching<'a, L: Language> {
+    raw: cursor::SyntaxElementChildrenMatching<&'a dyn Fn(SyntaxKind) -> bool>,
+    _p: PhantomData<L>,
+}
+
+impl<'a, L: Language> Iterator for SyntaxElementChildrenMatching<'a, L> {
+    type Item = SyntaxElement<L>;
+    fn next(&mut self) -> Option<Self::Item> {
+        self.raw.next().map(NodeOrToken::from)
+    }
+}
+
 pub struct Preorder<L: Language> {
     raw: cursor::Preorder,
     _p: PhantomData<L>,

src/cursor.rs

@@ -388,6 +388,24 @@ impl NodeData {
             })
         })
     }
+
+    fn next_sibling_matching(&self, matcher: &impl Fn(SyntaxKind) -> bool) -> Option<SyntaxNode> {
+        let mut siblings = self.green_siblings().enumerate();
+        let index = self.index() as usize;
+
+        siblings.nth(index);
+        siblings.find_map(|(index, child)| {
+            if !matcher(child.as_ref().kind()) {
+                return None;
+            }
+            child.as_ref().into_node().and_then(|green| {
+                let parent = self.parent_node()?;
+                let offset = parent.offset() + child.rel_offset();
+                Some(SyntaxNode::new_child(green, parent, index as u32, offset))
+            })
+        })
+    }
+
     fn prev_sibling(&self) -> Option<SyntaxNode> {
         let mut rev_siblings = self.green_siblings().enumerate().rev();
         let index = rev_siblings.len() - (self.index() as usize);
@@ -412,6 +430,22 @@ impl NodeData {
             Some(SyntaxElement::new(child.as_ref(), parent, index as u32, offset))
         })
     }
+
+    fn next_sibling_or_token_matching(&self, matcher: &impl Fn(SyntaxKind) -> bool) -> Option<SyntaxElement> {
+        let mut siblings = self.green_siblings().enumerate();
+        let index = self.index() as usize + 1;
+
+        siblings.nth(index);
+        siblings.find_map(|(index, child)| {
+            if !matcher(child.as_ref().kind()) {
+                return None;
+            }
+            let parent = self.parent_node()?;
+            let offset = parent.offset() + child.rel_offset();
+            Some(SyntaxElement::new(child.as_ref(), parent, index as u32, offset))
+        })
+    }
+
     fn prev_sibling_or_token(&self) -> Option<SyntaxElement> {
         let mut siblings = self.green_siblings().enumerate();
         let index = self.index().checked_sub(1)? as usize;
@@ -630,11 +664,21 @@ impl SyntaxNode {
         SyntaxNodeChildren::new(self.clone())
     }
 
+    #[inline]
+    pub fn children_matching<F: Fn(SyntaxKind) -> bool>(&self, matcher: F) -> SyntaxNodeChildrenMatching<F> {
+        SyntaxNodeChildrenMatching::new(self.clone(), matcher)
+    }
+
     #[inline]
     pub fn children_with_tokens(&self) -> SyntaxElementChildren {
         SyntaxElementChildren::new(self.clone())
     }
 
+    #[inline]
+    pub fn children_with_tokens_matching<F: Fn(SyntaxKind) -> bool>(&self, matcher: F) -> SyntaxElementChildrenMatching<F> {
+        SyntaxElementChildrenMatching::new(self.clone(), matcher)
+    }
+
     pub fn first_child(&self) -> Option<SyntaxNode> {
         self.green_ref().children().raw.enumerate().find_map(|(index, child)| {
             child.as_ref().into_node().map(|green| {
@@ -647,6 +691,23 @@ impl SyntaxNode {
             })
         })
     }
+
+    pub fn first_child_matching(&self, matcher: &impl Fn(SyntaxKind) -> bool) -> Option<SyntaxNode> {
+        self.green_ref().children().raw.enumerate().find_map(|(index, child)| {
+            if !matcher(child.as_ref().kind()) {
+                return None;
+            }
+            child.as_ref().into_node().map(|green| {
+                SyntaxNode::new_child(
+                    green,
+                    self.clone(),
+                    index as u32,
+                    self.offset() + child.rel_offset(),
+                )
+            })
+        })
+    }
+
     pub fn last_child(&self) -> Option<SyntaxNode> {
         self.green_ref().children().raw.enumerate().rev().find_map(|(index, child)| {
             child.as_ref().into_node().map(|green| {
@@ -665,6 +726,16 @@ impl SyntaxNode {
             SyntaxElement::new(child.as_ref(), self.clone(), 0, self.offset() + child.rel_offset())
         })
     }
+
+    pub fn first_child_or_token_matching(&self, matcher: &impl Fn(SyntaxKind) -> bool) -> Option<SyntaxElement> {
+        self.green_ref().children().raw.find_map(|child| {
+            if !matcher(child.as_ref().kind()) {
+                return None;
+            }
+            Some(SyntaxElement::new(child.as_ref(), self.clone(), 0, self.offset() + child.rel_offset()))
+        })
+    }
+
     pub fn last_child_or_token(&self) -> Option<SyntaxElement> {
         self.green_ref().children().raw.enumerate().next_back().map(|(index, child)| {
             SyntaxElement::new(
@@ -679,13 +750,23 @@ impl SyntaxNode {
     pub fn next_sibling(&self) -> Option<SyntaxNode> {
         self.data().next_sibling()
     }
+
+    pub fn next_sibling_matching(&self, matcher: &impl Fn(SyntaxKind) -> bool) -> Option<SyntaxNode> {
+        self.data().next_sibling_matching(matcher)
+    }
+
     pub fn prev_sibling(&self) -> Option<SyntaxNode> {
         self.data().prev_sibling()
     }
 
     pub fn next_sibling_or_token(&self) -> Option<SyntaxElement> {
         self.data().next_sibling_or_token()
     }
+
+    pub fn next_sibling_or_token_matching(&self, matcher: &impl Fn(SyntaxKind) -> bool) -> Option<SyntaxElement> {
+        self.data().next_sibling_or_token_matching(matcher)
+    }
+
     pub fn prev_sibling_or_token(&self) -> Option<SyntaxElement> {
         self.data().prev_sibling_or_token()
     }
@@ -910,6 +991,11 @@ impl SyntaxToken {
     pub fn next_sibling_or_token(&self) -> Option<SyntaxElement> {
         self.data().next_sibling_or_token()
     }
+
+    pub fn next_sibling_or_token_matching(&self, matcher: &impl Fn(SyntaxKind) -> bool) -> Option<SyntaxElement> {
+        self.data().next_sibling_or_token_matching(matcher)
+    }
+
     pub fn prev_sibling_or_token(&self) -> Option<SyntaxElement> {
         self.data().prev_sibling_or_token()
     }
@@ -1028,6 +1114,14 @@ impl SyntaxElement {
             NodeOrToken::Token(it) => it.next_sibling_or_token(),
         }
     }
+
+    pub fn next_sibling_or_token_matching(&self, matcher: &impl Fn(SyntaxKind) -> bool) -> Option<SyntaxElement> {
+        match self {
+            NodeOrToken::Node(it) => it.next_sibling_or_token_matching(matcher),
+            NodeOrToken::Token(it) => it.next_sibling_or_token_matching(matcher),
+        }
+    }
+
     pub fn prev_sibling_or_token(&self) -> Option<SyntaxElement> {
         match self {
             NodeOrToken::Node(it) => it.prev_sibling_or_token(),
@@ -1152,6 +1246,28 @@ impl Iterator for SyntaxNodeChildren {
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct SyntaxNodeChildrenMatching<F: Fn(SyntaxKind) -> bool> {
+    next: Option<SyntaxNode>,
+    matcher: F,
+}
+
+impl<F: Fn(SyntaxKind) -> bool> SyntaxNodeChildrenMatching<F> {
+    fn new(parent: SyntaxNode, matcher: F) -> SyntaxNodeChildrenMatching<F> {
+        SyntaxNodeChildrenMatching { next: parent.first_child_matching(&matcher), matcher }
+    }
+}
+
+impl<F: Fn(SyntaxKind) -> bool> Iterator for SyntaxNodeChildrenMatching<F> {
+    type Item = SyntaxNode;
+    fn next(&mut self) -> Option<SyntaxNode> {
+        self.next.take().map(|next| {
+            self.next = next.next_sibling_matching(&self.matcher);
+            next
+        })
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct SyntaxElementChildren {
     next: Option<SyntaxElement>,
@@ -1173,6 +1289,28 @@ impl Iterator for SyntaxElementChildren {
     }
 }
 
+#[derive(Clone, Debug)]
+pub struct SyntaxElementChildrenMatching<F: Fn(SyntaxKind) -> bool> {
+    next: Option<SyntaxElement>,
+    matcher: F,
+}
+
+impl<F: Fn(SyntaxKind) -> bool> SyntaxElementChildrenMatching<F> {
+    fn new(parent: SyntaxNode, matcher: F) -> SyntaxElementChildrenMatching<F> {
+        SyntaxElementChildrenMatching { next: parent.first_child_or_token_matching(&matcher), matcher }
+    }
+}
+
+impl<F: Fn(SyntaxKind) -> bool> Iterator for SyntaxElementChildrenMatching<F> {
+    type Item = SyntaxElement;
+    fn next(&mut self) -> Option<SyntaxElement> {
+        self.next.take().map(|next| {
+            self.next = next.next_sibling_or_token_matching(&self.matcher);
+            next
+        })
+    }
+}
+
 pub struct Preorder {
     start: SyntaxNode,
     next: Option<WalkEvent<SyntaxNode>>,

src/lib.rs

@@ -29,7 +29,7 @@ pub use text_size::{TextLen, TextRange, TextSize};
 
 pub use crate::{
     api::{
-        Language, SyntaxElement, SyntaxElementChildren, SyntaxNode, SyntaxNodeChildren, SyntaxToken,
+        Language, SyntaxElement, SyntaxElementChildren, SyntaxNode, SyntaxNodeChildren, SyntaxNodeChildrenMatching, SyntaxToken,
     },
     green::{
         Checkpoint, Children, GreenNode, GreenNodeBuilder, GreenNodeData, GreenToken,