From 113362a944e732f849ff2da4ac0aa549df5ce787 Mon Sep 17 00:00:00 2001
From: Dominik Nakamura <dnaka91@gmail.com>
Date: Tue, 1 Aug 2023 21:59:27 +0900
Subject: [PATCH] Use custom header IDs if present (#40)

Co-authored-by: Jan-Erik Rediger <janerik@fnordig.de>
---
 src/lib.rs              | 35 +++++++++++++++++++++--------------
 tests/attributes.in.md  | 11 +++++++++++
 tests/attributes.out.md | 14 ++++++++++++++
 tests/it.rs             | 19 +++++++++++++------
 4 files changed, 59 insertions(+), 20 deletions(-)
 create mode 100644 tests/attributes.in.md
 create mode 100644 tests/attributes.out.md

diff --git a/src/lib.rs b/src/lib.rs
index 41289aa..4015d0c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -140,11 +140,11 @@ fn add_toc(content: &str, cfg: &Config) -> Result<String> {
     let mut current_header_level: Option<u32> = None;
     let mut id_counter = HashMap::new();
 
-    let mut opts = Options::empty();
-    opts.insert(Options::ENABLE_TABLES);
-    opts.insert(Options::ENABLE_FOOTNOTES);
-    opts.insert(Options::ENABLE_STRIKETHROUGH);
-    opts.insert(Options::ENABLE_TASKLISTS);
+    let opts = Options::ENABLE_TABLES
+        | Options::ENABLE_FOOTNOTES
+        | Options::ENABLE_STRIKETHROUGH
+        | Options::ENABLE_TASKLISTS
+        | Options::ENABLE_HEADING_ATTRIBUTES;
 
     let mark: Vec<Event> = Parser::new(&cfg.marker).collect();
     log::trace!("Marker: {mark:?}");
@@ -179,20 +179,27 @@ fn add_toc(content: &str, cfg: &Config) -> Result<String> {
             current_header_level = Some(lvl as u32);
             continue;
         }
-        if let Event::End(Heading(..)) = e {
+        if let Event::End(Heading(_, fragment, _)) = e {
             // Skip if this header is nested too deeply.
             if let Some(level) = current_header_level.take() {
                 let header = current_header.clone();
-                let mut slug = mdbook::utils::normalize_id(&header);
-                let id_count = id_counter.entry(slug.clone()).or_insert(0);
+                let slug = if let Some(slug) = fragment {
+                    // If a fragment is defined, take it as is, not trying to append an extra ID
+                    // in case of duplicates (same behavior as mdBook)
+                    slug.to_owned()
+                } else {
+                    let mut slug = mdbook::utils::normalize_id(&header);
+                    let id_count = id_counter.entry(slug.clone()).or_insert(0);
 
-                // Append unique ID if multiple headers with the same name exist
-                // to follow what mdBook does
-                if *id_count > 0 {
-                    write!(slug, "-{id_count}").unwrap();
-                }
+                    // Append unique ID if multiple headers with the same name exist
+                    // to follow what mdBook does
+                    if *id_count > 0 {
+                        write!(slug, "-{id_count}").unwrap();
+                    }
 
-                *id_count += 1;
+                    *id_count += 1;
+                    slug
+                };
 
                 if level <= cfg.max_level {
                     toc_content.push((level, header, slug));
diff --git a/tests/attributes.in.md b/tests/attributes.in.md
new file mode 100644
index 0000000..cc39889
--- /dev/null
+++ b/tests/attributes.in.md
@@ -0,0 +1,11 @@
+# Chapter
+
+<!-- toc -->
+
+# Header 1
+
+## Header {#header-1-sub}
+
+# Header 2
+
+## Header {#header-2-sub .class1 .class2}
diff --git a/tests/attributes.out.md b/tests/attributes.out.md
new file mode 100644
index 0000000..ebb9ee9
--- /dev/null
+++ b/tests/attributes.out.md
@@ -0,0 +1,14 @@
+# Chapter
+
+* [Header 1](#header-1)
+  * [Header](#header-1-sub)
+* [Header 2](#header-2)
+  * [Header](#header-2-sub)
+
+# Header 1
+
+## Header {#header-1-sub}
+
+# Header 2
+
+## Header {#header-2-sub .class1 .class2}
diff --git a/tests/it.rs b/tests/it.rs
index 4ee922c..62caeda 100644
--- a/tests/it.rs
+++ b/tests/it.rs
@@ -7,15 +7,17 @@ fn default<T: Default>() -> T {
 }
 
 fn with_marker<S: Into<String>>(marker: S) -> Config {
-    let mut cfg = Config::default();
-    cfg.marker = marker.into();
-    cfg
+    Config {
+        marker: marker.into(),
+        ..Config::default()
+    }
 }
 
 fn with_max_level(level: u32) -> Config {
-    let mut cfg = Config::default();
-    cfg.max_level = level;
-    cfg
+    Config {
+        max_level: level,
+        ..Config::default()
+    }
 }
 
 trait FromContent {
@@ -169,3 +171,8 @@ fn empty_document() {
 fn crlf() {
     assert_toc!("crlf");
 }
+
+#[test]
+fn attributes() {
+    assert_toc!("attributes");
+}