futures_macro/
executor.rs

1use proc_macro::TokenStream;
2use proc_macro2::Span;
3use quote::{quote, quote_spanned, ToTokens};
4
5pub(crate) fn test(args: TokenStream, item: TokenStream) -> TokenStream {
6    if !args.is_empty() {
7        return syn::Error::new_spanned(proc_macro2::TokenStream::from(args), "invalid argument")
8            .to_compile_error()
9            .into();
10    }
11
12    let mut input = syn::parse_macro_input!(item as syn::ItemFn);
13
14    if input.sig.asyncness.take().is_none() {
15        return syn::Error::new_spanned(input.sig.fn_token, "Only async functions are supported")
16            .to_compile_error()
17            .into();
18    }
19
20    // If type mismatch occurs, the current rustc points to the last statement.
21    let (last_stmt_start_span, last_stmt_end_span) = {
22        let mut last_stmt = input
23            .block
24            .stmts
25            .last()
26            .map(ToTokens::into_token_stream)
27            .unwrap_or_default()
28            .into_iter();
29        // `Span` on stable Rust has a limitation that only points to the first
30        // token, not the whole tokens. We can work around this limitation by
31        // using the first/last span of the tokens like
32        // `syn::Error::new_spanned` does.
33        let start = last_stmt.next().map_or_else(Span::call_site, |t| t.span());
34        let end = last_stmt.last().map_or(start, |t| t.span());
35        (start, end)
36    };
37
38    let path = quote_spanned! {last_stmt_start_span=>
39        ::futures_test::__private
40    };
41    let body = &input.block;
42    input.block.stmts = vec![syn::Stmt::Expr(
43        syn::parse2(quote_spanned! {last_stmt_end_span=>
44            #path::block_on(async #body)
45        })
46        .unwrap(),
47        None,
48    )];
49
50    let gen = quote! {
51        #[::core::prelude::v1::test]
52        #input
53    };
54
55    gen.into()
56}