Fix borrowing issues

src/lib.rs

@@ -37,7 +37,10 @@
 extern crate proc_macro;
 
 use proc_macro::TokenStream;
-use quote::quote;
+use proc_macro2::{Ident, Span, TokenStream as TokenStream2};
+use quote::{quote, ToTokens};
+use syn::parse::{self, Parse, ParseStream};
+use syn::Token;
 
 /// Add context to errors from a function.
 ///
@@ -55,32 +58,67 @@ use quote::quote;
 ///     })().map_err(|err| err.context("context"))
 /// }
 /// ```
+///
+/// Sometimes you will receive borrowck errors, especially when returning references. These can
+/// often be fixed by setting the `move` option of the attribute macro. For example:
+///
+/// ```
+/// use fn_error_context::context;
+///
+/// #[context(move, "context")]
+/// fn returns_reference(val: &mut u32) -> anyhow::Result<&mut u32> {
+///     Ok(&mut *val)
+/// }
+/// ```
 #[proc_macro_attribute]
 pub fn context(args: TokenStream, input: TokenStream) -> TokenStream {
-    let args: proc_macro2::TokenStream = args.into();
+    let Args(move_token, format_args) = syn::parse_macro_input!(args);
     let mut input = syn::parse_macro_input!(input as syn::ItemFn);
 
     let body = &input.block;
     let return_ty = &input.sig.output;
-    if input.sig.asyncness.is_some() {
-        match return_ty {
+    let err = Ident::new("err", Span::mixed_site());
+    let new_body = if input.sig.asyncness.is_some() {
+        let return_ty = match return_ty {
             syn::ReturnType::Default => {
-                return syn::Error::new_spanned(return_ty, "function should return Result")
+                return syn::Error::new_spanned(input, "function should return Result")
                     .to_compile_error()
-                    .into()
-            }
-            syn::ReturnType::Type(_, return_ty) => {
-                input.block.stmts = syn::parse_quote!(
-                    let result: #return_ty = async { #body }.await;
-                    result.map_err(|err| err.context(format!(#args)).into())
-                );
+                    .into();
             }
+            syn::ReturnType::Type(_, return_ty) => return_ty,
+        };
+        let result = Ident::new("result", Span::mixed_site());
+        quote! {
+            let #result: #return_ty = async #move_token { #body }.await;
+            #result.map_err(|#err| #err.context(format!(#format_args)).into())
         }
     } else {
-        input.block.stmts = syn::parse_quote!(
-            (|| #return_ty #body)().map_err(|err| err.context(format!(#args)).into())
-        );
-    }
+        let force_fn_once = Ident::new("force_fn_once", Span::mixed_site());
+        quote! {
+            // Moving a non-`Copy` value into the closure tells borrowck to always treat the closure
+            // as a `FnOnce`, preventing some borrowing errors.
+            let #force_fn_once = ::core::iter::empty::<()>();
+            (#move_token || #return_ty {
+                ::core::mem::drop(#force_fn_once);
+                #body
+            })().map_err(|#err| #err.context(format!(#format_args)).into())
+        }
+    };
+    input.block.stmts = vec![syn::Stmt::Expr(syn::Expr::Verbatim(new_body))];
 
-    quote!(#input).into()
+    input.into_token_stream().into()
+}
+
+struct Args(Option<Token![move]>, TokenStream2);
+impl Parse for Args {
+    fn parse(input: ParseStream<'_>) -> parse::Result<Self> {
+        let move_token = if input.peek(Token![move]) {
+            let token = input.parse()?;
+            input.parse::<Token![,]>()?;
+            Some(token)
+        } else {
+            None
+        };
+        Ok(Self(move_token, input.parse()?))
+    }
 }

tests/async_borrowing.rs

@@ -0,0 +1,8 @@
+use fn_error_context::context;
+
+#[context("context")]
+async fn borrows(x: &mut u32) -> anyhow::Result<&mut u32> {
+    Ok(x)
+}
+
+fn main() {}

tests/async_move.rs

@@ -0,0 +1,8 @@
+use fn_error_context::context;
+
+#[context(move, "context")]
+async fn borrows(val: &mut u32) -> anyhow::Result<&u32> {
+    Ok(&*val)
+}
+
+fn main() {}

tests/async_no_move.rs

@@ -0,0 +1,9 @@
+use fn_error_context::context;
+
+#[context("{}", context.as_ref())]
+async fn no_move(context: impl AsRef<str>) -> anyhow::Result<()> {
+    context.as_ref();
+    Ok(())
+}
+
+fn main() {}

tests/async_without_return.stderr

@@ -1,7 +1,6 @@
 error: function should return Result
- --> $DIR/async_without_return.rs:3:1
+ --> $DIR/async_without_return.rs:4:1
   |
-3 | #[context("context")]
-  | ^^^^^^^^^^^^^^^^^^^^^
-  |
-  = note: this error originates in an attribute macro (in Nightly builds, run with -Z macro-backtrace for more info)
+4 | / async fn async_something() {
+5 | | }
+  | |_^

tests/move.rs

@@ -0,0 +1,8 @@
+use fn_error_context::context;
+
+#[context(move, "context")]
+fn foo(x: &mut u32) -> anyhow::Result<&u32> {
+    Ok(&*x)
+}
+
+fn main() {}

tests/no_move.rs

@@ -0,0 +1,9 @@
+use fn_error_context::context;
+
+#[context("{}", context.as_ref())]
+fn no_move(context: impl AsRef<str>) -> anyhow::Result<()> {
+    context.as_ref();
+    Ok(())
+}
+
+fn main() {}

tests/tests.rs

@@ -17,4 +17,9 @@ fn tests() {
     tests.pass("tests/fmt_named_arg.rs");
     tests.compile_fail("tests/async_without_return.rs");
     tests.compile_fail("tests/preserve_lint.rs");
+    tests.pass("tests/async_borrowing.rs");
+    tests.pass("tests/no_move.rs");
+    tests.pass("tests/async_no_move.rs");
+    tests.pass("tests/move.rs");
+    tests.pass("tests/async_move.rs");
 }