diff --git a/src/iter/ast_node_iter.rs b/src/iter/ast_node_iter.rs index ee92efa..d0e18ea 100644 --- a/src/iter/ast_node_iter.rs +++ b/src/iter/ast_node_iter.rs @@ -5,6 +5,7 @@ use super::macros::children_iter; use super::macros::empty_iter; use crate::types::Bold; use crate::types::Code; +use crate::types::Document; use crate::types::DocumentElement; use crate::types::Element; use crate::types::Heading; @@ -13,6 +14,7 @@ use crate::types::Object; use crate::types::PlainText; use crate::types::RadioLink; use crate::types::RegularLink; +use crate::types::Section; use crate::types::StrikeThrough; use crate::types::Underline; use crate::types::Verbatim; @@ -24,9 +26,9 @@ use crate::types::Verbatim; /// This only iterates over AST nodes, so an AstNodeIter::Heading would iterate over both the title and section contents, but it would not iterate over simple strings like the TODO keyword or priority. pub enum AstNodeIter<'r, 's> { // Document Nodes - // Document(DocumentIter<'r, 's>), + Document(DocumentIter<'r, 's>), Heading(HeadingIter<'r, 's>), - // Section(SectionIter<'r, 's>), + Section(SectionIter<'r, 's>), // Elements // Paragraph(ParagraphIter<'r, 's>), // PlainList(PlainListIter<'r, 's>), @@ -80,6 +82,35 @@ pub enum AstNodeIter<'r, 's> { // Timestamp(TimestampIter<'r, 's>), } +pub struct DocumentIter<'r, 's> { + zeroth_section_next: std::option::Iter<'r, Section<'s>>, + children_next: std::slice::Iter<'r, Heading<'s>>, +} + +impl<'r, 's> Iterator for DocumentIter<'r, 's> { + type Item = AstNode<'r, 's>; + + fn next(&mut self) -> Option { + self.zeroth_section_next + .next() + .map(Into::::into) + .or_else(|| self.children_next.next().map(Into::::into)) + } +} + +impl<'r, 's> IntoIterator for &'r Document<'s> { + type Item = AstNode<'r, 's>; + + type IntoIter = DocumentIter<'r, 's>; + + fn into_iter(self) -> Self::IntoIter { + DocumentIter { + zeroth_section_next: self.zeroth_section.iter(), + children_next: self.children.iter(), + } + } +} + pub struct HeadingIter<'r, 's> { title_next: std::slice::Iter<'r, Object<'s>>, children_next: std::slice::Iter<'r, DocumentElement<'s>>, @@ -109,6 +140,7 @@ impl<'r, 's> IntoIterator for &'r Heading<'s> { } } +children_iter!(Section<'s>, SectionIter, std::slice::Iter<'r, Element<'s>>); children_iter!(Bold<'s>, BoldIter, std::slice::Iter<'r, Object<'s>>); children_iter!(Italic<'s>, ItalicIter, std::slice::Iter<'r, Object<'s>>); children_iter!(