3 files changed, 65 insertions, 21 deletions

diff --git a/crates/libsyntax2/src/ast/mod.rs b/crates/libsyntax2/src/ast/mod.rs
index 881f380f3..0b6868547 100644
--- a/crates/libsyntax2/src/ast/mod.rs
+++ b/crates/libsyntax2/src/ast/mod.rs
@@ -1,10 +1,13 @@
 mod generated;
 
+use std::marker::PhantomData;
+
 use itertools::Itertools;
 use smol_str::SmolStr;
 
 use {
     SyntaxNodeRef, SyntaxKind::*,
+    yellow::{RefRoot, SyntaxNodeChildren},
 };
 pub use self::generated::*;
 
@@ -33,8 +36,8 @@ pub trait ArgListOwner<'a>: AstNode<'a> {
 }
 
 pub trait FnDefOwner<'a>: AstNode<'a> {
-    fn functions(self) -> Box<Iterator<Item=FnDef<'a>> + 'a> {
-        Box::new(children(self))
+    fn functions(self) -> AstNodeChildren<'a, FnDef<'a>> {
+        children(self)
     }
 }
 
@@ -49,8 +52,8 @@ pub trait TypeParamsOwner<'a>: AstNode<'a> {
 }
 
 pub trait AttrsOwner<'a>: AstNode<'a> {
-    fn attrs(self) -> Box<Iterator<Item=Attr<'a>> + 'a> {
-        Box::new(children(self))
+    fn attrs(self) -> AstNodeChildren<'a, Attr<'a>> {
+        children(self)
     }
 }
 
@@ -155,7 +158,7 @@ impl<'a> IfExpr<'a> {
     pub fn else_branch(self) -> Option<Block<'a>> {
         self.blocks().nth(1)
     }
-    fn blocks(self) -> impl Iterator<Item=Block<'a>> {
+    fn blocks(self) -> AstNodeChildren<'a, Block<'a>> {
         children(self)
     }
 }
@@ -164,8 +167,34 @@ fn child_opt<'a, P: AstNode<'a>, C: AstNode<'a>>(parent: P) -> Option<C> {
     children(parent).next()
 }
 
-fn children<'a, P: AstNode<'a>, C: AstNode<'a>>(parent: P) -> impl Iterator<Item=C> + 'a {
-    parent.syntax()
-        .children()
-        .filter_map(C::cast)
+fn children<'a, P: AstNode<'a>, C: AstNode<'a>>(parent: P) -> AstNodeChildren<'a, C> {
+    AstNodeChildren::new(parent.syntax())
+}
+
+
+#[derive(Debug)]
+pub struct AstNodeChildren<'a, N> {
+    inner: SyntaxNodeChildren<RefRoot<'a>>,
+    ph: PhantomData<N>,
+}
+
+impl<'a, N> AstNodeChildren<'a, N> {
+    fn new(parent: SyntaxNodeRef<'a>) -> Self {
+        AstNodeChildren {
+            inner: parent.children(),
+            ph: PhantomData,
+        }
+    }
+}
+
+impl<'a, N: AstNode<'a>> Iterator for AstNodeChildren<'a, N> {
+    type Item = N;
+    fn next(&mut self) -> Option<N> {
+        loop {
+            match N::cast(self.inner.next()?) {
+                Some(n) => return Some(n),
+                None => (),
+            }
+        }
+    }
 }
diff --git a/crates/libsyntax2/src/yellow/mod.rs b/crates/libsyntax2/src/yellow/mod.rs
index 82eda79d6..0596e702f 100644
--- a/crates/libsyntax2/src/yellow/mod.rs
+++ b/crates/libsyntax2/src/yellow/mod.rs
@@ -8,7 +8,7 @@ use std::{
     sync::Arc,
     ptr,
 };
-pub use self::syntax::{SyntaxNode, SyntaxNodeRef, SyntaxError};
+pub use self::syntax::{SyntaxNode, SyntaxNodeRef, SyntaxError, SyntaxNodeChildren};
 pub(crate) use self::{
     builder::GreenBuilder,
     green::GreenNode,
diff --git a/crates/libsyntax2/src/yellow/syntax.rs b/crates/libsyntax2/src/yellow/syntax.rs
index 444dbeb30..1d99cab4a 100644
--- a/crates/libsyntax2/src/yellow/syntax.rs
+++ b/crates/libsyntax2/src/yellow/syntax.rs
@@ -1,6 +1,7 @@
 use std::{
     fmt, sync::Arc,
     hash::{Hasher, Hash},
+    ops::Range,
 };
 
 use smol_str::SmolStr;
@@ -93,17 +94,11 @@ impl<R: TreeRoot> SyntaxNode<R> {
         SyntaxText::new(self.borrowed())
     }
 
-    pub fn children(&self) -> impl Iterator<Item = SyntaxNode<R>> {
-        let red = self.red;
-        let n_children = self.red().n_children();
-        let root = self.root.clone();
-        (0..n_children).map(move |i| {
-            let red = unsafe { red.get(root.syntax_root()) };
-            SyntaxNode {
-                root: root.clone(),
-                red: red.get_child(i).unwrap(),
-            }
-        })
+    pub fn children(&self) -> SyntaxNodeChildren<R> {
+        SyntaxNodeChildren {
+            parent: self.clone(),
+            iter: (0..self.red().n_children())
+        }
     }
 
     pub fn parent(&self) -> Option<SyntaxNode<R>> {
@@ -192,6 +187,26 @@ impl<R: TreeRoot> fmt::Debug for SyntaxNode<R> {
     }
 }
 
+#[derive(Debug)]
+pub struct SyntaxNodeChildren<R: TreeRoot> {
+    parent: SyntaxNode<R>,
+    iter: Range<usize>,
+}
+
+impl<R: TreeRoot> Iterator for SyntaxNodeChildren<R> {
+    type Item = SyntaxNode<R>;
+
+    fn next(&mut self) -> Option<SyntaxNode<R>> {
+        self.iter.next().map(|i| {
+            let red = self.parent.red();
+            SyntaxNode {
+                root: self.parent.root.clone(),
+                red: red.get_child(i).unwrap(),
+            }
+        })
+    }
+}
+
 fn has_short_text(kind: SyntaxKind) -> bool {
     match kind {
         IDENT | LIFETIME | INT_NUMBER | FLOAT_NUMBER => true,