1use mdbook::book::{Book, BookItem, Chapter};
2use mdbook::errors::Result;
3use mdbook::preprocess::{Preprocessor, PreprocessorContext};
4use pulldown_cmark::{CodeBlockKind::*, Event, Options, Parser, Tag};
5
6pub struct Wavedrom;
7
8impl Preprocessor for Wavedrom {
9    fn name(&self) -> &str {
10        "wavedrom"
11    }
12
13    fn run(&self, _ctx: &PreprocessorContext, mut book: Book) -> Result<Book> {
14        let mut res = None;
15        book.for_each_mut(|item: &mut BookItem| {
16            if let Some(Err(_)) = res {
17                return;
18            }
19
20            if let BookItem::Chapter(ref mut chapter) = *item {
21                res = Some(Wavedrom::add_wavedrom(chapter).map(|md| {
22                    chapter.content = md;
23                }));
24            }
25        });
26
27        res.unwrap_or(Ok(())).map(|_| book)
28    }
29
30    fn supports_renderer(&self, renderer: &str) -> bool {
31        renderer == "html"
32    }
33}
34
35fn escape_html(s: &str) -> String {
36    let mut output = String::new();
37    for c in s.chars() {
38        match c {
39            '<' => output.push_str("<"),
40            '>' => output.push_str(">"),
41            '"' => output.push_str("""),
42            '&' => output.push_str("&"),
43            _ => output.push(c),
44        }
45    }
46    output
47}
48
49fn add_wavedrom(content: &str) -> Result<String> {
50    let mut wavedrom_content = String::new();
51    let mut in_wavedrom_block = false;
52
53    let mut opts = Options::empty();
54    opts.insert(Options::ENABLE_TABLES);
55    opts.insert(Options::ENABLE_FOOTNOTES);
56    opts.insert(Options::ENABLE_STRIKETHROUGH);
57    opts.insert(Options::ENABLE_TASKLISTS);
58
59    let mut wavedrom_start = 0..0;
60
61    let mut wavedrom_blocks = vec![];
62
63    let events = Parser::new_ext(content, opts);
64    for (e, span) in events.into_offset_iter() {
65        if let Event::Start(Tag::CodeBlock(Fenced(code))) = e.clone() {
66            log::debug!("e={:?}, span={:?}", e, span);
67            if &*code == "wavedrom" {
68                wavedrom_start = span;
69                in_wavedrom_block = true;
70                wavedrom_content.clear();
71            }
72            continue;
73        }
74
75        if !in_wavedrom_block {
76            continue;
77        }
78
79        if let Event::End(Tag::CodeBlock(Fenced(code))) = e {
80            assert_eq!(
81                "wavedrom", &*code,
82                "After an opening wavedrom code block we expect it to close again"
83            );
84            in_wavedrom_block = false;
85            let pre = "```wavedrom\n";
86            let post = "```";
87
88            let wavedrom_content = &content[wavedrom_start.start + pre.len()..span.end - post.len()];
89            let wavedrom_content = escape_html(wavedrom_content);
90            let wavedrom_code = format!("<body onload=\"WaveDrom.ProcessAll()\">\n\n<script type=\"WaveDrom\">{}</script>\n\n", wavedrom_content);
91            wavedrom_blocks.push((wavedrom_start.start..span.end, wavedrom_code.clone()));
92        }
93    }
94
95    let mut content = content.to_string();
96    for (span, block) in wavedrom_blocks.iter().rev() {
97        let pre_content = &content[0..span.start];
98        let post_content = &content[span.end..];
99        content = format!("{}\n{}{}", pre_content, block, post_content);
100    }
101    Ok(content)
102}
103
104impl Wavedrom {
105    fn add_wavedrom(chapter: &mut Chapter) -> Result<String> {
106        add_wavedrom(&chapter.content)
107    }
108}
109
110#[cfg(test)]
111mod test {
112    use pretty_assertions::assert_eq;
113
114    use super::add_wavedrom;
115
116    #[test]
117    fn adds_wavedrom() {
118        let content = r#"# Chapter
119
120```wavedrom
121{signal: [
122  {name: 'clk', wave: 'p.....|...'}
123]}
124```
125
126Text
127"#;
128
129        let expected = r#"# Chapter
130
131
132<body onload="WaveDrom.ProcessAll()">
133
134<script type="WaveDrom">{signal: [
135  {name: 'clk', wave: 'p.....|...'}
136]}
137</script>
138
139
140
141Text
142"#;
143
144        assert_eq!(expected, add_wavedrom(content).unwrap());
145    }
146
147    #[test]
148    fn leaves_tables_untouched() {
149        let content = r#"# Heading
153
154| Head 1 | Head 2 |
155|--------|--------|
156| Row 1  | Row 2  |
157"#;
158
159        let expected = r#"# Heading
160
161| Head 1 | Head 2 |
162|--------|--------|
163| Row 1  | Row 2  |
164"#;
165
166        assert_eq!(expected, add_wavedrom(content).unwrap());
167    }
168
169    #[test]
170    fn leaves_html_untouched() {
171        let content = r#"# Heading
175
176<del>
177
178*foo*
179
180</del>
181"#;
182
183        let expected = r#"# Heading
184
185<del>
186
187*foo*
188
189</del>
190"#;
191
192        assert_eq!(expected, add_wavedrom(content).unwrap());
193    }
194
195    #[test]
196    fn html_in_list() {
197        let content = r#"# Heading
201
2021. paragraph 1
203   ```
204   code 1
205   ```
2062. paragraph 2
207"#;
208
209        let expected = r#"# Heading
210
2111. paragraph 1
212   ```
213   code 1
214   ```
2152. paragraph 2
216"#;
217
218        assert_eq!(expected, add_wavedrom(content).unwrap());
219    }
220
221    #[test]
222    fn escape_in_wavedrom_block() {
223        env_logger::init();
224        let content = r#"
225```wavedrom
226classDiagram
227    class PingUploader {
228        <<interface>>
229        +Upload() UploadResult
230    }
231```
232
233hello
234"#;
235
236        let expected = r#"
237
238<body onload="WaveDrom.ProcessAll()">
239
240<script type="WaveDrom">classDiagram
241    class PingUploader {
242        <<interface>>
243        +Upload() UploadResult
244    }
245</script>
246
247
248
249hello
250"#;
251
252        assert_eq!(expected, add_wavedrom(content).unwrap());
253    }
254
255}