1 files changed, 68 insertions, 4 deletions

diff --git a/src/directive/toc.rs b/src/directive/toc.rs
index e00cb24..f027931 100644
--- a/src/directive/toc.rs
+++ b/src/directive/toc.rs
@@ -1,22 +1,86 @@
 use crate::directive::{DirectiveImplementation, Processed};
 use crate::error::SiteError;
+use crate::html::{Content, Element, ElementTag};
 use crate::page::PageMeta;
 use crate::site::Site;
 use crate::wikitext::ParsedDirective;
 
 #[derive(Debug, Default, Eq, PartialEq)]
-pub struct Toc {}
+pub struct Toc {
+    levels: String,
+}
 
 impl DirectiveImplementation for Toc {
     const REQUIRED: &'static [&'static str] = &[];
     const ALLOWED: &'static [&'static str] = &["levels"];
     const ALLOW_ANY_UNNAMED: bool = true;
 
-    fn from_parsed(_: &ParsedDirective) -> Self {
-        Self::default()
+    fn from_parsed(p: &ParsedDirective) -> Self {
+        let args = p.args();
+        let levels = args.get("levels").unwrap_or(&"9");
+        Self::new(levels.to_string())
     }
 
     fn process(&self, _site: &Site, _meta: &mut PageMeta) -> Result<Processed, SiteError> {
-        Err(SiteError::UnimplementedDirective("toc".into()))
+        let levels: usize = self
+            .levels
+            .parse()
+            .map_err(|e| SiteError::LevelsParse(self.levels.clone(), e))?;
+        Ok(Processed::Toc(levels))
+    }
+}
+
+impl Toc {
+    fn new(levels: String) -> Self {
+        Self { levels }
+    }
+
+    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 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");
+            }
+            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");
+        }
+        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),
+            }
+        }
     }
 }