aquamarine/
attrs.rs

1use itertools::Itertools;
2use proc_macro2::TokenStream;
3use proc_macro_error::{abort, emit_call_site_warning};
4use quote::quote;
5use std::iter;
6use syn::{Attribute, Ident, MetaNameValue};
7
8const MERMAID_JS_LOCAL: &str = "../mermaid.min.js";
9const MERMAID_JS_CDN: &str = "https://unpkg.com/mermaid@9.1.5/dist/mermaid.min.js";
10
11const UNEXPECTED_ATTR_ERROR: &str =
12    "unexpected attribute inside a diagram definition: only #[doc] is allowed";
13
14#[derive(Clone, Default)]
15pub struct Attrs(Vec<Attr>);
16
17#[derive(Clone)]
18pub enum Attr {
19    /// Attribute that is to be forwarded as-is
20    Forward(Attribute),
21    /// Doc comment that cannot be forwarded as-is
22    DocComment(Ident, String),
23    /// Diagram start token
24    DiagramStart(Ident),
25    /// Diagram entry (line)
26    DiagramEntry(Ident, String),
27    /// Diagram end token
28    DiagramEnd(Ident),
29}
30
31impl Attr {
32    pub fn as_ident(&self) -> Option<&Ident> {
33        match self {
34            Attr::Forward(attr) => attr.path.get_ident(),
35            Attr::DocComment(ident, _) => Some(ident),
36            Attr::DiagramStart(ident) => Some(ident),
37            Attr::DiagramEntry(ident, _) => Some(ident),
38            Attr::DiagramEnd(ident) => Some(ident),
39        }
40    }
41
42    pub fn is_diagram_end(&self) -> bool {
43        match self {
44            Attr::DiagramEnd(_) => true,
45            _ => false
46        }
47    }
48    
49    pub fn is_diagram_start(&self) -> bool {
50        match self {
51            Attr::DiagramStart(_) => true,
52            _ => false
53        }
54    }
55    
56    pub fn expect_diagram_entry_text(&self) -> &str {
57        match self {
58            Attr::DiagramEntry(_, body) => body.as_str(),
59            _ => abort!(self.as_ident(), UNEXPECTED_ATTR_ERROR),
60        }
61    }
62}
63
64impl From<Vec<Attribute>> for Attrs {
65    fn from(attrs: Vec<Attribute>) -> Self {
66        let mut out = Attrs::default();
67        out.push_attrs(attrs);
68        out
69    }
70}
71
72impl quote::ToTokens for Attrs {
73    fn to_tokens(&self, tokens: &mut TokenStream) {
74        let mut attrs = self.0.iter();
75        while let Some(attr) = attrs.next() {
76            match attr {
77                Attr::Forward(attr) => attr.to_tokens(tokens),
78                Attr::DocComment(_, comment) => tokens.extend(quote! {
79                    #[doc = #comment]
80                }),
81                Attr::DiagramStart(_) => {
82                    let diagram = attrs
83                        .by_ref()
84                        .take_while(|x| !x.is_diagram_end())
85                        .map(Attr::expect_diagram_entry_text);
86
87                    tokens.extend(generate_diagram_rustdoc(diagram));
88                }
89                // If that happens, then the parsing stage is faulty: doc comments outside of
90                // in between Start and End tokens are to be emitted as Attr::Forward or Attr::DocComment
91                Attr::DiagramEntry(_, body) => {
92                    emit_call_site_warning!("encountered an unexpected attribute that's going to be ignored, this is a bug! ({})", body);
93                }
94                Attr::DiagramEnd(_) => (),
95            }
96        }
97    }
98}
99
100const MERMAID_INIT_SCRIPT: &str = r#"
101    var amrn_mermaid_theme = 'default';
102    if(typeof currentTheme !== 'undefined') {
103        let docs_theme = currentTheme.href;
104        let is_dark = /.*(dark|ayu).*\.css/.test(docs_theme)
105        if(is_dark) {
106            amrn_mermaid_theme = 'dark'
107        }
108    } else {
109        console.log("currentTheme is undefined, are we not inside rustdoc?");
110    }
111    mermaid.initialize({'startOnLoad':'true', 'theme': amrn_mermaid_theme, 'logLevel': 3 });
112"#;
113
114fn generate_diagram_rustdoc<'a>(parts: impl Iterator<Item = &'a str>) -> TokenStream {
115    let preamble = iter::once(r#"<div class="mermaid">"#);
116    let postamble = iter::once("</div>");
117
118    let mermaid_js_load_primary = format!(r#"<script src="{}"></script>"#, MERMAID_JS_LOCAL);
119    let mermaid_js_load_fallback = format!(r#"<script>window.mermaid || document.write('<script src="{}" crossorigin="anonymous"><\/script>')</script>"#, MERMAID_JS_CDN);
120
121    let mermaid_js_init = format!(r#"<script>{}</script>"#, MERMAID_INIT_SCRIPT);
122
123    let body = preamble.chain(parts).chain(postamble).join("\n");
124
125    quote! {
126        #[doc = #mermaid_js_load_primary]
127        #[doc = #mermaid_js_load_fallback]
128        #[doc = #mermaid_js_init]
129        #[doc = #body]
130    }
131}
132
133impl Attrs {
134    pub fn push_attrs(&mut self, attrs: Vec<Attribute>) {
135        use syn::Lit::*;
136        use syn::Meta::*;
137
138        let mut current_location = Location::OutsideDiagram;
139        let mut diagram_start_ident = None;
140
141        for attr in attrs {
142            match &attr.parse_meta() {
143                Ok(NameValue(MetaNameValue {
144                    lit: Str(s), path, ..
145                })) if path.is_ident("doc") => {
146                    let ident = path.get_ident().unwrap();
147                    for attr in split_attr_body(ident, &s.value(), &mut current_location) {
148                        if attr.is_diagram_start() {
149                            diagram_start_ident.replace(ident.clone());
150                        }
151                        self.0.push(attr);
152                    }
153                }
154                _ => {
155                    if let Location::InsideDiagram = current_location {
156                        abort!(attr, UNEXPECTED_ATTR_ERROR)
157                    } else {
158                        self.0.push(Attr::Forward(attr))
159                    }
160                }
161            }
162        }
163
164        if current_location.is_inside() {
165            abort!(diagram_start_ident, "diagram code block is not terminated");
166        }
167    }
168}
169
170#[derive(Debug, Copy, Clone, Eq, PartialEq)]
171enum Location {
172    OutsideDiagram,
173    InsideDiagram,
174}
175
176impl Location {
177    fn is_inside(self) -> bool {
178        match self {
179            Location::InsideDiagram => true, 
180            _ => false
181        }
182    }
183}
184
185fn split_attr_body(ident: &Ident, input: &str, loc: &mut Location) -> Vec<Attr> {
186    use self::Location::*;
187
188    const TICKS: &str = "```";
189    const MERMAID: &str = "mermaid";
190
191    let mut tokens = tokenize_doc_str(input).peekable();
192
193    // Special case: empty strings outside the diagram span should be still generated
194    if tokens.peek().is_none() && !loc.is_inside() {
195        return vec![Attr::DocComment(ident.clone(), String::new())];
196    };
197
198    // To aid rustc with type inference in closures
199    #[derive(Default)]
200    struct Ctx<'a> {
201        attrs: Vec<Attr>,
202        buffer: Vec<&'a str>,
203    }
204
205    let mut ctx = Default::default();
206
207    let flush_buffer_as_doc_comment = |ctx: &mut Ctx| {
208        if !ctx.buffer.is_empty() {
209            ctx.attrs.push(Attr::DocComment(
210                ident.clone(),
211                ctx.buffer.drain(..).join(" "),
212            ));
213        }
214    };
215
216    let flush_buffer_as_diagram_entry = |ctx: &mut Ctx| {
217        let s = ctx.buffer.drain(..).join(" ");
218        if !s.trim().is_empty() {
219            ctx.attrs.push(Attr::DiagramEntry(ident.clone(), s));
220        }
221    };
222
223    while let Some(token) = tokens.next() {
224        match (*loc, token, tokens.peek()) {
225            // Flush the buffer, then open the diagram code block
226            (OutsideDiagram, TICKS, Some(&MERMAID)) => {
227                tokens.next();
228                *loc = InsideDiagram;
229                flush_buffer_as_doc_comment(&mut ctx);
230                ctx.attrs.push(Attr::DiagramStart(ident.clone()));
231            }
232            // Flush the buffer, close the code block
233            (InsideDiagram, TICKS, _) => {
234                *loc = OutsideDiagram;
235                flush_buffer_as_diagram_entry(&mut ctx);
236                ctx.attrs.push(Attr::DiagramEnd(ident.clone()))
237            }
238            _ => ctx.buffer.push(token),
239        }
240    }
241
242    if !ctx.buffer.is_empty() {
243        if loc.is_inside() {
244            flush_buffer_as_diagram_entry(&mut ctx);
245        } else {
246            flush_buffer_as_doc_comment(&mut ctx);
247        };
248    }
249
250    ctx.attrs
251}
252
253fn tokenize_doc_str(input: &str) -> impl Iterator<Item = &str> {
254    const TICKS: &str = "```";
255    split_inclusive(input, TICKS).flat_map(|token| {
256        // not str::split_whitespace because we don't wanna filter-out the whitespace tokens
257        token.split(' ')
258    })
259}
260
261// TODO: remove once str::split_inclusive is stable
262fn split_inclusive<'a, 'b: 'a>(input: &'a str, delim: &'b str) -> impl Iterator<Item = &'a str> {
263    let mut tokens = vec![];
264    let mut prev = 0;
265
266    for (idx, matches) in input.match_indices(delim) {
267        tokens.extend(nonempty(&input[prev..idx]));
268
269        prev = idx + matches.len();
270
271        tokens.push(matches);
272    }
273
274    if prev < input.len() {
275        tokens.push(&input[prev..]);
276    }
277
278    tokens.into_iter()
279}
280
281fn nonempty(s: &str) -> Option<&str> {
282    if s.is_empty() {
283        None
284    } else {
285        Some(s)
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use super::{split_inclusive, Attr};
292    use std::fmt;
293
294    #[cfg(test)]
295    impl fmt::Debug for Attr {
296        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
297            match self {
298                Attr::Forward(..) => f.write_str("Attr::Forward"),
299                Attr::DocComment(_, body) => write!(f, "Attr::DocComment({:?})", body),
300                Attr::DiagramStart(..) => f.write_str("Attr::DiagramStart"),
301                Attr::DiagramEntry(_, body) => write!(f, "Attr::DiagramEntry({:?})", body),
302                Attr::DiagramEnd(..) => f.write_str("Attr::DiagramEnd"),
303            }
304        }
305    }
306
307    #[cfg(test)]
308    impl Eq for Attr {}
309
310    #[cfg(test)]
311    impl PartialEq for Attr {
312        fn eq(&self, other: &Self) -> bool {
313            use std::mem::discriminant;
314            use Attr::*;
315            match (self, other) {
316                (DocComment(_, a), DocComment(_, b)) => a == b,
317                (DiagramEntry(_, a), DiagramEntry(_, b)) => a == b,
318                (a, b) => discriminant(a) == discriminant(b),
319            }
320        }
321    }
322
323    #[test]
324    fn temp_split_inclusive() {
325        let src = "```";
326        let out: Vec<_> = split_inclusive(src, "```").collect();
327        assert_eq!(&out, &["```",]);
328
329        let src = "```abcd```";
330        let out: Vec<_> = split_inclusive(src, "```").collect();
331        assert_eq!(&out, &["```", "abcd", "```"]);
332
333        let src = "left```abcd```right";
334        let out: Vec<_> = split_inclusive(src, "```").collect();
335        assert_eq!(&out, &["left", "```", "abcd", "```", "right",]);
336    }
337
338    mod split_attr_body_tests {
339        use super::super::*;
340
341        use proc_macro2::Ident;
342        use proc_macro2::Span;
343
344        use pretty_assertions::assert_eq;
345
346        fn i() -> Ident {
347            Ident::new("fake", Span::call_site())
348        }
349
350        struct TestCase<'a> {
351            ident: Ident,
352            location: Location,
353            input: &'a str,
354            expect_location: Location,
355            expect_attrs: Vec<Attr>,
356        }
357
358        fn check(case: TestCase) {
359            let mut loc = case.location;
360            let attrs = split_attr_body(&case.ident, case.input, &mut loc);
361            assert_eq!(loc, case.expect_location);
362            assert_eq!(attrs, case.expect_attrs);
363        }
364
365        #[test]
366        fn one_line_one_diagram() {
367            let case = TestCase {
368                ident: i(),
369                location: Location::OutsideDiagram,
370                input: "```mermaid abcd```",
371                expect_location: Location::OutsideDiagram,
372                expect_attrs: vec![
373                    Attr::DiagramStart(i()),
374                    Attr::DiagramEntry(i(), "abcd".into()),
375                    Attr::DiagramEnd(i()),
376                ],
377            };
378
379            check(case)
380        }
381
382        #[test]
383        fn one_line_multiple_diagrams() {
384            let case = TestCase {
385                ident: i(),
386                location: Location::OutsideDiagram,
387                input: "```mermaid abcd``` ```mermaid efgh``` ```mermaid ijkl```",
388                expect_location: Location::OutsideDiagram,
389                expect_attrs: vec![
390                    Attr::DiagramStart(i()),
391                    Attr::DiagramEntry(i(), "abcd".into()),
392                    Attr::DiagramEnd(i()),
393                    Attr::DocComment(i(), " ".into()),
394                    Attr::DiagramStart(i()),
395                    Attr::DiagramEntry(i(), "efgh".into()),
396                    Attr::DiagramEnd(i()),
397                    Attr::DocComment(i(), " ".into()),
398                    Attr::DiagramStart(i()),
399                    Attr::DiagramEntry(i(), "ijkl".into()),
400                    Attr::DiagramEnd(i()),
401                ],
402            };
403
404            check(case)
405        }
406
407        #[test]
408        fn other_snippet() {
409            let case = TestCase {
410                ident: i(),
411                location: Location::OutsideDiagram,
412                input: "```rust panic!()```",
413                expect_location: Location::OutsideDiagram,
414                expect_attrs: vec![Attr::DocComment(i(), "``` rust panic!() ```".into())],
415            };
416
417            check(case)
418        }
419
420        #[test]
421        fn carry_over() {
422            let case = TestCase {
423                ident: i(),
424                location: Location::OutsideDiagram,
425                input: "left```mermaid abcd```right",
426                expect_location: Location::OutsideDiagram,
427                expect_attrs: vec![
428                    Attr::DocComment(i(), "left".into()),
429                    Attr::DiagramStart(i()),
430                    Attr::DiagramEntry(i(), "abcd".into()),
431                    Attr::DiagramEnd(i()),
432                    Attr::DocComment(i(), "right".into()),
433                ],
434            };
435
436            check(case)
437        }
438
439        #[test]
440        fn multiline_termination() {
441            let case = TestCase {
442                ident: i(),
443                location: Location::InsideDiagram,
444                input: "abcd```",
445                expect_location: Location::OutsideDiagram,
446                expect_attrs: vec![
447                    Attr::DiagramEntry(i(), "abcd".into()),
448                    Attr::DiagramEnd(i()),
449                ],
450            };
451
452            check(case)
453        }
454
455        #[test]
456        fn multiline_termination_single_token() {
457            let case = TestCase {
458                ident: i(),
459                location: Location::InsideDiagram,
460                input: "```",
461                expect_location: Location::OutsideDiagram,
462                expect_attrs: vec![Attr::DiagramEnd(i())],
463            };
464
465            check(case)
466        }
467
468        #[test]
469        fn multiline_termination_carry() {
470            let case = TestCase {
471                ident: i(),
472                location: Location::InsideDiagram,
473                input: "abcd```right",
474                expect_location: Location::OutsideDiagram,
475                expect_attrs: vec![
476                    Attr::DiagramEntry(i(), "abcd".into()),
477                    Attr::DiagramEnd(i()),
478                    Attr::DocComment(i(), "right".into()),
479                ],
480            };
481
482            check(case)
483        }
484    }
485}
aquamarine/attrs.rs

aquamarine/
attrs.rs
