From ee5ed17c20f3ef157439d41453a324643ac7e4b8 Mon Sep 17 00:00:00 2001
From: Tom Alexander <tom@fizz.buzz>
Date: Tue, 10 Oct 2023 15:19:42 -0400
Subject: [PATCH] Implement ExactSizeIterator for DocumentIter.

---
 src/compare/diff.rs       | 19 ++++++++-----------
 src/compare/util.rs       | 10 ++--------
 src/iter/ast_node_iter.rs | 37 +++++++++++++++++++++++++++++--------
 3 files changed, 39 insertions(+), 27 deletions(-)

diff --git a/src/compare/diff.rs b/src/compare/diff.rs
index 154bf31..838b0e2 100644
--- a/src/compare/diff.rs
+++ b/src/compare/diff.rs
@@ -447,17 +447,14 @@ fn _compare_document<'b, 's>(
     let mut child_status = Vec::new();
     let mut message = None;
 
-    // compare_children_iter(
-    //     source,
-    //     emacs,
-    //     rust.zeroth_section
-    //         .iter()
-    //         .map(Into::<AstNode>::into)
-    //         .chain(rust.children.iter().map(Into::<AstNode>::into)),
-    //     &mut child_status,
-    //     &mut this_status,
-    //     &mut message,
-    // )?;
+    compare_children_iter(
+        source,
+        emacs,
+        rust.into_iter(),
+        &mut child_status,
+        &mut this_status,
+        &mut message,
+    )?;
 
     for diff in compare_properties!(
         source,
diff --git a/src/compare/util.rs b/src/compare/util.rs
index e3e02d5..47dd8e5 100644
--- a/src/compare/util.rs
+++ b/src/compare/util.rs
@@ -284,13 +284,7 @@ where
     Ok(())
 }
 
-pub(crate) fn compare_children_iter<
-    'b,
-    's,
-    'x,
-    RC: 'x,
-    RI: Iterator<Item = &'x RC> + ExactSizeIterator,
->(
+pub(crate) fn compare_children_iter<'b, 's, RC, RI: Iterator<Item = RC> + ExactSizeIterator>(
     source: &'s str,
     emacs: &'b Token<'s>,
     rust_children: RI,
@@ -299,7 +293,7 @@ pub(crate) fn compare_children_iter<
     message: &mut Option<String>,
 ) -> Result<(), Box<dyn std::error::Error>>
 where
-    AstNode<'b, 's>: From<&'x RC>,
+    AstNode<'b, 's>: From<RC>,
 {
     let emacs_children = emacs.as_list()?;
     let emacs_children_length = emacs_children.len() - 2;
diff --git a/src/iter/ast_node_iter.rs b/src/iter/ast_node_iter.rs
index 1c4210d..851666e 100644
--- a/src/iter/ast_node_iter.rs
+++ b/src/iter/ast_node_iter.rs
@@ -202,14 +202,35 @@ impl<'r, 's> AstNodeIter<'r, 's> {
     }
 }
 
-multi_field_iter!(
-    Document<'s>,
-    DocumentIter,
-    zeroth_section,
-    std::option::Iter<'r, Section<'s>>,
-    children,
-    std::slice::Iter<'r, Heading<'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!(
     Heading<'s>,
     HeadingIter,