From c44e7d642f22e78dc46b4cbcbf83cd0c2e2910a9 Mon Sep 17 00:00:00 2001
From: Tom Alexander <tom@fizz.buzz>
Date: Mon, 24 Apr 2023 22:10:24 -0400
Subject: [PATCH] Starting code for iterating over the parsed ast.

---
 src/parser/document.rs | 29 +++++++++++++++++++++++++++++
 src/parser/mod.rs      |  1 +
 src/parser/token.rs    | 12 ++++++++++++
 3 files changed, 42 insertions(+)
 create mode 100644 src/parser/token.rs

diff --git a/src/parser/document.rs b/src/parser/document.rs
index 9954fab..98e570c 100644
--- a/src/parser/document.rs
+++ b/src/parser/document.rs
@@ -18,6 +18,7 @@ use super::element::Element;
 use super::object::Object;
 use super::parser_with_context::parser_with_context;
 use super::source::Source;
+use super::token::Token;
 use super::util::exit_matcher_parser;
 use super::util::get_consumed;
 use super::util::start_of_line;
@@ -255,3 +256,31 @@ fn headline<'r, 's>(
 fn headline_end<'r, 's>(_context: Context<'r, 's>, input: &'s str) -> Res<&'s str, &'s str> {
     line_ending(input)
 }
+
+impl<'s> Document<'s> {
+    pub fn iter_tokens<'r>(&'r self) -> impl Iterator<Item = Token<'r, 's>> {
+        self.zeroth_section
+            .iter()
+            .map(Token::Section)
+            .chain(self.children.iter().map(Token::Heading))
+    }
+}
+
+impl<'s> Heading<'s> {
+    pub fn iter_tokens<'r>(&'r self) -> impl Iterator<Item = Token<'r, 's>> {
+        self.title.iter().map(Token::Object).chain(self.children.iter().map(
+            |de| {
+                match de {
+                    DocumentElement::Heading(obj) => Token::Heading(obj),
+                    DocumentElement::Section(obj) => Token::Section(obj),
+                }
+            }
+        ))
+    }
+}
+
+impl<'s> Section<'s> {
+    pub fn iter_tokens<'r>(&'r self) -> impl Iterator<Item = Token<'r, 's>> {
+        self.children.iter().map(Token::Element)
+    }
+}
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index df09653..dcd61c6 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -35,6 +35,7 @@ pub mod sexp;
 mod source;
 mod table;
 mod text_markup;
+mod token;
 mod util;
 pub use document::document;
 pub use document::Document;
diff --git a/src/parser/token.rs b/src/parser/token.rs
new file mode 100644
index 0000000..6685773
--- /dev/null
+++ b/src/parser/token.rs
@@ -0,0 +1,12 @@
+use super::Document;
+use super::Element;
+use super::Heading;
+use super::Object;
+use super::Section;
+
+pub enum Token<'r, 's> {
+    Heading(&'r Heading<'s>),
+    Section(&'r Section<'s>),
+    Object(&'r Object<'s>),
+    Element(&'r Element<'s>),
+}