From 0073af19e283910c8ad337329fdcb256f327d2bf Mon Sep 17 00:00:00 2001
From: Tom Alexander <tom@fizz.buzz>
Date: Fri, 14 Jul 2023 19:54:41 -0400
Subject: [PATCH] Running into an issue returning different iterators from the
 same function.

---
 src/parser/element.rs |  1 +
 src/parser/token.rs   | 27 +++++++++++++++++++++++++++
 toy_language.txt      |  4 +---
 3 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/src/parser/element.rs b/src/parser/element.rs
index 3052ece..27efa21 100644
--- a/src/parser/element.rs
+++ b/src/parser/element.rs
@@ -20,6 +20,7 @@ use super::lesser_element::SrcBlock;
 use super::lesser_element::VerseBlock;
 use super::source::SetSource;
 use super::source::Source;
+use super::token::Token;
 use super::Drawer;
 
 #[derive(Debug)]
diff --git a/src/parser/token.rs b/src/parser/token.rs
index 6685773..7f02a93 100644
--- a/src/parser/token.rs
+++ b/src/parser/token.rs
@@ -3,10 +3,37 @@ use super::Element;
 use super::Heading;
 use super::Object;
 use super::Section;
+use crate::parser::DocumentElement;
 
 pub enum Token<'r, 's> {
+    Document(&'r Document<'s>),
     Heading(&'r Heading<'s>),
     Section(&'r Section<'s>),
     Object(&'r Object<'s>),
     Element(&'r Element<'s>),
 }
+
+impl<'r, 's> Token<'r, 's> {
+    pub fn iter_tokens(&self) -> impl Iterator<Item = Token<'r, 's>> {
+        match self {
+            Token::Document(document) => document
+                .zeroth_section
+                .iter()
+                .map(Token::Section)
+                .chain(document.children.iter().map(Token::Heading)),
+            Token::Heading(heading) => {
+                heading
+                    .title
+                    .iter()
+                    .map(Token::Object)
+                    .chain(heading.children.iter().map(|de| match de {
+                        DocumentElement::Heading(ref obj) => Token::Heading(obj),
+                        DocumentElement::Section(ref obj) => Token::Section(obj),
+                    }))
+            }
+            Token::Section(section) => section.children.iter().map(Token::Element),
+            Token::Object(_) => panic!(),
+            Token::Element(_) => panic!(),
+        }
+    }
+}
diff --git a/toy_language.txt b/toy_language.txt
index 64ee871..a29169b 100644
--- a/toy_language.txt
+++ b/toy_language.txt
@@ -1,3 +1 @@
-foo *bar /baz *lorem* ipsum/ dolar* alpha
-
-foo *bar /baz _lorem_ ipsum/ dolar* alpha
+foo <<<*bar* baz>>> lorem ipsum *bar* baz dolar.