spandoc_attribute/
lib.rs

1//! proc macro crate that implements the `#[spandoc::spandoc]` attribute. See
2//[`spandoc`](https://docs.rs/spandoc) documentation for details.
3#![doc(html_root_url = "https://docs.rs/spandoc-attribute/0.1.0")]
4#![cfg_attr(docsrs, feature(doc_cfg))]
5#![warn(
6    missing_docs,
7    missing_doc_code_examples,
8    rust_2018_idioms,
9    unreachable_pub,
10    bad_style,
11    const_err,
12    dead_code,
13    improper_ctypes,
14    non_shorthand_field_patterns,
15    no_mangle_generic_items,
16    overflowing_literals,
17    path_statements,
18    patterns_in_fns_without_body,
19    private_in_public,
20    unconditional_recursion,
21    unused,
22    unused_allocation,
23    unused_comparisons,
24    unused_parens,
25    while_true
26)]
27
28use proc_macro::TokenStream;
29use proc_macro2::Ident;
30use quote::quote_spanned;
31use syn::{
32    fold::Fold, spanned::Spanned, Attribute, AttributeArgs, Block, ExprAsync, ExprAwait, ItemFn,
33    Meta, Signature,
34};
35
36#[proc_macro_attribute]
37/// entrypoint for spandoc attribute proc macro
38pub fn spandoc(args: TokenStream, item: TokenStream) -> TokenStream {
39    let input: ItemFn = syn::parse_macro_input!(item as ItemFn);
40    let _args = syn::parse_macro_input!(args as AttributeArgs);
41
42    let span = input.span();
43    let ItemFn {
44        attrs,
45        vis,
46        block,
47        sig,
48        ..
49    } = input;
50
51    let Signature { ref ident, .. } = sig;
52
53    let block = SpanInstrumentedExpressions {
54        ident: ident.clone(),
55    }
56    .fold_block(*block);
57
58    quote_spanned!( span =>
59        #(#attrs) *
60        #[allow(clippy::cognitive_complexity)]
61        #vis #sig
62        #block
63    )
64    .into()
65}
66
67struct InstrumentAwaits;
68
69impl Fold for InstrumentAwaits {
70    fn fold_expr_async(&mut self, i: ExprAsync) -> ExprAsync {
71        i
72    }
73
74    fn fold_expr_await(&mut self, i: ExprAwait) -> ExprAwait {
75        let mut i = syn::fold::fold_expr_await(self, i);
76        let span = i.span();
77        let base = i.base;
78        let base = quote_spanned! { span => __fancy_guard.wrap(#base) };
79
80        let base = syn::parse2(base).unwrap();
81        i.base = Box::new(base);
82        i
83    }
84}
85
86struct SpanInstrumentedExpressions {
87    ident: Ident,
88}
89
90impl Fold for SpanInstrumentedExpressions {
91    fn fold_block(&mut self, block: Block) -> Block {
92        if block.stmts.is_empty() {
93            return block;
94        }
95
96        let block_span = block.span();
97        let mut block = syn::fold::fold_block(self, block);
98
99        let stmts = block.stmts;
100        let mut new_stmts = proc_macro2::TokenStream::new();
101        let last = stmts.len() - 1;
102
103        for (i, mut stmt) in stmts.into_iter().enumerate() {
104            let stmt_span = stmt.span();
105
106            let as_span = |attr: Attribute| {
107                let meta = attr.parse_meta().ok()?;
108                let lit = match meta {
109                    Meta::NameValue(syn::MetaNameValue {
110                        lit: syn::Lit::Str(lit),
111                        ..
112                    }) => lit,
113                    _ => return None,
114                };
115
116                let (lit, args) = args::split(lit)?;
117                let span_name = format!("{}::comment", self.ident);
118
119                let span = match args {
120                    Some(args) => {
121                        quote_spanned! { lit.span() =>
122                            tracing::span!(tracing::Level::ERROR, #span_name, #args, text = %#lit)
123                        }
124                    }
125                    None => quote_spanned! { lit.span() =>
126                        tracing::span!(tracing::Level::ERROR, #span_name, text = %#lit)
127                    },
128                };
129
130                Some(span)
131            };
132
133            let attrs = if let Some(attrs) = attr::from_stmt(&mut stmt) {
134                attrs
135            } else {
136                new_stmts.extend(quote_spanned! { stmt_span => #stmt });
137                continue;
138            };
139
140            let ind = if let Some(ind) = attr::find_doc_attr_ind(attrs) {
141                ind
142            } else {
143                new_stmts.extend(quote_spanned! { stmt_span => #stmt });
144                continue;
145            };
146
147            let attr = attrs[ind].clone();
148            let span = as_span(attr);
149
150            let stmt = if span.is_some() {
151                attrs.remove(ind);
152                InstrumentAwaits.fold_stmt(stmt)
153            } else {
154                stmt
155            };
156
157            let stmts = match span {
158                Some(span) if i == last => {
159                    quote_spanned! { stmt_span =>
160                        let __dummy_span = #span;
161                        let __fancy_guard = spandoc::FancyGuard::new(&__dummy_span);
162                        #stmt
163                    }
164                }
165                Some(span) => {
166                    quote_spanned! { stmt_span =>
167                        let __dummy_span = #span;
168                        let __fancy_guard = spandoc::FancyGuard::new(&__dummy_span);
169                        #stmt
170                        drop(__fancy_guard);
171                        drop(__dummy_span);
172                    }
173                }
174                _ => quote_spanned! { stmt_span => #stmt },
175            };
176
177            new_stmts.extend(stmts);
178        }
179
180        let new_block = quote_spanned! { block_span =>
181            {
182                #new_stmts
183            }
184        };
185
186        let new_block: Block = syn::parse2(new_block).unwrap();
187
188        block.stmts = new_block.stmts;
189        block
190    }
191}
192
193mod attr {
194    use syn::{Attribute, Expr, Stmt};
195
196    pub(crate) fn from_stmt(stmt: &mut Stmt) -> Option<&mut Vec<Attribute>> {
197        match stmt {
198            syn::Stmt::Local(local) => Some(&mut local.attrs),
199            syn::Stmt::Item(_) => None,
200            syn::Stmt::Expr(expr) => from_expr(expr),
201            syn::Stmt::Semi(expr, ..) => from_expr(expr),
202        }
203    }
204
205    fn from_expr(expr: &mut Expr) -> Option<&mut Vec<Attribute>> {
206        match expr {
207            Expr::Array(e) => Some(&mut e.attrs),
208            Expr::Assign(e) => Some(&mut e.attrs),
209            Expr::AssignOp(e) => Some(&mut e.attrs),
210            Expr::Async(e) => Some(&mut e.attrs),
211            Expr::Await(e) => Some(&mut e.attrs),
212            Expr::Binary(e) => Some(&mut e.attrs),
213            Expr::Block(e) => Some(&mut e.attrs),
214            Expr::Box(e) => Some(&mut e.attrs),
215            Expr::Break(e) => Some(&mut e.attrs),
216            Expr::Call(e) => Some(&mut e.attrs),
217            Expr::Cast(e) => Some(&mut e.attrs),
218            Expr::Closure(e) => Some(&mut e.attrs),
219            Expr::Continue(e) => Some(&mut e.attrs),
220            Expr::Field(e) => Some(&mut e.attrs),
221            Expr::ForLoop(e) => Some(&mut e.attrs),
222            Expr::Group(e) => Some(&mut e.attrs),
223            Expr::If(e) => Some(&mut e.attrs),
224            Expr::Index(e) => Some(&mut e.attrs),
225            Expr::Let(e) => Some(&mut e.attrs),
226            Expr::Lit(e) => Some(&mut e.attrs),
227            Expr::Loop(e) => Some(&mut e.attrs),
228            Expr::Macro(e) => Some(&mut e.attrs),
229            Expr::Match(e) => Some(&mut e.attrs),
230            Expr::MethodCall(e) => Some(&mut e.attrs),
231            Expr::Paren(e) => Some(&mut e.attrs),
232            Expr::Path(e) => Some(&mut e.attrs),
233            Expr::Range(e) => Some(&mut e.attrs),
234            Expr::Reference(e) => Some(&mut e.attrs),
235            Expr::Repeat(e) => Some(&mut e.attrs),
236            Expr::Return(e) => Some(&mut e.attrs),
237            Expr::Struct(e) => Some(&mut e.attrs),
238            Expr::Try(e) => Some(&mut e.attrs),
239            Expr::TryBlock(e) => Some(&mut e.attrs),
240            Expr::Tuple(e) => Some(&mut e.attrs),
241            Expr::Type(e) => Some(&mut e.attrs),
242            Expr::Unary(e) => Some(&mut e.attrs),
243            Expr::Unsafe(e) => Some(&mut e.attrs),
244            Expr::Verbatim(_) => None,
245            Expr::While(e) => Some(&mut e.attrs),
246            Expr::Yield(e) => Some(&mut e.attrs),
247            _ => None,
248            // some variants omitted
249        }
250    }
251
252    pub(crate) fn find_doc_attr_ind(attrs: &mut Vec<Attribute>) -> Option<usize> {
253        attrs.iter().position(|attr| attr.path.is_ident("doc"))
254    }
255}
256
257mod args {
258    use core::ops::Range;
259    use syn::LitStr;
260
261    pub(crate) fn split(lit: LitStr) -> Option<(LitStr, Option<proc_macro2::TokenStream>)> {
262        let text = lit.value();
263        let text = text.trim();
264        let span = lit.span();
265
266        let text = if !text.starts_with("SPANDOC: ") {
267            return None;
268        } else {
269            text.trim_start_matches("SPANDOC: ")
270        };
271
272        if let Some((text_range, args_range)) = get_ranges(text) {
273            let args = &text[args_range];
274            let text = &text[text_range].trim();
275
276            let lit = LitStr::new(text, span);
277            let args = LitStr::new(args, span);
278            let args = args.parse().unwrap();
279
280            Some((lit, Some(args)))
281        } else {
282            let lit = LitStr::new(text, span);
283            Some((lit, None))
284        }
285    }
286
287    fn get_ranges(text: &str) -> Option<(Range<usize>, Range<usize>)> {
288        let mut depth = 0;
289
290        if !text.ends_with('}') {
291            return None;
292        }
293
294        let chars = text.chars().collect::<Vec<_>>();
295        let len = chars.len();
296
297        for (ind, c) in chars.into_iter().enumerate().rev() {
298            match c {
299                '}' => depth += 1,
300                '{' => depth -= 1,
301                _ => (),
302            }
303
304            if depth == 0 {
305                let end = len - 1;
306                return Some((0..ind, ind + 1..end));
307            }
308        }
309
310        None
311    }
312
313    #[cfg(test)]
314    pub fn split_str(text: &str) -> (&str, Option<&str>) {
315        match get_ranges(text) {
316            Some((text_range, args_range)) => {
317                let args = &text[args_range];
318                let text = &text[text_range].trim();
319
320                (text, Some(args))
321            }
322            _ => (text, None),
323        }
324    }
325}
326
327#[cfg(test)]
328mod tests {
329    use super::*;
330
331    #[test]
332    fn no_args() {
333        let input = "This doesn't have args";
334        let (text, args) = args::split_str(input);
335        assert_eq!(input, text);
336        assert_eq!(None, args);
337    }
338
339    #[test]
340    fn with_args() {
341        let input = "This doesn't have args {but, this, does}";
342        let (text, args) = args::split_str(input);
343        assert_eq!("This doesn't have args", text);
344        assert_eq!(Some("but, this, does"), args);
345    }
346}
spandoc_attribute/lib.rs

spandoc_attribute/
lib.rs
