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        .chars()
111        .filter_map(|ch| {
112            if ch.is_alphanumeric() || ch == '_' || ch == '-' {
113                Some(ch.to_ascii_lowercase())
114            } else if ch.is_whitespace() {
115                Some('-')
116            } else {
117                None
118            }
119        })
120        .collect::<String>()
121}
122
123fn add_toc(content: &str, cfg: &Config) -> Result<String> {
124    let mut toc_found = false;
125
126    let mut toc_content = vec![];
127    let mut current_header = None;
128    let mut current_header_level: Option<(u32, Option<String>)> = None;
129    let mut id_counter = HashMap::new();
130
131    let opts = Options::ENABLE_TABLES
132        | Options::ENABLE_FOOTNOTES
133        | Options::ENABLE_STRIKETHROUGH
134        | Options::ENABLE_TASKLISTS
135        | Options::ENABLE_HEADING_ATTRIBUTES;
136
137    let mark: Vec<Event> = Parser::new(&cfg.marker).collect();
138    log::trace!("Marker: {mark:?}");
139    let mut mark_start = None;
140    let mut mark_end = 0..0;
141    let mut mark_loc = 0;
142
143    let content = content.replace("\r\n", "\n");
144    for (e, span) in Parser::new_ext(&content, opts).into_offset_iter() {
145        log::trace!(
146            "Event: {e:?} (span: {span:?}, content: {:?})",
147            &content[span.start..span.end]
148        );
149        if !toc_found {
150            log::trace!("TOC not found yet. Location: {mark_loc}, Start: {mark_start:?}");
151            if e == mark[mark_loc] {
152                if mark_start.is_none() {
153                    mark_start = Some(span.clone());
154                }
155                mark_loc += 1;
156                if mark_loc >= mark.len() {
157                    mark_end = span.clone();
158                    toc_found = true
159                }
160            } else if mark_loc > 0 {
161                mark_loc = 0;
162                mark_start = None;
163            } else {
164                continue;
165            }
166        }
167        log::trace!("TOC found. Location: {mark_loc}, Start: {mark_start:?}");
168
169        if let Event::Start(Heading { level, id, .. }) = e {
170            log::trace!("Header(lvl={level}, fragment={id:?})");
171            let id = id.map(|s| s.to_string());
172            current_header_level = Some((level as u32, id));
173
174            // Find start of the header after `#`
175            let header_content = content[span.start..span.end].trim_end();
176            let idx = header_content.find(|c: char| c != '#' && !c.is_ascii_whitespace());
177            current_header = Some((span.start + idx.unwrap_or(0), 0));
178            continue;
179        }
180        // Headers might consist of text and code. pulldown_cmark unescapes `\\`, so we try to find
181        // the correct span and extract the text ourselves later.
182        // We enabled `HEADING_ATTRIBUTES` so attributes within `{ }` won't be in the emitted event
183        if let Some(current_header) = &mut current_header {
184            if let Event::Text(_) = &e
185                && span.end > current_header.1
186            {
187                current_header.1 = span.end;
188            }
189            if let Event::Code(_) = &e
190                && span.end > current_header.1
191            {
192                current_header.1 = span.end;
193            }
194        }
195        if let Event::End(TagEnd::Heading(header_lvl)) = e {
196            // Skip if this header is nested too deeply.
197            if let Some((level, id)) = current_header_level.take() {
198                assert!(header_lvl as u32 == level);
199                let header_span = current_header.take().unwrap();
200                let header = content[header_span.0..header_span.1].trim_end();
201                let slug = if let Some(slug) = id {
202                    // If a fragment is defined, take it as is, not trying to append an extra ID
203                    // in case of duplicates (same behavior as mdBook)
204                    slug.to_owned()
205                } else {
206                    let mut slug = normalize_id(header);
207                    let id_count = id_counter.entry(slug.clone()).or_insert(0);
208
209                    // Append unique ID if multiple headers with the same name exist
210                    // to follow what mdBook does
211                    if *id_count > 0 {
212                        write!(slug, "-{id_count}").unwrap();
213                    }
214
215                    *id_count += 1;
216                    slug
217                };
218
219                if level <= cfg.max_level {
220                    toc_content.push((level, header.to_string(), slug));
221                }
222            }
223            continue;
224        }
225        if current_header_level.is_none() {
226            continue;
227        }
228    }
229
230    let toc = build_toc(&toc_content);
231    log::trace!("Built TOC: {toc:?}");
232    log::trace!("toc_found={toc_found} mark_start={mark_start:?} mark_end={mark_end:?}");
233
234    let content = if toc_found {
235        let mark_start = mark_start.unwrap();
236        let content_before_toc = &content[0..mark_start.start];
237        let content_after_toc = &content[mark_end.end..];
238        log::trace!("content_before_toc={content_before_toc:?}");
239        log::trace!("content_after_toc={content_after_toc:?}");
240        // Multiline markers might have consumed trailing newlines,
241        // we ensure there's always one before the content.
242        let extra = if content_after_toc.is_empty() || content_after_toc.as_bytes()[0] == b'\n' {
243            ""
244        } else {
245            "\n"
246        };
247        format!("{content_before_toc}{toc}{extra}{content_after_toc}")
248    } else {
249        content.to_string()
250    };
251
252    Ok(content)
253}
254
255impl Toc {
256    /// Add a table of contents to the given chapter.
257    pub fn add_toc(chapter: &Chapter, cfg: &Config) -> Result<String> {
258        add_toc(&chapter.content, cfg)
259    }
260}
mdbook_toc/lib.rs

mdbook_toc/
lib.rs
