11use proc_macro:: TokenStream ;
2- use quote:: quote;
2+ use proc_macro2:: Span ;
3+ use quote:: { quote, quote_spanned, ToTokens } ;
34
45pub ( crate ) fn test ( args : TokenStream , item : TokenStream ) -> TokenStream {
56 if !args. is_empty ( ) {
@@ -9,23 +10,45 @@ pub(crate) fn test(args: TokenStream, item: TokenStream) -> TokenStream {
910 }
1011
1112 let mut input = syn:: parse_macro_input!( item as syn:: ItemFn ) ;
12- let attrs = & input. attrs ;
13- let vis = & input. vis ;
14- let sig = & mut input. sig ;
15- let body = & input. block ;
1613
17- if sig. asyncness . take ( ) . is_none ( ) {
18- return syn:: Error :: new_spanned ( sig. fn_token , "Only async functions are supported" )
14+ if input . sig . asyncness . take ( ) . is_none ( ) {
15+ return syn:: Error :: new_spanned ( input . sig . fn_token , "Only async functions are supported" )
1916 . to_compile_error ( )
2017 . into ( ) ;
2118 }
2219
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+ ) ] ;
48+
2349 let gen = quote ! {
2450 #[ :: core:: prelude:: v1:: test]
25- #( #attrs) *
26- #vis #sig {
27- :: futures_test:: __private:: block_on( async move #body)
28- }
51+ #input
2952 } ;
3053
3154 gen. into ( )
0 commit comments