1
1
use proc_macro:: TokenStream ;
2
- use quote:: quote;
2
+ use proc_macro2:: Span ;
3
+ use quote:: { quote, quote_spanned, ToTokens } ;
3
4
4
5
pub ( crate ) fn test ( args : TokenStream , item : TokenStream ) -> TokenStream {
5
6
if !args. is_empty ( ) {
@@ -9,23 +10,45 @@ pub(crate) fn test(args: TokenStream, item: TokenStream) -> TokenStream {
9
10
}
10
11
11
12
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 ;
16
13
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" )
19
16
. to_compile_error ( )
20
17
. into ( ) ;
21
18
}
22
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
+ ) ] ;
48
+
23
49
let gen = quote ! {
24
50
#[ :: core:: prelude:: v1:: test]
25
- #( #attrs) *
26
- #vis #sig {
27
- :: futures_test:: __private:: block_on( async move #body)
28
- }
51
+ #input
29
52
} ;
30
53
31
54
gen. into ( )
0 commit comments