mdbook_toc/
lib.rs

1use std::cmp::Ordering;
2use std::collections::HashMap;
3use std::fmt::Write;
4
5use mdbook_preprocessor::book::{Book, BookItem, Chapter};
6use mdbook_preprocessor::errors::Result;
7use mdbook_preprocessor::{Preprocessor, PreprocessorContext};
8use pulldown_cmark::{Event, Options, Parser};
9use pulldown_cmark::{Tag::*, TagEnd};
10
11pub struct Toc;
12
13static DEFAULT_MARKER: &str = "<!-- toc -->\n";
14static DEFAULT_MAX_LEVEL: u32 = 4;
15
16/// Configuration for Table of Contents generation
17pub struct Config {
18    /// Marker to use, defaults to `<!-- toc -->\n`
19    pub marker: String,
20    /// The maximum level of headers to include in the table of contents.
21    /// Defaults to `4`.
22    pub max_level: u32,
23}
24
25impl Default for Config {
26    fn default() -> Config {
27        Config {
28            marker: DEFAULT_MARKER.into(),
29            max_level: DEFAULT_MAX_LEVEL,
30        }
31    }
32}
33
34impl Preprocessor for Toc {
35    fn name(&self) -> &str {
36        "toc"
37    }
38
39    fn run(&self, ctx: &PreprocessorContext, mut book: Book) -> Result<Book> {
40        let mut res = None;
41        let cfg_key = |key| format!("preprocessor.{}.{}", self.name(), key);
42        let cfg = Config {
43            marker: ctx
44                .config
45                .get(&cfg_key("marker"))?
46                .unwrap_or_else(|| DEFAULT_MARKER.into()),
47            max_level: ctx
48                .config
49                .get(&cfg_key("max_level"))?
50                .unwrap_or(DEFAULT_MAX_LEVEL),
51        };
52
53        book.for_each_mut(|item: &mut BookItem| {
54            if let Some(Err(_)) = res {
55                return;
56            }
57
58            if let BookItem::Chapter(ref mut chapter) = *item {
59                res = Some(Toc::add_toc(chapter, &cfg).map(|md| {
60                    chapter.content = md;
61                }));
62            }
63        });
64
65        res.unwrap_or(Ok(())).map(|_| book)
66    }
67}
68
69fn build_toc(toc: &[(u32, String, String)]) -> String {
70    log::trace!("ToC from {toc:?}");
71    let mut result = String::new();
72
73    // "Normalize" header levels.
74    // If headers skip a level, we need to normalize them to avoid the skip.
75    // Otherwise the markdown render will escape nested levels.
76    //
77    // This is a rough approximation only.
78    let mut toc_iter = toc.iter().peekable();
79
80    // Start from the level of the first header.
81    let min_level = toc.iter().map(|(lvl, _, _)| *lvl).min().unwrap_or(1);
82    let mut last_lower = match toc_iter.peek() {
83        Some((lvl, _, _)) => *lvl,
84        None => 0,
85    };
86    let toc = toc.iter().map(|(lvl, name, slug)| {
87        let lvl = *lvl;
88        let lvl = match (last_lower + 1).cmp(&lvl) {
89            Ordering::Less => last_lower + 1,
90            _ => {
91                last_lower = lvl;
92                lvl
93            }
94        };
95        (lvl, name, slug)
96    });
97
98    for (level, name, slug) in toc {
99        let width = 2 * (level - min_level) as usize;
100        writeln!(result, "{:width$}* [{name}](#{slug})", "").unwrap();
101    }
102
103    result
104}
105
106/// Convert the given string to a valid HTML element ID.
107/// The only restriction is that the ID must not contain any ASCII whitespace.
108fn normalize_id(content: &str) -> String {
109    content
110        .trim()
111        .to_lowercase()
112        .chars()
113        .filter_map(|ch| {
114            if ch.is_alphanumeric() || ch == '_' || ch == '-' {
115                Some(ch)
116            } else if ch.is_whitespace() {
117                Some('-')
118            } else {
119                None
120            }
121        })
122        .collect::<String>()
123}
124
125fn add_toc(content: &str, cfg: &Config) -> Result<String> {
126    let mut toc_found = false;
127
128    let mut toc_content = vec![];
129    let mut current_header = None;
130    let mut current_header_level: Option<(u32, Option<String>)> = None;
131    let mut id_counter = HashMap::new();
132
133    let opts = Options::ENABLE_TABLES
134        | Options::ENABLE_FOOTNOTES
135        | Options::ENABLE_STRIKETHROUGH
136        | Options::ENABLE_TASKLISTS
137        | Options::ENABLE_HEADING_ATTRIBUTES;
138
139    let mark: Vec<Event> = Parser::new(&cfg.marker).collect();
140    log::trace!("Marker: {mark:?}");
141    let mut mark_start = None;
142    let mut mark_end = 0..0;
143    let mut mark_loc = 0;
144
145    let content = content.replace("\r\n", "\n");
146    for (e, span) in Parser::new_ext(&content, opts).into_offset_iter() {
147        log::trace!(
148            "Event: {e:?} (span: {span:?}, content: {:?})",
149            &content[span.start..span.end]
150        );
151        if !toc_found {
152            log::trace!("TOC not found yet. Location: {mark_loc}, Start: {mark_start:?}");
153            if e == mark[mark_loc] {
154                if mark_start.is_none() {
155                    mark_start = Some(span.clone());
156                }
157                mark_loc += 1;
158                if mark_loc >= mark.len() {
159                    mark_end = span.clone();
160                    toc_found = true
161                }
162            } else if mark_loc > 0 {
163                mark_loc = 0;
164                mark_start = None;
165            } else {
166                continue;
167            }
168        }
169        log::trace!("TOC found. Location: {mark_loc}, Start: {mark_start:?}");
170
171        if let Event::Start(Heading { level, id, .. }) = e {
172            log::trace!("Header(lvl={level}, fragment={id:?})");
173            let id = id.map(|s| s.to_string());
174            current_header_level = Some((level as u32, id));
175
176            // Find start of the header after `#`
177            let header_content = content[span.start..span.end].trim_end();
178            let idx = header_content.find(|c: char| c != '#' && !c.is_ascii_whitespace());
179            current_header = Some((span.start + idx.unwrap_or(0), 0));
180            continue;
181        }
182        // Headers might consist of text and code. pulldown_cmark unescapes `\\`, so we try to find
183        // the correct span and extract the text ourselves later.
184        // We enabled `HEADING_ATTRIBUTES` so attributes within `{ }` won't be in the emitted event
185        if let Some(current_header) = &mut current_header {
186            if let Event::Text(_) = &e
187                && span.end > current_header.1
188            {
189                current_header.1 = span.end;
190            }
191            if let Event::Code(_) = &e
192                && span.end > current_header.1
193            {
194                current_header.1 = span.end;
195            }
196        }
197        if let Event::End(TagEnd::Heading(header_lvl)) = e {
198            // Skip if this header is nested too deeply.
199            if let Some((level, id)) = current_header_level.take() {
200                assert!(header_lvl as u32 == level);
201                let header_span = current_header.take().unwrap();
202                let header = content[header_span.0..header_span.1].trim_end();
203                let slug = if let Some(slug) = id {
204                    // If a fragment is defined, take it as is, not trying to append an extra ID
205                    // in case of duplicates (same behavior as mdBook)
206                    slug.to_owned()
207                } else {
208                    let mut slug = normalize_id(header);
209                    let id_count = id_counter.entry(slug.clone()).or_insert(0);
210
211                    // Append unique ID if multiple headers with the same name exist
212                    // to follow what mdBook does
213                    if *id_count > 0 {
214                        write!(slug, "-{id_count}").unwrap();
215                    }
216
217                    *id_count += 1;
218                    slug
219                };
220
221                if level <= cfg.max_level {
222                    toc_content.push((level, header.to_string(), slug));
223                }
224            }
225            continue;
226        }
227        if current_header_level.is_none() {
228            continue;
229        }
230    }
231
232    let toc = build_toc(&toc_content);
233    log::trace!("Built TOC: {toc:?}");
234    log::trace!("toc_found={toc_found} mark_start={mark_start:?} mark_end={mark_end:?}");
235
236    let content = if toc_found {
237        let mark_start = mark_start.unwrap();
238        let content_before_toc = &content[0..mark_start.start];
239        let content_after_toc = &content[mark_end.end..];
240        log::trace!("content_before_toc={content_before_toc:?}");
241        log::trace!("content_after_toc={content_after_toc:?}");
242        // Multiline markers might have consumed trailing newlines,
243        // we ensure there's always one before the content.
244        let extra = if content_after_toc.is_empty() || content_after_toc.as_bytes()[0] == b'\n' {
245            ""
246        } else {
247            "\n"
248        };
249        format!("{content_before_toc}{toc}{extra}{content_after_toc}")
250    } else {
251        content.to_string()
252    };
253
254    Ok(content)
255}
256
257impl Toc {
258    /// Add a table of contents to the given chapter.
259    pub fn add_toc(chapter: &Chapter, cfg: &Config) -> Result<String> {
260        add_toc(&chapter.content, cfg)
261    }
262}
mdbook_toc/lib.rs

mdbook_toc/
lib.rs
