1 files changed, 76 insertions, 40 deletions

diff --git a/src/directive/toc.rs b/src/directive/toc.rs
index fa115b6..c4ffc9e 100644
--- a/src/directive/toc.rs
+++ b/src/directive/toc.rs
@@ -1,8 +1,41 @@
 use crate::directive::{DirectiveError, DirectiveImplementation, Processed};
-use crate::html::{Content, Element, ElementTag};
 use crate::page::PageMeta;
 use crate::site::Site;
 use crate::wikitext::ParsedDirective;
+use html_page::{Element, Tag, Visitor};
+use std::cmp::Ordering;
+
+#[derive(Debug, Default)]
+struct Headings {
+    headings: Vec<(usize, Element)>,
+}
+
+impl Visitor for Headings {
+    fn visit_text(&mut self, _: &str) {}
+    fn visit_element(&mut self, e: &Element) {
+        match e.tag() {
+            Tag::H1 => self.headings.push((1, e.clone())),
+            Tag::H2 => self.headings.push((2, e.clone())),
+            Tag::H3 => self.headings.push((3, e.clone())),
+            Tag::H4 => self.headings.push((4, e.clone())),
+            Tag::H5 => self.headings.push((5, e.clone())),
+            Tag::H6 => self.headings.push((6, e.clone())),
+            _ => (),
+        }
+    }
+}
+
+#[derive(Debug, Default)]
+struct Text {
+    text: String,
+}
+
+impl Visitor for Text {
+    fn visit_text(&mut self, s: &str) {
+        self.text.push_str(s);
+    }
+    fn visit_element(&mut self, _: &Element) {}
+}
 
 #[derive(Debug, Default, Eq, PartialEq)]
 pub struct Toc {
@@ -35,51 +68,54 @@ impl Toc {
     }
 
     pub fn post_process(html: &Element, levels: usize) -> String {
-        let headings: Vec<(usize, &[Content])> = html
-            .children()
-            .iter()
-            .filter_map(|c| match c {
-                Content::Elt(e) => Some(e),
-                _ => None,
-            })
-            .filter_map(|e| match e.tag() {
-                ElementTag::H1 => Some((1, e.children())),
-                ElementTag::H2 => Some((2, e.children())),
-                ElementTag::H3 => Some((3, e.children())),
-                ElementTag::H4 => Some((4, e.children())),
-                ElementTag::H5 => Some((5, e.children())),
-                ElementTag::H6 => Some((6, e.children())),
-                _ => None,
-            })
-            .collect();
-        let mut html = String::new();
+        let mut headings = Headings::default();
+        headings.visit(html);
+
+        let mut stack = vec![];
         let mut prev_level: usize = 0;
-        for (level, text) in headings.iter() {
-            if *level > levels {
-                continue;
-            } else if *level > prev_level {
-                html.push_str("<ol>");
-            } else if *level < prev_level {
-                html.push_str("</ol>\n");
+        for (level, heading) in headings
+            .headings
+            .iter()
+            .filter(|(level, _)| *level < levels)
+        {
+            match level.cmp(&prev_level) {
+                Ordering::Greater => {
+                    let list = Element::new(Tag::Ul);
+                    stack.push(list);
+                }
+                Ordering::Less => {
+                    assert!(!stack.is_empty());
+                    let ending = stack.pop().unwrap();
+
+                    assert!(!stack.is_empty());
+                    let parent = stack.last_mut().unwrap();
+                    parent.push_child(&ending);
+                }
+                Ordering::Equal => (),
+            }
+            if let Some(ol) = stack.last_mut() {
+                let mut text = Text::default();
+                text.visit(heading);
+                let li = Element::new(Tag::Li).with_text(&text.text);
+                ol.push_child(&li);
             }
-            html.push_str("<li>");
-            Self::stringify(&mut html, text);
-            html.push_str("</li>\n");
             prev_level = *level;
         }
-        for _ in 0..prev_level {
-            html.push_str("</ol>\n");
+
+        while stack.len() > 1 {
+            assert!(!stack.is_empty());
+            let ending = stack.pop().unwrap();
+
+            assert!(!stack.is_empty());
+            let parent = stack.last_mut().unwrap();
+            parent.push_child(&ending);
         }
-        html
-    }
 
-    fn stringify(buf: &mut String, bits: &[Content]) {
-        for c in bits.iter() {
-            match c {
-                Content::Text(s) => buf.push_str(s),
-                Content::Elt(e) => Self::stringify(buf, e.children()),
-                Content::Html(h) => buf.push_str(h),
-            }
+        let mut toc =
+            Element::new(Tag::Div).with_child(Element::new(Tag::H1).with_text("Contents"));
+        if let Some(e) = stack.get(0) {
+            toc.push_child(e);
         }
+        format!("{}", toc)
     }
 }