mdx_gen/
extensions.rs

1//! Extension functionality for the MDX Gen library.
2//!
3//! This module provides utilities for enhancing Markdown processing,
4//! including custom block handling and table formatting.
5//! Syntax highlighting has moved to the `highlight` module, which
6//! the default `syntax_highlighting` feature provides.
7
8use crate::error::MarkdownError;
9use comrak::nodes::{NodeHtmlBlock, NodeValue};
10use regex::Regex;
11use std::cell::RefCell;
12use std::collections::HashMap;
13use std::str::FromStr;
14use std::sync::LazyLock;
15
16// ── Headings / table-of-contents ────────────────────────────────────
17
18/// A single heading discovered in a Markdown document.
19///
20/// Returned by [`collect_headings`] and by
21/// [`crate::process_markdown_with_toc`].
22#[derive(Debug, Clone, PartialEq, Eq)]
23pub struct Heading {
24    /// Heading level (1-6 for ATX, 1-2 for setext).
25    pub level: u8,
26    /// Plain-text content of the heading (markup stripped).
27    pub text: String,
28    /// The anchor id that comrak emits for this heading. Computed
29    /// with [`comrak::Anchorizer`] so it matches the `id="…"` value
30    /// produced when `MarkdownOptions::header_ids` is set.
31    pub id: String,
32}
33
34/// Walks the AST and returns one [`Heading`] per heading node, in
35/// document order.
36///
37/// `prefix` mirrors `MarkdownOptions::header_ids`:
38/// * `None` or `Some("")` → bare slug (`"introduction"`).
39/// * `Some(p)` → `format!("{p}{slug}")` (`"user-content-introduction"`).
40///
41/// Uses comrak's own `Anchorizer` so dedup behaviour (`-1`, `-2`
42/// suffixes for repeated headings) matches the rendered HTML.
43pub fn collect_headings<'a>(
44    root: comrak::nodes::Node<'a>,
45    prefix: Option<&str>,
46) -> Vec<Heading> {
47    let mut anchorizer = comrak::Anchorizer::new();
48    let mut out = Vec::new();
49    for node in root.descendants() {
50        let level = match node.data.borrow().value {
51            NodeValue::Heading(h) => h.level,
52            _ => continue,
53        };
54        let text = extract_text(node);
55        let slug = anchorizer.anchorize(&text);
56        let id = match prefix {
57            Some(p) if !p.is_empty() => format!("{p}{slug}"),
58            _ => slug,
59        };
60        out.push(Heading { level, text, id });
61    }
62    out
63}
64
65/// Recursively concatenates the plain-text content of a node's subtree,
66/// stripping all Markdown and HTML markup. Blocks are separated by
67/// whitespace to prevent words from merging.
68///
69/// Captures both inline code (`` `foo` ``) and fenced code blocks —
70/// useful for search indexing where readers may query terms that
71/// only appear inside code samples.
72pub fn collect_all_text<'a>(root: comrak::nodes::Node<'a>) -> String {
73    let mut buf = String::new();
74    for d in root.descendants() {
75        match &d.data.borrow().value {
76            NodeValue::Text(t) => buf.push_str(t),
77            NodeValue::Code(c) => buf.push_str(&c.literal),
78            NodeValue::CodeBlock(cb) => {
79                if !buf.is_empty() && !buf.ends_with(' ') {
80                    buf.push(' ');
81                }
82                buf.push_str(&cb.literal);
83            }
84            NodeValue::SoftBreak | NodeValue::LineBreak
85                if !buf.is_empty() && !buf.ends_with(' ') =>
86            {
87                buf.push(' ');
88            }
89            // Ensure space between structural elements.
90            NodeValue::Paragraph
91            | NodeValue::Heading(_)
92            | NodeValue::Item(_)
93            | NodeValue::BlockQuote
94            | NodeValue::Table(_)
95            | NodeValue::TableRow(_)
96            | NodeValue::TableCell
97                if !buf.is_empty() && !buf.ends_with(' ') =>
98            {
99                buf.push(' ');
100            }
101            _ => {}
102        }
103    }
104    buf.trim().to_string()
105}
106
107/// Recursively concatenates the text content of a node's subtree,
108/// matching what comrak renders inside `<h*>` tags (text, inline
109/// code, image alt text). Raw inline HTML is skipped.
110fn extract_text<'a>(node: comrak::nodes::Node<'a>) -> String {
111    let mut buf = String::new();
112    for d in node.descendants() {
113        match &d.data.borrow().value {
114            NodeValue::Text(t) => buf.push_str(t),
115            NodeValue::Code(c) => buf.push_str(&c.literal),
116            NodeValue::Image(img) => buf.push_str(&img.title),
117            _ => {}
118        }
119    }
120    buf
121}
122
123// ── Table regexes (cached, for legacy process_tables) ───────────────
124//
125// Opening and closing `<table>` tags are literal substrings, handled
126// by `str::replace` in `process_tables` below — no regex needed. The
127// `<td …>` rewrite does need a pattern to capture the attribute run,
128// so it stays as a cached `Regex`.
129
130static TABLE_CELL_RE: LazyLock<Regex> =
131    LazyLock::new(|| Regex::new(r"<td([^>]*)>").unwrap());
132
133/// Regex matching known custom block div elements inside HTML blocks.
134static CUSTOM_BLOCK_RE: LazyLock<Regex> = LazyLock::new(|| {
135    Regex::new(
136        r#"(?si)<div\s+class=["']?(note|warning|tip|info|important|caution)["']?>(.*?)</div>"#,
137    )
138    .unwrap()
139});
140
141// ── Column alignment ────────────────────────────────────────────────
142
143/// Alignment options for table columns.
144#[derive(Debug, Clone, Copy, PartialEq)]
145pub enum ColumnAlignment {
146    /// Align the column to the left.
147    Left,
148    /// Align the column to the center.
149    Center,
150    /// Align the column to the right.
151    Right,
152}
153
154// ── Custom block types ──────────────────────────────────────────────
155
156/// Represents different types of custom blocks.
157#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
158pub enum CustomBlockType {
159    /// A note block.
160    Note,
161    /// A warning block.
162    Warning,
163    /// A tip block.
164    Tip,
165    /// An info block.
166    Info,
167    /// An important block.
168    Important,
169    /// A caution block.
170    Caution,
171}
172
173impl CustomBlockType {
174    /// Returns the default Bootstrap alert class.
175    pub fn default_alert_class(&self) -> &'static str {
176        match self {
177            Self::Note => "alert-info",
178            Self::Warning => "alert-warning",
179            Self::Tip => "alert-success",
180            Self::Info => "alert-primary",
181            Self::Important => "alert-danger",
182            Self::Caution => "alert-secondary",
183        }
184    }
185
186    /// Returns the default human-readable title.
187    pub fn default_title(&self) -> &'static str {
188        match self {
189            Self::Note => "Note",
190            Self::Warning => "Warning",
191            Self::Tip => "Tip",
192            Self::Info => "Info",
193            Self::Important => "Important",
194            Self::Caution => "Caution",
195        }
196    }
197
198    /// Returns the default Bootstrap alert class for this block type.
199    pub fn get_alert_class(&self) -> &'static str {
200        self.default_alert_class()
201    }
202
203    /// Returns the default title for this block type.
204    pub fn get_title(&self) -> &'static str {
205        self.default_title()
206    }
207
208    /// Returns the alert class, respecting config overrides.
209    pub fn alert_class_with<'a>(
210        &self,
211        config: &'a CustomBlockConfig,
212    ) -> &'a str {
213        config
214            .class_overrides
215            .get(self)
216            .map(|s| s.as_str())
217            .unwrap_or_else(move || self.default_alert_class())
218    }
219
220    /// Returns the title, respecting config overrides.
221    pub fn title_with<'a>(
222        &self,
223        config: &'a CustomBlockConfig,
224    ) -> &'a str {
225        config
226            .title_overrides
227            .get(self)
228            .map(|s| s.as_str())
229            .unwrap_or_else(move || self.default_title())
230    }
231}
232
233impl FromStr for CustomBlockType {
234    type Err = MarkdownError;
235
236    fn from_str(block_type: &str) -> Result<Self, Self::Err> {
237        match block_type.to_lowercase().as_str() {
238            "note" => Ok(Self::Note),
239            "warning" => Ok(Self::Warning),
240            "tip" => Ok(Self::Tip),
241            "info" => Ok(Self::Info),
242            "important" => Ok(Self::Important),
243            "caution" => Ok(Self::Caution),
244            _ => Err(MarkdownError::CustomBlockError(format!(
245                "Unknown block type: {block_type}"
246            ))),
247        }
248    }
249}
250
251// ── Custom block configuration ──────────────────────────────────────
252
253/// Configuration for custom block rendering.
254///
255/// Allows overriding the default CSS class and title for each
256/// block type, enabling use with CSS frameworks other than Bootstrap.
257#[derive(Debug, Clone, Default)]
258pub struct CustomBlockConfig {
259    /// Override the CSS alert class per block type.
260    pub class_overrides: HashMap<CustomBlockType, String>,
261    /// Override the display title per block type.
262    pub title_overrides: HashMap<CustomBlockType, String>,
263}
264
265impl CustomBlockConfig {
266    /// Creates a new empty configuration (uses all defaults).
267    pub fn new() -> Self {
268        Self::default()
269    }
270
271    /// Overrides the CSS class for a specific block type.
272    pub fn with_class(
273        mut self,
274        block_type: CustomBlockType,
275        class: impl Into<String>,
276    ) -> Self {
277        self.class_overrides.insert(block_type, class.into());
278        self
279    }
280
281    /// Overrides the display title for a specific block type.
282    pub fn with_title(
283        mut self,
284        block_type: CustomBlockType,
285        title: impl Into<String>,
286    ) -> Self {
287        self.title_overrides.insert(block_type, title.into());
288        self
289    }
290}
291
292// ── AST-level custom block processing ───────────────────────────────
293
294/// Walks the comrak AST and transforms `HtmlBlock` nodes that contain
295/// known custom block divs into styled alert HTML.
296///
297/// This is safer than regex on rendered HTML because it only touches
298/// nodes the parser explicitly identified as raw HTML blocks.
299pub fn process_custom_block_nodes<'a>(
300    root: comrak::nodes::Node<'a>,
301    config: &CustomBlockConfig,
302) {
303    for node in root.descendants() {
304        let mut ast = node.data.borrow_mut();
305        if let NodeValue::HtmlBlock(ref mut block) = ast.value {
306            block.literal =
307                transform_custom_blocks(&block.literal, config);
308        }
309    }
310}
311
312/// Transforms custom block divs in a raw HTML string.
313fn transform_custom_blocks(
314    html: &str,
315    config: &CustomBlockConfig,
316) -> String {
317    CUSTOM_BLOCK_RE
318        .replace_all(html, |caps: &regex::Captures| {
319            let block_type = CustomBlockType::from_str(
320                caps.get(1).unwrap().as_str(),
321            )
322            .expect("regex only matches known block types");
323            generate_custom_block_html(block_type, &caps[2], config)
324        })
325        .to_string()
326}
327
328/// Generates the HTML for a custom block.
329fn generate_custom_block_html(
330    block_type: CustomBlockType,
331    content: &str,
332    config: &CustomBlockConfig,
333) -> String {
334    format!(
335        r#"<div class="alert {}" role="alert"><strong>{}:</strong> {}</div>"#,
336        block_type.alert_class_with(config),
337        block_type.title_with(config),
338        content
339    )
340}
341
342// ── AST-level table enhancement ─────────────────────────────────────
343
344/// Walks the comrak AST and replaces `Table` nodes with `HtmlBlock`
345/// nodes containing responsive-wrapped, class-enhanced table HTML.
346///
347/// This eliminates the last regex pass over rendered HTML.
348pub fn enhance_table_nodes<'a>(
349    root: comrak::nodes::Node<'a>,
350    arena: &'a comrak::Arena<'a>,
351    options: &comrak::Options,
352) {
353    // Collect table nodes first to avoid borrow issues during mutation
354    let table_nodes: Vec<comrak::nodes::Node<'a>> = root
355        .descendants()
356        .filter(|node| {
357            matches!(node.data.borrow().value, NodeValue::Table(_))
358        })
359        .collect();
360
361    for table_node in table_nodes {
362        // Render this table subtree to HTML
363        let mut table_html = String::new();
364        if comrak::format_html(table_node, options, &mut table_html)
365            .is_err()
366        {
367            continue;
368        }
369
370        // Apply the responsive wrapper and alignment classes
371        let enhanced = process_tables(&table_html);
372
373        // Create a replacement HtmlBlock node
374        let start = comrak::nodes::LineColumn { line: 0, column: 0 };
375        let replacement = arena.alloc(comrak::nodes::AstNode::new(
376            RefCell::new(comrak::nodes::Ast::new(
377                NodeValue::HtmlBlock(NodeHtmlBlock {
378                    block_type: 6, // generic block
379                    literal: enhanced,
380                }),
381                start,
382            )),
383        ));
384
385        // Insert replacement and remove original
386        table_node.insert_before(replacement);
387        table_node.detach();
388    }
389}
390
391// ── Legacy string-level custom block processing ─────────────────────
392
393/// Processes custom blocks in an HTML string.
394///
395/// Provided for backward compatibility. Prefer
396/// [`process_custom_block_nodes`] for AST-level processing.
397pub fn process_custom_blocks(content: &str) -> String {
398    transform_custom_blocks(content, &CustomBlockConfig::default())
399}
400
401// ── Table post-processing ───────────────────────────────────────────
402
403/// Processes tables, enhancing them with responsive design and alignment classes.
404pub fn process_tables(table_html: &str) -> String {
405    let table_html = table_html.replace(
406        "<table>",
407        r#"<div class="table-responsive"><table class="table">"#,
408    );
409    let table_html = table_html.replace("</table>", "</table></div>");
410
411    TABLE_CELL_RE
412        .replace_all(&table_html, |caps: &regex::Captures| {
413            let attrs = &caps[1];
414            if attrs.contains("align=\"center\"") {
415                format!(r#"<td{attrs} class="text-center">"#)
416            } else if attrs.contains("align=\"right\"") {
417                format!(r#"<td{attrs} class="text-right">"#)
418            } else {
419                format!(r#"<td{attrs} class="text-left">"#)
420            }
421        })
422        .to_string()
423}
424
425#[cfg(test)]
426mod tests {
427    use super::*;
428
429    #[test]
430    fn test_process_custom_blocks_default_config() {
431        let input = r#"
432            <div class="note">This is a note.</div>
433            <div class="WARNING">This is a warning.</div>
434            <div class="Tip">This is a tip.</div>
435        "#;
436        let processed = process_custom_blocks(input);
437        assert!(processed.contains(r#"alert alert-info"#));
438        assert!(processed.contains(r#"alert alert-warning"#));
439        assert!(processed.contains(r#"alert alert-success"#));
440    }
441
442    #[test]
443    fn test_custom_block_config_overrides() {
444        let config = CustomBlockConfig::new()
445            .with_class(CustomBlockType::Note, "callout-info")
446            .with_title(CustomBlockType::Note, "Did you know?");
447
448        let html = generate_custom_block_html(
449            CustomBlockType::Note,
450            "test content",
451            &config,
452        );
453        assert!(html.contains("callout-info"));
454        assert!(html.contains("Did you know?:"));
455    }
456
457    #[test]
458    fn test_unknown_block_passthrough() {
459        let input =
460            r#"<div class="unknown">Should pass through.</div>"#;
461        let processed = process_custom_blocks(input);
462        assert_eq!(processed, input);
463    }
464
465    #[test]
466    fn test_process_tables() {
467        let input = r#"<table><tr><td align="center">Center</td><td align="right">Right</td><td>Left</td></tr></table>"#;
468        let processed = process_tables(input);
469        assert!(processed.contains(r#"table-responsive"#));
470        assert!(processed.contains(r#"text-center"#));
471        assert!(processed.contains(r#"text-right"#));
472        assert!(processed.contains(r#"text-left"#));
473    }
474
475    #[test]
476    fn test_process_multiple_tables() {
477        let input = "<table><tr><td>A</td></tr></table>\n<table><tr><td>B</td></tr></table>";
478        let processed = process_tables(input);
479        assert_eq!(processed.matches("table-responsive").count(), 2);
480    }
481
482    #[test]
483    fn test_unknown_block_type_from_str() {
484        let result = CustomBlockType::from_str("unknown");
485        assert!(result.is_err());
486        let err = result.unwrap_err();
487        assert!(
488            err.to_string().contains("Unknown block type: unknown"),
489            "Error message should contain the unknown type"
490        );
491    }
492
493    #[test]
494    fn test_unknown_block_type_from_str_various() {
495        for name in ["foobar", "alert", "danger", "success", ""] {
496            let result = CustomBlockType::from_str(name);
497            assert!(
498                result.is_err(),
499                "'{name}' should not parse as a valid block type"
500            );
501        }
502    }
503}
mdx_gen/extensions.rs

mdx_gen/
extensions.rs
