From 282417ee944948445e1fced8334f25478ca91c55 Mon Sep 17 00:00:00 2001
From: Tom Alexander <tom@fizz.buzz>
Date: Wed, 27 Sep 2023 18:00:30 -0400
Subject: [PATCH] Implementing HeadingIter but I do not think it can be generic
 enough for a macro.

Hopefully most types won't need so much care.
---
 src/iter/ast_node_iter.rs | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/src/iter/ast_node_iter.rs b/src/iter/ast_node_iter.rs
index dc181da4..ee92efa5 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::Item> {
-        self.next.next().map(Into::<AstNode>::into)
+        self.title_next
+            .next()
+            .map(Into::<AstNode>::into)
+            .or_else(|| self.children_next.next().map(Into::<AstNode>::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::<AstNode>::into)
-                .chain(self.children.iter().map(Into::<AstNode>::into)),
+            title_next: self.title.iter(),
+            children_next: self.children.iter(),
         }
     }
 }