diff --git a/src/iter/ast_node_iter.rs b/src/iter/ast_node_iter.rs index dc181da..ee92efa 100644 --- a/src/iter/ast_node_iter.rs +++ b/src/iter/ast_node_iter.rs @@ -5,6 +5,8 @@ use super::macros::children_iter; use super::macros::empty_iter; use crate::types::Bold; use crate::types::Code; +use crate::types::DocumentElement; +use crate::types::Element; use crate::types::Heading; use crate::types::Italic; use crate::types::Object; @@ -79,14 +81,18 @@ pub enum AstNodeIter<'r, 's> { } pub struct HeadingIter<'r, 's> { - next: std::slice::Iter<'r, Object<'s>>, + title_next: std::slice::Iter<'r, Object<'s>>, + children_next: std::slice::Iter<'r, DocumentElement<'s>>, } impl<'r, 's> Iterator for HeadingIter<'r, 's> { type Item = AstNode<'r, 's>; fn next(&mut self) -> Option { - self.next.next().map(Into::::into) + self.title_next + .next() + .map(Into::::into) + .or_else(|| self.children_next.next().map(Into::::into)) } } @@ -97,11 +103,8 @@ impl<'r, 's> IntoIterator for &'r Heading<'s> { fn into_iter(self) -> Self::IntoIter { HeadingIter { - next: self - .title - .iter() - .map(Into::::into) - .chain(self.children.iter().map(Into::::into)), + title_next: self.title.iter(), + children_next: self.children.iter(), } } }