From ab46a9e5c6cf44abc72bbac521070e428ad324d2 Mon Sep 17 00:00:00 2001
From: Tom Alexander <tom@fizz.buzz>
Date: Wed, 27 Sep 2023 15:56:45 -0400
Subject: [PATCH] Ran into issue with heading, naming this type is going to be
 a nightmare.

---
 src/iter/ast_node.rs      | 10 +++++++
 src/iter/ast_node_iter.rs | 62 +++++++++++++++++----------------------
 2 files changed, 37 insertions(+), 35 deletions(-)

diff --git a/src/iter/ast_node.rs b/src/iter/ast_node.rs
index 0704ab3..4ff87a6 100644
--- a/src/iter/ast_node.rs
+++ b/src/iter/ast_node.rs
@@ -9,6 +9,7 @@ use crate::types::Comment;
 use crate::types::CommentBlock;
 use crate::types::DiarySexp;
 use crate::types::Document;
+use crate::types::DocumentElement;
 use crate::types::Drawer;
 use crate::types::DynamicBlock;
 use crate::types::Element;
@@ -111,6 +112,15 @@ pub enum AstNode<'r, 's> {
     Timestamp(&'r Timestamp<'s>),
 }
 
+impl<'r, 's> From<&'r DocumentElement<'s>> for AstNode<'r, 's> {
+    fn from(value: &'r DocumentElement<'s>) -> Self {
+        match value {
+            DocumentElement::Heading(inner) => inner.into(),
+            DocumentElement::Section(inner) => inner.into(),
+        }
+    }
+}
+
 impl<'r, 's> From<&'r Element<'s>> for AstNode<'r, 's> {
     fn from(value: &'r Element<'s>) -> Self {
         match value {
diff --git a/src/iter/ast_node_iter.rs b/src/iter/ast_node_iter.rs
index 39d5b58..dc181da 100644
--- a/src/iter/ast_node_iter.rs
+++ b/src/iter/ast_node_iter.rs
@@ -5,8 +5,12 @@ use super::macros::children_iter;
 use super::macros::empty_iter;
 use crate::types::Bold;
 use crate::types::Code;
+use crate::types::Heading;
 use crate::types::Italic;
 use crate::types::Object;
+use crate::types::PlainText;
+use crate::types::RadioLink;
+use crate::types::RegularLink;
 use crate::types::StrikeThrough;
 use crate::types::Underline;
 use crate::types::Verbatim;
@@ -19,7 +23,7 @@ use crate::types::Verbatim;
 pub enum AstNodeIter<'r, 's> {
     // Document Nodes
     // Document(DocumentIter<'r, 's>),
-    // Heading(HeadingIter<'r, 's>),
+    Heading(HeadingIter<'r, 's>),
     // Section(SectionIter<'r, 's>),
     // Elements
     // Paragraph(ParagraphIter<'r, 's>),
@@ -51,9 +55,9 @@ pub enum AstNodeIter<'r, 's> {
     StrikeThrough(StrikeThroughIter<'r, 's>),
     Code(CodeIter<'r, 's>),
     Verbatim(VerbatimIter<'r, 's>),
-    // PlainText(PlainTextIter<'r, 's>),
-    // RegularLink(RegularLinkIter<'r, 's>),
-    // RadioLink(RadioLinkIter<'r, 's>),
+    PlainText(PlainTextIter<'r, 's>),
+    RegularLink(RegularLinkIter<'r, 's>),
+    RadioLink(RadioLinkIter<'r, 's>),
     // RadioTarget(RadioTargetIter<'r, 's>),
     // PlainLink(PlainLinkIter<'r, 's>),
     // AngleLink(AngleLinkIter<'r, 's>),
@@ -74,11 +78,11 @@ pub enum AstNodeIter<'r, 's> {
     // Timestamp(TimestampIter<'r, 's>),
 }
 
-pub struct BoldIter<'r, 's> {
+pub struct HeadingIter<'r, 's> {
     next: std::slice::Iter<'r, Object<'s>>,
 }
 
-impl<'r, 's> Iterator for BoldIter<'r, 's> {
+impl<'r, 's> Iterator for HeadingIter<'r, 's> {
     type Item = AstNode<'r, 's>;
 
     fn next(&mut self) -> Option<Self::Item> {
@@ -86,18 +90,23 @@ impl<'r, 's> Iterator for BoldIter<'r, 's> {
     }
 }
 
-impl<'r, 's> IntoIterator for &'r Bold<'s> {
+impl<'r, 's> IntoIterator for &'r Heading<'s> {
     type Item = AstNode<'r, 's>;
 
-    type IntoIter = BoldIter<'r, 's>;
+    type IntoIter = HeadingIter<'r, 's>;
 
     fn into_iter(self) -> Self::IntoIter {
-        BoldIter {
-            next: self.children.iter(),
+        HeadingIter {
+            next: self
+                .title
+                .iter()
+                .map(Into::<AstNode>::into)
+                .chain(self.children.iter().map(Into::<AstNode>::into)),
         }
     }
 }
 
+children_iter!(Bold<'s>, BoldIter, std::slice::Iter<'r, Object<'s>>);
 children_iter!(Italic<'s>, ItalicIter, std::slice::Iter<'r, Object<'s>>);
 children_iter!(
     Underline<'s>,
@@ -109,29 +118,12 @@ children_iter!(
     StrikeThroughIter,
     std::slice::Iter<'r, Object<'s>>
 );
-
-pub struct CodeIter<'r, 's> {
-    phantom: PhantomData<&'r Code<'s>>,
-}
-
-impl<'r, 's> Iterator for CodeIter<'r, 's> {
-    type Item = AstNode<'r, 's>;
-
-    fn next(&mut self) -> Option<Self::Item> {
-        None
-    }
-}
-
-impl<'r, 's> IntoIterator for &'r Code<'s> {
-    type Item = AstNode<'r, 's>;
-
-    type IntoIter = CodeIter<'r, 's>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        CodeIter {
-            phantom: PhantomData,
-        }
-    }
-}
-
+empty_iter!(Code<'s>, CodeIter);
 empty_iter!(Verbatim<'s>, VerbatimIter);
+empty_iter!(PlainText<'s>, PlainTextIter);
+empty_iter!(RegularLink<'s>, RegularLinkIter);
+children_iter!(
+    RadioLink<'s>,
+    RadioLinkIter,
+    std::slice::Iter<'r, Object<'s>>
+);