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}
57