Allow custom ToC markers through configuration

2 files changed, 218 insertions, 31 deletions

diff --git a/README.md b/README.md
index 2387d7f..be735c9 100644
--- a/README.md
+++ b/README.md
@@ -6,7 +6,7 @@ A preprocessor for [mdbook][] to add inline Table of Contents support.
 
 It turns this:
 
-```
+```md
 <!-- toc -->
 ```
 
@@ -16,13 +16,13 @@ into a Table of Contents based on all top- and second-level headings of the chap
 
 If you want to use only this preprocessor, install the tool:
 
-```
+```sh
 cargo install mdbook-toc
 ```
 
 Add it as a preprocessor to your `book.toml`:
 
-```
+```toml
 [preprocessor.toc]
 command = "mdbook-toc"
 renderer = ["html"]
@@ -30,10 +30,47 @@ renderer = ["html"]
 
 Finally, build your book as normal:
 
-```
+```sh
 mdbook path/to/book
 ```
 
+## Custom TOC marker
+
+The default marker is:
+
+```md
+<!-- toc -->
+```
+
+If you wish to use a different, such as the GitLab marker `[[_TOC_]]`, you must add the following settings to your `book.toml`.
+
+```toml
+[preprocessor.toc]
+marker = "[[_TOC_]]"
+```
+
+You can also use multi-line markers such as the GitHub marker, which is:
+
+```md
+* auto-gen TOC;
+{:toc}
+```
+
+Configure the string with a newline:
+
+```toml
+[preprocessor.toc]
+marker = "* auto-gen TOC;\n{:toc}"
+```
+
+or with multi-line strings:
+
+```toml
+[preprocessor.toc]
+marker = """* auto-gen TOC;
+{:toc}"""
+```
+
 ## License
 
 MPL. See [LICENSE](LICENSE).  
diff --git a/src/lib.rs b/src/lib.rs
index 0fbe50c..a6c3f50 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -10,20 +10,40 @@ use pulldown_cmark_to_cmark::{cmark_with_options, Options as COptions};
 
 pub struct Toc;
 
+static DEFAULT_MARKER: &str = "<!-- toc -->\n";
+
 impl Preprocessor for Toc {
     fn name(&self) -> &str {
         "toc"
     }
 
-    fn run(&self, _ctx: &PreprocessorContext, mut book: Book) -> Result<Book> {
+    fn run(&self, ctx: &PreprocessorContext, mut book: Book) -> Result<Book> {
         let mut res = None;
+        let toc_marker = if let Some(cfg) = ctx.config.get_preprocessor(self.name()) {
+            if let Some(marker) = cfg.get("marker") {
+                match marker.as_str() {
+                    Some(m) => m,
+                    None => {
+                        return Err(Error::msg(format!(
+                            "Marker {:?} is not a valid string",
+                            marker
+                        )))
+                    }
+                }
+            } else {
+                DEFAULT_MARKER
+            }
+        } else {
+            DEFAULT_MARKER
+        };
+
         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| {
+                res = Some(Toc::add_toc(chapter, &toc_marker).map(|md| {
                     chapter.content = md;
                 }));
             }
@@ -69,7 +89,7 @@ fn build_toc(toc: &[(u32, String, String)]) -> String {
     result
 }
 
-fn add_toc(content: &str) -> Result<String> {
+fn add_toc(content: &str, marker: &str) -> Result<String> {
     let mut buf = String::with_capacity(content.len());
     let mut toc_found = false;
 
@@ -84,17 +104,29 @@ fn add_toc(content: &str) -> Result<String> {
     opts.insert(Options::ENABLE_STRIKETHROUGH);
     opts.insert(Options::ENABLE_TASKLISTS);
 
+    let mark: Vec<Event> = Parser::new(marker).collect();
+    let mut mark_start = -1;
+    let mut mark_loc = 0;
+    let mut c = -1;
+
     for e in Parser::new_ext(&content, opts) {
+        c += 1;
         log::trace!("Event: {:?}", e);
-
-        if let Event::Html(html) = e {
-            if &*html == "<!-- toc -->\n" {
-                toc_found = true;
-            }
-            continue;
-        }
         if !toc_found {
-            continue;
+            if e == mark[mark_loc] {
+                if mark_start == -1 {
+                    mark_start = c;
+                }
+                mark_loc += 1;
+                if mark_loc >= mark.len() {
+                    toc_found = true
+                }
+            } else if mark_loc > 0 {
+                mark_loc = 0;
+                mark_start = -1;
+            } else {
+                continue;
+            }
         }
 
         if let Event::Start(Heading(lvl)) = e {
@@ -138,14 +170,23 @@ fn add_toc(content: &str) -> Result<String> {
     let toc = build_toc(&toc_content);
     let toc_events = Parser::new(&toc).collect::<Vec<_>>();
 
+    let mut c = -1;
     let events = Parser::new_ext(&content, opts)
         .map(|e| {
-            if let Event::Html(html) = e.clone() {
-                if &*html == "<!-- toc -->\n" {
-                    return toc_events.clone();
-                }
+            c += 1;
+            if c > mark_start && c < mark_start + (mark.len() as i32) {
+                vec![]
+            } else if c == mark_start {
+                toc_events.clone()
+            } else {
+                vec![e]
             }
-            vec![e]
+            // if let Event::Html(html) = e.clone() {
+            //     if &*html == marker {
+            //         return toc_events.clone();
+            //     }
+            // }
+            // vec![e]
         })
         .flatten();
 
@@ -156,14 +197,14 @@ fn add_toc(content: &str) -> Result<String> {
 }
 
 impl Toc {
-    fn add_toc(chapter: &mut Chapter) -> Result<String> {
-        add_toc(&chapter.content)
+    fn add_toc(chapter: &mut Chapter, marker: &str) -> Result<String> {
+        add_toc(&chapter.content, marker)
     }
 }
 
 #[cfg(test)]
 mod test {
-    use super::add_toc;
+    use super::{add_toc, DEFAULT_MARKER};
     use pretty_assertions::assert_eq;
 
     #[test]
@@ -207,7 +248,10 @@ mod test {
 
 ### Header 2.2.1"#;
 
-        assert_eq!(expected, add_toc(content).unwrap());
+        assert_eq!(
+            expected,
+            add_toc(content, DEFAULT_MARKER).unwrap()
+        );
     }
 
     #[test]
@@ -240,7 +284,10 @@ mod test {
 
 ## Header 2.1"#;
 
-        assert_eq!(expected, add_toc(content).unwrap());
+        assert_eq!(
+            expected,
+            add_toc(content, DEFAULT_MARKER).unwrap()
+        );
     }
 
     #[test]
@@ -262,7 +309,10 @@ mod test {
 |------|------|
 |Row 1|Row 2|"#;
 
-        assert_eq!(expected, add_toc(content).unwrap());
+        assert_eq!(
+            expected,
+            add_toc(content, DEFAULT_MARKER).unwrap()
+        );
     }
 
     #[test]
@@ -311,7 +361,10 @@ mod test {
 
 # Another header `with inline` code"#;
 
-        assert_eq!(expected, add_toc(content).unwrap());
+        assert_eq!(
+            expected,
+            add_toc(content, DEFAULT_MARKER).unwrap()
+        );
     }
 
     #[test]
@@ -358,7 +411,10 @@ mod test {
 
 ## User Preferences"#;
 
-        assert_eq!(expected, add_toc(content).unwrap());
+        assert_eq!(
+            expected,
+            add_toc(content, DEFAULT_MARKER).unwrap()
+        );
     }
 
     #[test]
@@ -375,7 +431,7 @@ mod test {
 
 text"#;
 
-    let expected = r#"# Heading
+        let expected = r#"# Heading
 
 * [Level 1.1](#level-11)
   * [Level 1.1.1](#level-111)
@@ -395,7 +451,55 @@ text"#;
 
 text"#;
 
-        assert_eq!(expected, add_toc(content).unwrap());
+        assert_eq!(
+            expected,
+            add_toc(content, DEFAULT_MARKER).unwrap()
+        );
+    }
+
+    #[test]
+    fn add_toc_with_gitlab_marker() {
+        let marker = "[[_TOC_]]".to_owned();
+        let content = r#"# Chapter
+
+[[_TOC_]]
+
+# Header 1
+
+## Header 1.1
+
+# Header 2
+
+## Header 2.1
+
+## Header 2.2
+
+### Header 2.2.1
+
+"#;
+
+        let expected = r#"# Chapter
+
+* [Header 1](#header-1)
+  * [Header 1.1](#header-11)
+* [Header 2](#header-2)
+  * [Header 2.1](#header-21)
+  * [Header 2.2](#header-22)
+    * [Header 2.2.1](#header-221)
+
+# Header 1
+
+## Header 1.1
+
+# Header 2
+
+## Header 2.1
+
+## Header 2.2
+
+### Header 2.2.1"#;
+
+        assert_eq!(expected, add_toc(content, &marker).unwrap());
     }
 
     #[test]
@@ -431,6 +535,52 @@ text"#;
 
 ## Duplicate"#;
 
-        assert_eq!(expected, add_toc(content).unwrap());
+        assert_eq!(expected, add_toc(content, DEFAULT_MARKER).unwrap());
+    }
+
+    #[test]
+    fn add_toc_with_github_marker() {
+        let marker = "* auto-gen TOC:\n{:toc}".to_owned();
+        let content = r#"# Chapter
+
+* auto-gen TOC:
+{:toc}
+
+# Header 1
+
+## Header 1.1
+
+# Header 2
+
+## Header 2.1
+
+## Header 2.2
+
+### Header 2.2.1
+
+"#;
+
+        let expected = r#"# Chapter
+
+* [Header 1](#header-1)
+  * [Header 1.1](#header-11)
+* [Header 2](#header-2)
+  * [Header 2.1](#header-21)
+  * [Header 2.2](#header-22)
+    * [Header 2.2.1](#header-221)
+
+# Header 1
+
+## Header 1.1
+
+# Header 2
+
+## Header 2.1
+
+## Header 2.2
+
+### Header 2.2.1"#;
+
+        assert_eq!(expected, add_toc(content, &marker).unwrap());
     }
 }