async_backtrace_attributes/
lib.rs

1use proc_macro2::TokenStream;
2use quote::ToTokens;
3use syn::parse::{Parse, ParseStream};
4use syn::{Attribute, Block, ItemFn, Signature, Visibility};
5
6mod expand;
7
8#[proc_macro_attribute]
9pub fn framed(
10    args: proc_macro::TokenStream,
11    item: proc_macro::TokenStream,
12) -> proc_macro::TokenStream {
13    assert!(args.is_empty());
14    // Cloning a `TokenStream` is cheap since it's reference counted internally.
15    instrument_precise(item.clone()).unwrap_or_else(|_err| instrument_speculative(item))
16}
17
18/// Instrument the function, without parsing the function body (instead using
19/// the raw tokens).
20fn instrument_speculative(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
21    let input = syn::parse_macro_input!(item as MaybeItemFn);
22    let instrumented_function_name = input.sig.ident.to_string();
23    expand::gen_function(input.as_ref(), instrumented_function_name.as_str(), None).into()
24}
25
26/// Instrument the function, by fully parsing the function body,
27/// which allows us to rewrite some statements related to async-like patterns.
28fn instrument_precise(
29    item: proc_macro::TokenStream,
30) -> Result<proc_macro::TokenStream, syn::Error> {
31    let input = syn::parse::<ItemFn>(item)?;
32    let instrumented_function_name = input.sig.ident.to_string();
33
34    // check for async_trait-like patterns in the block, and instrument
35    // the future instead of the wrapper
36    if let Some(async_like) = expand::AsyncInfo::from_fn(&input) {
37        return Ok(async_like.gen_async(instrumented_function_name.as_str()));
38    }
39
40    Ok(expand::gen_function((&input).into(), instrumented_function_name.as_str(), None).into())
41}
42
43/// This is a more flexible/imprecise `ItemFn` type,
44/// which's block is just a `TokenStream` (it may contain invalid code).
45#[derive(Debug, Clone)]
46struct MaybeItemFn {
47    attrs: Vec<Attribute>,
48    vis: Visibility,
49    sig: Signature,
50    block: TokenStream,
51}
52
53impl MaybeItemFn {
54    fn as_ref(&self) -> MaybeItemFnRef<'_, TokenStream> {
55        MaybeItemFnRef {
56            attrs: &self.attrs,
57            vis: &self.vis,
58            sig: &self.sig,
59            block: &self.block,
60        }
61    }
62}
63
64/// This parses a `TokenStream` into a `MaybeItemFn`
65/// (just like `ItemFn`, but skips parsing the body).
66impl Parse for MaybeItemFn {
67    fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
68        let attrs = input.call(syn::Attribute::parse_outer)?;
69        let vis: Visibility = input.parse()?;
70        let sig: Signature = input.parse()?;
71        let block: TokenStream = input.parse()?;
72        Ok(Self {
73            attrs,
74            vis,
75            sig,
76            block,
77        })
78    }
79}
80
81/// A generic reference type for `MaybeItemFn`,
82/// that takes a generic block type `B` that implements `ToTokens` (eg.
83/// `TokenStream`, `Block`).
84#[derive(Debug, Clone)]
85struct MaybeItemFnRef<'a, B: ToTokens> {
86    attrs: &'a Vec<Attribute>,
87    vis: &'a Visibility,
88    sig: &'a Signature,
89    block: &'a B,
90}
91
92impl<'a> From<&'a ItemFn> for MaybeItemFnRef<'a, Box<Block>> {
93    fn from(val: &'a ItemFn) -> Self {
94        MaybeItemFnRef {
95            attrs: &val.attrs,
96            vis: &val.vis,
97            sig: &val.sig,
98            block: &val.block,
99        }
100    }
101}
async_backtrace_attributes/lib.rs

async_backtrace_attributes/
lib.rs
