Skip to main content

socketry_markdown/mdast/
headings.rs

1// Released under the MIT License.
2// Copyright, 2026, by Samuel Williams.
3
4//! Helpers for extracting headings from an mdast tree.
5use super::Node;
6use alloc::{
7    collections::{BTreeMap, BTreeSet},
8    string::{String, ToString},
9    vec::Vec,
10};
11
12/// Options for extracting headings from a document.
13#[derive(Clone, Copy, Debug, Eq, PartialEq)]
14pub struct HeadingOptions {
15    /// The shallowest heading level to include (between `1` and `6`).
16    pub min_level: u8,
17    /// The deepest heading level to include (between `1` and `6`).
18    pub max_level: u8,
19}
20
21impl Default for HeadingOptions {
22    fn default() -> Self {
23        Self {
24            min_level: 1,
25            max_level: 6,
26        }
27    }
28}
29
30/// A heading extracted from an mdast tree.
31#[derive(Clone, Debug, Eq, PartialEq)]
32pub struct HeadingEntry<'a> {
33    /// The original heading node.
34    pub node: &'a Node,
35    /// The heading level, between `1` and `6`.
36    pub level: u8,
37    /// The heading's text content.
38    pub text: String,
39    /// A unique, lowercase anchor derived from the heading text.
40    pub anchor: String,
41}
42
43/// Headings extracted from a Markdown document.
44///
45/// This helper is useful for building a table of contents. It walks the tree
46/// in document order, converts each heading to plain text, and assigns unique
47/// anchors to repeated headings. Anchor assignment includes headings filtered
48/// out by [`HeadingOptions`], so the remaining anchors still match rendered
49/// headings when [`CompileOptions::heading_ids`][crate::CompileOptions::heading_ids]
50/// is enabled.
51#[derive(Clone, Debug, Default, Eq, PartialEq)]
52pub struct Headings<'a> {
53    entries: Vec<HeadingEntry<'a>>,
54}
55
56impl<'a> Headings<'a> {
57    /// Extract all headings from a document tree.
58    #[must_use]
59    pub fn extract(root: &'a Node) -> Self {
60        Self::extract_with_options(root, &HeadingOptions::default())
61    }
62
63    /// Extract headings from a document tree with a level range.
64    #[must_use]
65    pub fn extract_with_options(root: &'a Node, options: &HeadingOptions) -> Self {
66        let mut result = Self::default();
67        let mut anchors = AnchorGenerator::default();
68
69        root.walk(|node| {
70            if let Node::Heading(heading) = node {
71                let text = chomp_line_ending(&node.text_content()).to_string();
72                let anchor = anchors.anchor_for(&text);
73
74                if heading.depth >= options.min_level && heading.depth <= options.max_level {
75                    result.entries.push(HeadingEntry {
76                        node,
77                        level: heading.depth,
78                        text,
79                        anchor,
80                    });
81                }
82            }
83        });
84
85        result
86    }
87
88    /// Return the extracted entries as a slice.
89    #[must_use]
90    pub fn as_slice(&self) -> &[HeadingEntry<'a>] {
91        &self.entries
92    }
93
94    /// Iterate over the extracted entries in document order.
95    pub fn iter(&self) -> core::slice::Iter<'_, HeadingEntry<'a>> {
96        self.entries.iter()
97    }
98
99    /// Return the number of extracted headings.
100    #[must_use]
101    pub fn len(&self) -> usize {
102        self.entries.len()
103    }
104
105    /// Whether the document contains no headings in the selected level range.
106    #[must_use]
107    pub fn is_empty(&self) -> bool {
108        self.entries.is_empty()
109    }
110}
111
112impl<'a, 'b> IntoIterator for &'b Headings<'a> {
113    type Item = &'b HeadingEntry<'a>;
114    type IntoIter = core::slice::Iter<'b, HeadingEntry<'a>>;
115
116    fn into_iter(self) -> Self::IntoIter {
117        self.entries.iter()
118    }
119}
120
121impl<'a> IntoIterator for Headings<'a> {
122    type Item = HeadingEntry<'a>;
123    type IntoIter = alloc::vec::IntoIter<HeadingEntry<'a>>;
124
125    fn into_iter(self) -> Self::IntoIter {
126        self.entries.into_iter()
127    }
128}
129
130#[derive(Default)]
131struct AnchorGenerator {
132    next_suffix: BTreeMap<String, usize>,
133    used: BTreeSet<String>,
134}
135
136impl AnchorGenerator {
137    fn anchor_for(&mut self, text: &str) -> String {
138        let base = slug(text);
139        let mut suffix = self.next_suffix.get(&base).copied().unwrap_or(2);
140        let mut anchor = base.clone();
141
142        while self.used.contains(&anchor) {
143            anchor = alloc::format!("{base}-{suffix}");
144            suffix += 1;
145        }
146
147        self.next_suffix.insert(base, suffix);
148        self.used.insert(anchor.clone());
149        anchor
150    }
151}
152
153fn slug(text: &str) -> String {
154    let mut result = String::new();
155    let mut previous_was_whitespace = false;
156    let lowercase = text.to_lowercase();
157
158    for character in lowercase.chars() {
159        if character.is_whitespace() {
160            if !previous_was_whitespace {
161                result.push('-');
162            }
163            previous_was_whitespace = true;
164        } else {
165            result.push(character);
166            previous_was_whitespace = false;
167        }
168    }
169
170    result
171}
172
173fn chomp_line_ending(text: &str) -> &str {
174    if let Some(text) = text.strip_suffix("\r\n") {
175        text
176    } else if let Some(text) = text.strip_suffix('\n') {
177        text
178    } else if let Some(text) = text.strip_suffix('\r') {
179        text
180    } else {
181        text
182    }
183}