From d9fd1bc7a3c0ab38245c54f88b0e3a3ad949a30b Mon Sep 17 00:00:00 2001
From: Jake Goulding <jake.goulding@gmail.com>
Date: Mon, 11 Mar 2024 13:33:50 -0400
Subject: [PATCH] Improve how the `report` macro handles return type inference

---
 compatibility-tests/futures/src/report.rs | 25 +++++++++++++++++++++++
 snafu-derive/src/report.rs                | 16 +++++++--------
 tests/report.rs                           | 24 ++++++++++++++++++++++
 3 files changed, 57 insertions(+), 8 deletions(-)

diff --git a/compatibility-tests/futures/src/report.rs b/compatibility-tests/futures/src/report.rs
index ba1442f..fb33d4b 100644
--- a/compatibility-tests/futures/src/report.rs
+++ b/compatibility-tests/futures/src/report.rs
@@ -56,3 +56,28 @@ fn async_std_main_attribute_last() {
 async fn async_std_test_attribute_last() -> Result<(), Error> {
     Ok(())
 }
+
+#[test]
+fn works_with_tough_inference() {
+    #[derive(Debug, Snafu)]
+    struct InnerError;
+
+    #[derive(Debug, Snafu)]
+    struct OuterError {
+        source: InnerError,
+    }
+
+    fn inner() -> Result<(), InnerError> {
+        InnerSnafu.fail()
+    }
+
+    #[snafu::report]
+    #[tokio::main(flavor = "current_thread")]
+    async fn mainlike() -> Result<(), OuterError> {
+        loop {
+            inner().context(OuterSnafu)?;
+        }
+    }
+
+    let _: Report<_> = mainlike();
+}
diff --git a/snafu-derive/src/report.rs b/snafu-derive/src/report.rs
index 4c92621..08358e6 100644
--- a/snafu-derive/src/report.rs
+++ b/snafu-derive/src/report.rs
@@ -56,15 +56,15 @@ pub fn body(
         if cfg!(feature = "rust_1_61") {
             quote! {
                 {
-                    let __snafu_body = async #block;
-                    <::snafu::Report<_> as ::core::convert::From<_>>::from(__snafu_body.await)
+                    let __snafu_body: #output_ty = async #block.await;
+                    <::snafu::Report<_> as ::core::convert::From<_>>::from(__snafu_body)
                 }
             }
         } else {
             quote! {
                 {
-                    let __snafu_body = async #block;
-                    ::core::result::Result::map_err(__snafu_body.await, ::snafu::Report::from_error)
+                    let __snafu_body: #output_ty = async #block.await;
+                    ::core::result::Result::map_err(__snafu_body, ::snafu::Report::from_error)
                 }
             }
         }
@@ -72,15 +72,15 @@ pub fn body(
         if cfg!(feature = "rust_1_61") {
             quote! {
                 {
-                    let __snafu_body = || #block;
-                    <::snafu::Report<_> as ::core::convert::From<_>>::from(__snafu_body())
+                    let __snafu_body: #output_ty = (|| #block)();
+                    <::snafu::Report<_> as ::core::convert::From<_>>::from(__snafu_body)
                 }
             }
         } else {
             quote! {
                 {
-                    let __snafu_body = || #block;
-                    ::core::result::Result::map_err(__snafu_body(), ::snafu::Report::from_error)
+                    let __snafu_body: #output_ty = (|| #block)();
+                    ::core::result::Result::map_err(__snafu_body, ::snafu::Report::from_error)
                 }
             }
         }
diff --git a/tests/report.rs b/tests/report.rs
index cd5c61d..866c3b7 100644
--- a/tests/report.rs
+++ b/tests/report.rs
@@ -155,6 +155,30 @@ fn procedural_macro_works_with_result_return_type() {
     let _: Report<Error> = mainlike_result();
 }
 
+#[test]
+fn procedural_macro_works_with_tough_inference() {
+    #[derive(Debug, Snafu)]
+    struct InnerError;
+
+    #[derive(Debug, Snafu)]
+    struct OuterError {
+        source: InnerError,
+    }
+
+    fn inner() -> Result<(), InnerError> {
+        InnerSnafu.fail()
+    }
+
+    #[snafu::report]
+    fn mainlike_result() -> Result<(), OuterError> {
+        loop {
+            inner().context(OuterSnafu)?;
+        }
+    }
+
+    let _: Report<_> = mainlike_result();
+}
+
 #[test]
 fn termination_returns_failure_code() {
     use std::process::Termination;