37
37
extern crate proc_macro;
38
38
39
39
use proc_macro:: TokenStream ;
40
- use quote:: quote;
40
+ use proc_macro2:: TokenStream as TokenStream2 ;
41
+ use quote:: { quote, ToTokens } ;
42
+ use syn:: parse:: { self , Parse , ParseStream } ;
43
+ use syn:: Token ;
41
44
42
45
/// Add context to errors from a function.
43
46
///
@@ -55,32 +58,56 @@ use quote::quote;
55
58
/// })().map_err(|err| err.context("context"))
56
59
/// }
57
60
/// ```
61
+ ///
62
+ /// Sometimes you will receive borrowck errors, especially when returning references. These can
63
+ /// often be fixed by setting the `move` option of the attribute macro. For example:
64
+ ///
65
+ /// ```
66
+ /// #[context(move, "context")]
67
+ /// fn returns_reference(val: &mut u32) -> anyhow::Result<&mut u32> {
68
+ /// Ok(&mut *val)
69
+ /// }
70
+ /// ```
58
71
#[ proc_macro_attribute]
59
72
pub fn context ( args : TokenStream , input : TokenStream ) -> TokenStream {
60
- let args : proc_macro2 :: TokenStream = args . into ( ) ;
73
+ let Args ( move_token , format_args ) = syn :: parse_macro_input! ( args ) ;
61
74
let mut input = syn:: parse_macro_input!( input as syn:: ItemFn ) ;
62
75
63
76
let body = & input. block ;
64
77
let return_ty = & input. sig . output ;
65
- if input. sig . asyncness . is_some ( ) {
66
- match return_ty {
78
+ let new_body = if input. sig . asyncness . is_some ( ) {
79
+ let return_ty = match return_ty {
67
80
syn:: ReturnType :: Default => {
68
- return syn:: Error :: new_spanned ( return_ty , "function should return Result" )
81
+ return syn:: Error :: new_spanned ( input , "function should return Result" )
69
82
. to_compile_error ( )
70
- . into ( )
71
- }
72
- syn:: ReturnType :: Type ( _, return_ty) => {
73
- input. block . stmts = syn:: parse_quote!(
74
- let result: #return_ty = async move { #body } . await ;
75
- result. map_err( |err| err. context( format!( #args) ) . into( ) )
76
- ) ;
83
+ . into ( ) ;
77
84
}
85
+ syn:: ReturnType :: Type ( _, return_ty) => return_ty,
86
+ } ;
87
+ quote ! {
88
+ let result: #return_ty = async #move_token { #body } . await ;
89
+ result. map_err( |err| err. context( format!( #format_args) ) . into( ) )
78
90
}
79
91
} else {
80
- input. block . stmts = syn:: parse_quote!(
81
- ( || #return_ty #body) ( ) . map_err( |err| err. context( format!( #args) ) . into( ) )
82
- ) ;
83
- }
92
+ quote ! {
93
+ ( #move_token || #return_ty #body) ( ) . map_err( |err| err. context( format!( #format_args) ) . into( ) )
94
+ }
95
+ } ;
96
+ input. block . stmts = vec ! [ syn:: Stmt :: Expr ( syn:: Expr :: Verbatim ( new_body) ) ] ;
84
97
85
- quote ! ( #input) . into ( )
98
+ input. into_token_stream ( ) . into ( )
99
+ }
100
+
101
+ struct Args ( Option < Token ! [ move] > , TokenStream2 ) ;
102
+ impl Parse for Args {
103
+ fn parse ( input : ParseStream < ' _ > ) -> parse:: Result < Self > {
104
+ let move_token = if input. peek ( Token ! [ move] ) {
105
+ let token = input. parse ( ) ?;
106
+ input. parse :: < Token ! [ , ] > ( ) ?;
107
+ Some ( token)
108
+ } else {
109
+ None
110
+ } ;
111
+ Ok ( Self ( move_token, input. parse ( ) ?) )
112
+ }
86
113
}
0 commit comments