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,