rune_tracing_macros/
instrument.rs

1use proc_macro2::TokenStream;
2use quote::quote;
3use syn::Token;
4
5/// An internal call to the macro.
6#[derive(Default)]
7pub struct Attr {
8    span: Option<syn::Expr>,
9}
10
11impl syn::parse::Parse for Attr {
12    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
13        let mut attr = Attr::default();
14        let mut last = false;
15
16        while !input.is_empty() {
17            let ident = input.parse::<syn::Ident>()?;
18
19            if ident == "span" {
20                input.parse::<Token![=]>()?;
21                attr.span = Some(input.parse()?);
22            } else {
23                return Err(syn::Error::new_spanned(ident, "Unsupported attribute"));
24            }
25
26            if last {
27                break;
28            }
29
30            last = input.parse::<Option<Token![,]>>()?.is_none();
31        }
32
33        Ok(attr)
34    }
35}
36
37/// An internal call to the macro.
38pub struct Expander {
39    attrs: Vec<syn::Attribute>,
40    vis: syn::Visibility,
41    sig: syn::Signature,
42    remaining: TokenStream,
43}
44
45impl syn::parse::Parse for Expander {
46    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
47        Ok(Self {
48            attrs: input.call(syn::Attribute::parse_outer)?,
49            vis: input.parse()?,
50            sig: input.parse()?,
51            remaining: input.parse()?,
52        })
53    }
54}
55
56impl Expander {
57    pub(super) fn expand(self, attr: &Attr) -> TokenStream {
58        let mut it = self.sig.inputs.iter();
59
60        let first = match it.next() {
61            Some(syn::FnArg::Typed(ty)) => match &*ty.pat {
62                syn::Pat::Ident(ident) => Some(&ident.ident),
63                _ => None,
64            },
65            _ => None,
66        };
67
68        let ident = &self.sig.ident;
69
70        let log = match first {
71            Some(a) => {
72                let ident = syn::LitStr::new(&ident.to_string(), ident.span());
73
74                match &attr.span {
75                    Some(span) => Some(quote! {
76                        let _instrument_span = ::tracing::span!(::tracing::Level::TRACE, #ident);
77                        let _instrument_enter = _instrument_span.enter();
78
79                        if ::tracing::enabled!(::tracing::Level::TRACE) {
80                            if let Some(source) = #a.q.sources.source(#a.source_id, Spanned::span(&#span)) {
81                                ::tracing::trace!("{:?}", source);
82                            }
83                        }
84                    }),
85                    None => Some(quote! {
86                        let _instrument_span = ::tracing::span!(::tracing::Level::TRACE, #ident);
87                        let _instrument_enter = _instrument_span.enter();
88                        ::tracing::trace!("entering");
89                    }),
90                }
91            }
92            _ => None,
93        };
94
95        let attrs = &self.attrs;
96        let vis = &self.vis;
97        let sig = &self.sig;
98        let remaining = &self.remaining;
99
100        quote! {
101            #(#attrs)*
102            #vis #sig {
103                #log
104                #remaining
105            }
106        }
107    }
108}