3737extern crate proc_macro;
3838
3939use 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 ;
4144
4245/// Add context to errors from a function.
4346///
@@ -55,32 +58,56 @@ use quote::quote;
5558/// })().map_err(|err| err.context("context"))
5659/// }
5760/// ```
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+ /// ```
5871#[ proc_macro_attribute]
5972pub 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 ) ;
6174 let mut input = syn:: parse_macro_input!( input as syn:: ItemFn ) ;
6275
6376 let body = & input. block ;
6477 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 {
6780 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" )
6982 . 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 ( ) ;
7784 }
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( ) )
7890 }
7991 } 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) ) ] ;
8497
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+ }
86113}
0 commit comments