From c47b4301351956e37a6d1b85201aafc36d67be93 Mon Sep 17 00:00:00 2001
From: Domenic Quirl <DomenicQuirl@pm.me>
Date: Sun, 21 Feb 2021 20:04:17 +0100
Subject: [PATCH] add methods for arity of syntax nodes

---
 src/green/element.rs |  4 ++--
 src/green/node.rs    |  5 +++++
 src/syntax.rs        | 14 ++++++++++++++
 3 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/green/element.rs b/src/green/element.rs
index 3a62f47..8b7af64 100644
--- a/src/green/element.rs
+++ b/src/green/element.rs
@@ -13,7 +13,7 @@ pub(super) type GreenElement = NodeOrToken<GreenNode, GreenToken>;
 pub(crate) type GreenElementRef<'a> = NodeOrToken<&'a GreenNode, &'a GreenToken>;
 
 #[repr(transparent)]
-pub(super) struct PackedGreenElement {
+pub(crate) struct PackedGreenElement {
     ptr: ErasedPtr,
 }
 
@@ -113,7 +113,7 @@ impl From<PackedGreenElement> for GreenElement {
 }
 
 impl PackedGreenElement {
-    fn is_node(&self) -> bool {
+    pub(crate) fn is_node(&self) -> bool {
         self.ptr as usize & 1 == 0
     }
 
diff --git a/src/green/node.rs b/src/green/node.rs
index 46a4296..1d7f632 100644
--- a/src/green/node.rs
+++ b/src/green/node.rs
@@ -115,6 +115,11 @@ impl GreenNode {
         self.data.header.header.text_len
     }
 
+    #[inline]
+    pub(crate) fn iter(&self) -> slice::Iter<'_, PackedGreenElement> {
+        self.data.slice.iter()
+    }
+
     /// Iterator over all children of this node.
     #[inline]
     pub fn children(&self) -> Children<'_> {
diff --git a/src/syntax.rs b/src/syntax.rs
index fb96173..c6baaab 100644
--- a/src/syntax.rs
+++ b/src/syntax.rs
@@ -695,6 +695,20 @@ impl<L: Language, D, R> SyntaxNode<L, D, R> {
         }
     }
 
+    /// The number of child nodes (!) of this node.
+    ///
+    /// If you want to also consider leafs, see [`arity_with_tokens`](SyntaxNode::arity_with_tokens).
+    #[inline]
+    pub fn arity(&self) -> usize {
+        self.green().iter().filter(|&child| child.is_node()).count()
+    }
+
+    /// The number of children of this node.
+    #[inline]
+    pub fn arity_with_tokens(&self) -> usize {
+        self.data().children.len()
+    }
+
     /// Returns an iterator along the chain of parents of this node.
     #[inline]
     pub fn ancestors(&self) -> impl Iterator<Item = &SyntaxNode<L, D, R>> {