Eager proc macro expansion

Question

Given a TokenStream, can I produce a new TokenStream with all the attributes pre-expanded?

I suspect this would dramatically increase compile time, or require multiple compile steps.

I've tried to take a stab at writing a proc macro that can execute a closure on a function exit. The exact implementation looks like the following

#[proc_macro_attribute]
pub fn inspect(attr: TokenStream, item: TokenStream) -> TokenStream {
    let mut callback = parse_macro_input!(attr as ExprClosure);
    let function = parse_macro_input!(item as ItemFn);
    let sig = &function.sig;
    let body = &function.block;

    if let Some(compile_error) = validate_inspect(&callback, &function) {
        return compile_error;
    }

    let out = if sig.asyncness.is_some() {
        quote! { let out = async move { #body }.await; }
    } else {
        quote! { let out = (move || { #body })(); }
    };

    quote! {
        #sig {
            #out
            (#callback)(&out);
            out
        }
    }
    .into()
}

It may be used as follows:

#[inspect(|out| println!("increment returns {i}"))]
async fn increment(value: i32) -> i32 {
    value + 1
}

The problem is that this does not play well with async_trait as the proc_macro for async_trait removes the asyncness of the function.

#[async_trait]
trait Foo {
    async fn foo() -> i32;
}

struct Bar;

/// ! Compile error
#[async_trait]
impl Foo for Bar {
    #[inspect(|out| println!("foo returns {}", i))]
    async fn foo() -> i32 {
        42
    }
}

The following example does not compile because proc-macros are expanded outside in.

Without my proc macro,

#[async_trait]
impl Foo for Bar {
    async fn foo() -> i32 {
        42
    }
}

async_trait expands to roughly something like this¹:

impl Foo for Bar {
    fn foo<'async_trait>() -> Pin<
        Box<dyn Future<Output = i32> + Send + 'async_trait>,
    > {
        Box::pin(async move {
            42
        })
    }
}

As you can see the function is converted to a synchronous function returning a dynamic future since async traits are not yet stabilized.

I would like to have a proc macro called invert_expansion_trait which expands the attributes of items inside a trait implementation first, before running it's own subsequent attributes. Would this be possible to implement today?

I tried creating a eager proc macro expansion that expands the attributes of an ItemImpl but couldn't get it to work because I lack knowledge / resources to understand the expansion process.

---

_{1: some code removed from expansion for ease of viewing}

Answer 1

Per @chayim-friedman, I went with option 1, to eagerly inspect and expand the proc macros manually. It looks like this:

/// Applies internal proc macro expansions eagerly onto trait implementations
///
/// Proc macros work outside in, therefore, proc macros which affect the function signature will
/// cause `inspect` and `trap` to fail. In order to avoid this, apply a `hook_trait` at the beginning
/// of an impl block. The macro will manually expand the internal proc macros first, before external
/// proc macros may affect the signature or implementation.
#[proc_macro_attribute]
pub fn hook_trait(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let mut item_impl = parse_macro_input!(item as ItemImpl);

    for impl_item_fn in item_impl.items.iter_mut().filter_map(|item| match item {
        ImplItem::Fn(impl_item_fn) => Some(impl_item_fn),
        _ => None,
    }) {
        // Extract the proc macro attributes
        let mut attrs = vec![];
        impl_item_fn.attrs = impl_item_fn
            .attrs
            .drain(..)
            .filter_map(|attr| {
                if attr.path().is_ident(stringify!(inspect)) {
                    attrs.push((inspect as ProcMacroFn, attr));
                    None
                } else if attr.path().is_ident(stringify!(trap)) {
                    attrs.push((trap as ProcMacroFn, attr));
                    None
                } else {
                    Some(attr)
                }
            })
            .collect();

        // Eagerly apply the proc macro to each implementation
        for (proc_macro, attr) in attrs {
            let Meta::List(meta_list) = attr.meta else {
                return Error::new(attr.span(), "Attributes requires a metadata parameters")
                    .to_compile_error()
                    .into();
            };

            let tokens = proc_macro(
                meta_list.tokens.to_token_stream().into(),
                impl_item_fn.to_token_stream().into(),
            );
            match parse(tokens.clone()) {
                Ok(new_impl_item_fn) => *impl_item_fn = new_impl_item_fn,
                Err(_) => return tokens,
            }
        }
    }

    item_impl.to_token_stream().into()
}

Answer 2

You cannot expand that within the macro.

On nightly, there exists TokenStream::expand_expr(), but it only supports expanding expressions (not items). Relaxing this is an unresolved question in the tracking issue, but it is not implemented currently.

There are two possible solutions that exist in the ecosystem.

The first is to define your macro on the trait too, with attributes on the specific functions to know which macros to expand. Since macros are expanded in order, if you place your attribute above #[async_trait] it will expand before. You probably need to document that too.

The second, more limited but also more powerful, is to do the same, but in addition, have your macro capture the #[async_trait] and expand it too. This can work even if async-trait is unable to work with the expansion of your macro, however of course it is limited to attributes your macro knows.