From c578bb45af3c2b1c5e7f302f6adcfbff4891bf48 Mon Sep 17 00:00:00 2001
From: Tom Alexander <tom@fizz.buzz>
Date: Tue, 10 Oct 2023 15:27:03 -0400
Subject: [PATCH] Implement ExactSizeIterator for the other node types.

---
 src/iter/ast_node_iter.rs | 37 ++++++++-----------------------------
 src/iter/macros.rs        | 13 +++++++++++++
 2 files changed, 21 insertions(+), 29 deletions(-)

diff --git a/src/iter/ast_node_iter.rs b/src/iter/ast_node_iter.rs
index 851666e..1c4210d 100644
--- a/src/iter/ast_node_iter.rs
+++ b/src/iter/ast_node_iter.rs
@@ -202,35 +202,14 @@ impl<'r, 's> AstNodeIter<'r, 's> {
     }
 }
 
-pub struct DocumentIter<'r, 's> {
-    zeroth_section: std::option::Iter<'r, Section<'s>>,
-    children: 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::Item> {
-        self.zeroth_section
-            .next()
-            .map(Into::<AstNode>::into)
-            .or_else(|| self.children.next().map(Into::<AstNode>::into))
-    }
-
-    fn size_hint(&self) -> (usize, Option<usize>) {
-        let size = self.zeroth_section.len() + self.children.len();
-        (size, Some(size))
-    }
-}
-impl<'r, 's> ExactSizeIterator for DocumentIter<'r, 's> {}
-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: self.zeroth_section.iter(),
-            children: self.children.iter(),
-        }
-    }
-}
+multi_field_iter!(
+    Document<'s>,
+    DocumentIter,
+    zeroth_section,
+    std::option::Iter<'r, Section<'s>>,
+    children,
+    std::slice::Iter<'r, Heading<'s>>
+);
 multi_field_iter!(
     Heading<'s>,
     HeadingIter,
diff --git a/src/iter/macros.rs b/src/iter/macros.rs
index 18a265d..3f9333b 100644
--- a/src/iter/macros.rs
+++ b/src/iter/macros.rs
@@ -11,8 +11,15 @@ macro_rules! children_iter {
             fn next(&mut self) -> Option<Self::Item> {
                 self.next.next().map(Into::<AstNode>::into)
             }
+
+            fn size_hint(&self) -> (usize, Option<usize>) {
+                let size = self.next.len();
+                (size, Some(size))
+            }
         }
 
+        impl<'r, 's> ExactSizeIterator for $itertype<'r, 's> {}
+
         impl<'r, 's> IntoIterator for &'r $astnodetype {
             type Item = AstNode<'r, 's>;
 
@@ -42,8 +49,14 @@ macro_rules! empty_iter {
             fn next(&mut self) -> Option<Self::Item> {
                 None
             }
+
+            fn size_hint(&self) -> (usize, Option<usize>) {
+                (0, Some(0))
+            }
         }
 
+        impl<'r, 's> ExactSizeIterator for $itertype<'r, 's> {}
+
         impl<'r, 's> IntoIterator for &'r $astnodetype {
             type Item = AstNode<'r, 's>;