1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107

//! [D2] diagram generator [`Preprocessor`] library for [`MdBook`](https://rust-lang.github.io/mdBook/).

#![deny(
    clippy::all,
    missing_debug_implementations,
    missing_copy_implementations,
    missing_docs,
    clippy::cargo
)]
#![warn(clippy::pedantic, clippy::nursery)]

use mdbook::book::{Book, Chapter};
use mdbook::errors::Error;
use mdbook::preprocess::{Preprocessor, PreprocessorContext};
use mdbook::BookItem;
use pulldown_cmark::{CodeBlockKind, CowStr, Event, Options, Parser, Tag};
use pulldown_cmark_to_cmark::cmark;

mod backend;
use backend::{Backend, RenderContext};

mod config;

/// [D2] diagram generator [`Preprocessor`] for [`MdBook`](https://rust-lang.github.io/mdBook/).
#[derive(Default, Clone, Copy, Debug)]
pub struct D2;

impl Preprocessor for D2 {
    fn name(&self) -> &str {
        "d2"
    }

    fn run(&self, ctx: &PreprocessorContext, mut book: Book) -> Result<Book, Error> {
        let backend = Backend::from_context(ctx);

        book.for_each_mut(|section| {
            if let BookItem::Chapter(chapter) = section {
                let events = process_events(
                    &backend,
                    chapter,
                    Parser::new_ext(&chapter.content, Options::all()),
                );

                // create a buffer in which we can place the markdown
                let mut buf = String::with_capacity(chapter.content.len() + 128);

                // convert it back to markdown and replace the original chapter's content
                cmark(events, &mut buf).unwrap();
                chapter.content = buf;
            }
        });

        Ok(book)
    }
}

fn process_events<'a>(
    backend: &'a Backend,
    chapter: &'a Chapter,
    events: impl Iterator<Item = Event<'a>> + 'a,
) -> impl Iterator<Item = Event<'a>> + 'a {
    let mut in_block = false;
    // if Windows crlf line endings are used, a code block will consist
    // of many different Text blocks, thus we need to buffer them in here
    // see https://github.com/raphlinus/pulldown-cmark/issues/507
    let mut diagram = String::new();
    let mut diagram_index = 0;

    events.flat_map(move |event| {
        match (&event, in_block) {
            // check if we are entering a d2 codeblock
            (
                Event::Start(Tag::CodeBlock(CodeBlockKind::Fenced(CowStr::Borrowed("d2")))),
                false,
            ) => {
                in_block = true;
                diagram.clear();
                diagram_index += 1;
                vec![]
            }
            // check if we are currently inside a d2 block
            (Event::Text(content), true) => {
                diagram.push_str(content);
                vec![]
            }
            // check if we are exiting a d2 block
            (Event::End(Tag::CodeBlock(CodeBlockKind::Fenced(CowStr::Borrowed("d2")))), true) => {
                in_block = false;
                let render_context = RenderContext::new(
                    chapter.source_path.as_ref().unwrap(),
                    &chapter.name,
                    chapter.number.as_ref(),
                    diagram_index,
                );
                match backend.render(&render_context, &diagram) {
                    Ok(events) => events,
                    Err(e) => {
                        eprintln!("{e}");
                        vec![]
                    }
                }
            }
            // if nothing matches, change nothing
            _ => vec![event],
        }
    })
}