init

.gitignore

@@ -0,0 +1,2 @@
+/target
+**/*.rs.bk

Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "mdbook-toc"
+version = "0.1.0"
+authors = ["Jan-Erik Rediger <janerik@fnordig.de>"]
+
+[dependencies]
+mdbook = "0.1.8"
+pulldown-cmark = "0.1.2"
+pulldown-cmark-to-cmark = "1.1.0"
+env_logger = "0.5.10"
+log = "0.4.3"
+clap = "2.32.0"

src/lib.rs

@@ -0,0 +1,101 @@
+extern crate mdbook;
+extern crate pulldown_cmark;
+extern crate pulldown_cmark_to_cmark;
+
+use std::borrow::Cow;
+use mdbook::errors::{Error, Result};
+use mdbook::book::{Book, BookItem, Chapter};
+use mdbook::preprocess::{Preprocessor, PreprocessorContext};
+use pulldown_cmark::{Event, Parser};
+use pulldown_cmark::Tag::*;
+use pulldown_cmark_to_cmark::fmt::cmark;
+
+pub struct Toc;
+
+impl Preprocessor for Toc {
+    fn name(&self) -> &str {
+        "toc"
+    }
+
+    fn run(&self, _ctx: &PreprocessorContext, book: &mut Book) -> Result<()> {
+        let mut res: Option<_> = None;
+        book.for_each_mut(|item: &mut BookItem| {
+            if let Some(Err(_)) = res {
+                return;
+            }
+            if let BookItem::Chapter(ref mut chapter) = *item {
+                res = Some(Toc::add_toc(chapter).map(|md| {
+                    chapter.content = md;
+                }));
+            }
+        });
+        res.unwrap_or(Ok(()))
+    }
+}
+
+fn build_toc<'a>(toc: &[(i32, Cow<'a, str>)]) -> String {
+    let mut result = String::new();
+
+    for (level, name) in toc {
+        let width = 2*(level-1) as usize;
+        let slug = mdbook::utils::normalize_id(&name);
+        let entry = format!("{1:0$}* [{2}](#{3})\n", width, "", name, slug);
+        result.push_str(&entry);
+    }
+
+    result
+}
+
+impl Toc {
+    fn add_toc(chapter: &mut Chapter) -> Result<String> {
+        let mut buf = String::with_capacity(chapter.content.len());
+        let mut toc_found = false;
+
+        let mut toc_content = vec![];
+        let mut current_header_level : Option<i32> = None;
+
+        for e in Parser::new(&chapter.content) {
+            if let Event::Html(html) = e {
+                if html == "<!-- toc -->\n" {
+                    toc_found = true;
+                }
+                continue;
+            }
+            if !toc_found {
+                continue;
+            }
+
+            if let Event::Start(Header(lvl)) = e {
+                current_header_level = Some(lvl);
+                continue;
+            }
+            if let Event::End(Header(_)) = e {
+                current_header_level = None;
+                continue;
+            }
+            if current_header_level.is_none() {
+                continue;
+            }
+
+            if let Event::Text(header) = e {
+                toc_content.push((current_header_level.unwrap(), header));
+            }
+        }
+
+        let toc_events = build_toc(&toc_content);
+        let toc_events = Parser::new(&toc_events).collect::<Vec<_>>();
+
+        let events = Parser::new(&chapter.content).map(|e| {
+            if let Event::Html(html) = e.clone() {
+                if html == "<!-- toc -->\n" {
+                    return toc_events.clone();
+                }
+            }
+            vec![e]
+        }).flat_map(|e| e);
+
+        cmark(events, &mut buf, None)
+            .map(|_| buf)
+            .map_err(|err| Error::from(format!("Markdown serialization failed: {}", err)))
+    }
+}

src/main.rs

@@ -0,0 +1,64 @@
+extern crate env_logger;
+extern crate mdbook;
+extern crate mdbook_toc;
+extern crate clap;
+
+use mdbook::MDBook;
+use mdbook::errors::Result;
+use mdbook_toc::Toc;
+use clap::{App, ArgMatches};
+
+use std::env;
+use std::process;
+use std::path::{Path, PathBuf};
+
+fn get_book_dir(args: &ArgMatches) -> PathBuf {
+    if let Some(dir) = args.value_of("dir") {
+        // Check if path is relative from current dir, or absolute...
+        let p = Path::new(dir);
+        if p.is_relative() {
+            env::current_dir().unwrap().join(dir)
+        } else {
+            p.to_path_buf()
+        }
+    } else {
+        env::current_dir().expect("Unable to determine the current directory")
+    }
+}
+
+pub fn make_app<'a, 'b>() -> App<'a, 'b> {
+    App::new("mdbook-toc")
+        .about("Build the book from the markdown files with ToC support")
+        .arg_from_usage(
+            "-d, --dest-dir=[dest-dir] 'The output directory for your book{n}(Defaults to ./book \
+             when omitted)'",
+        )
+        .arg_from_usage(
+            "[dir] 'A directory for your book{n}(Defaults to Current Directory when omitted)'",
+        )
+}
+
+pub fn execute(args: &ArgMatches) -> Result<()> {
+    let book_dir = get_book_dir(args);
+    let mut book = MDBook::load(&book_dir)?;
+
+    if let Some(dest_dir) = args.value_of("dest-dir") {
+        book.config.build.build_dir = PathBuf::from(dest_dir);
+    }
+
+    book.with_preprecessor(Toc);
+    book.build()?;
+
+    Ok(())
+}
+
+fn main() {
+    env_logger::init();
+    let app = make_app();
+    let matches = app.get_matches();
+
+    if let Err(e) = execute(&matches) {
+        eprintln!("{}", e);
+        process::exit(1);
+    }
+}