Repository: emit-rs/emit (rust)

Issue:
emit::span doesn't work on methods returning impl Trait
Using `emit::span` on a method returning `impl Trait` produces:

```
error[E0562]: `impl Trait` is not allowed in the type of variable bindings
```

We shouldn't try infer the type of the function block unless we're going to match it as a result with the `ok_lvl` or `err_lvl` attributes.

Pull request description:
Fix up impl trait as inferred type
Closes #109

Patch:
diff --git a/macros/src/span.rs b/macros/src/span.rs
index 970151f..718066a 100644
--- a/macros/src/span.rs
+++ b/macros/src/span.rs
@@ -211,25 +211,35 @@ fn inject_sync(
     let template_tokens = template.template_tokens().to_ref_tokens();
     let template_literal_tokens = template.template_literal_tokens();
 
-    // Capture control flow statements within the body
-    let body_tokens = quote!((move || {
-        let __r: #body_ret_ty_tokens = #body_tokens;
-        __r
-    })());
-
     let Completion {
         body_tokens,
         span_evt_props_tokens,
         default_completion_tokens,
-    } = completion(
-        body_tokens,
-        default_lvl_tokens,
-        ok_lvl_tokens,
-        err_lvl_tokens,
-        span_guard,
-        rt_tokens,
-        &template_tokens,
-    )?;
+    } = if ok_lvl_tokens.is_some() || err_lvl_tokens.is_some() {
+        // Wrap the body in a closure so we can rely on code running afterwards
+        // without control flow statements like `return` getting in the way
+        //
+        // We can't use a drop guard here because we need to match on the result
+        //
+        // We also need to specify the return type, otherwise inference seems to fail.
+        // We might be able to avoid this in the future
+        let body_tokens = quote!((move || {
+            let __r: #body_ret_ty_tokens = #body_tokens;
+            __r
+        })());
+
+        result_completion(
+            body_tokens,
+            default_lvl_tokens,
+            ok_lvl_tokens,
+            err_lvl_tokens,
+            span_guard,
+            rt_tokens,
+            &template_tokens,
+        )?
+    } else {
+        completion(body_tokens, default_lvl_tokens, rt_tokens, &template_tokens)?
+    };
 
     Ok(quote!({
         let (mut __ctxt, __span_guard) = emit::__private::__private_begin_span(
@@ -267,24 +277,31 @@ fn inject_async(
     let template_tokens = template.template_tokens().to_ref_tokens();
     let template_literal_tokens = template.template_literal_tokens();
 
-    let body_tokens = quote!(async move {
-        let __r: #body_ret_ty_tokens = #body_tokens;
-        __r
-    }.await);
-
     let Completion {
         body_tokens,
         span_evt_props_tokens,
         default_completion_tokens,
-    } = completion(
-        body_tokens,
-        default_lvl_tokens,
-        ok_lvl_tokens,
-        err_lvl_tokens,
-        span_guard,
-        rt_tokens,
-        &template_tokens,
-    )?;
+    } = if ok_lvl_tokens.is_some() || err_lvl_tokens.is_some() {
+        // Like the sync case, ensure control flow doesn't interrupt
+        // our matching of the result, and provide a concrete type
+        // for inference
+        let body_tokens = quote!(async move {
+            let __r: #body_ret_ty_tokens = #body_tokens;
+            __r
+        }.await);
+
+        result_completion(
+            body_tokens,
+            default_lvl_tokens,
+            ok_lvl_tokens,
+            err_lvl_tokens,
+            span_guard,
+            rt_tokens,
+            &template_tokens,
+        )?
+    } else {
+        completion(body_tokens, default_lvl_tokens, rt_tokens, &template_tokens)?
+    };
 
     Ok(quote!({
         let (__ctxt, __span_guard) = emit::__private::__private_begin_span(
@@ -319,7 +336,7 @@ struct Completion {
     default_completion_tokens: TokenStream,
 }
 
-fn completion(
+fn result_completion(
     body_tokens: TokenStream,
     default_lvl_tokens: Option<TokenStream>,
     ok_lvl_tokens: Option<TokenStream>,
@@ -328,69 +345,74 @@ fn completion(
     rt_tokens: &TokenStream,
     template_tokens: &TokenStream,
 ) -> Result<Completion, syn::Error> {
-    let body_tokens = if ok_lvl_tokens.is_some() || err_lvl_tokens.is_some() {
-        // If the span is applied to a Result-returning function then wrap the body
-        // We'll attach the error to the span if the call fails and set the appropriate level
-
-        let ok_branch = {
-            let mut evt_props = Props::new();
-            push_evt_props(
-                &mut evt_props,
-                ok_lvl_tokens.or_else(|| default_lvl_tokens.clone()),
-            )?;
-            let span_evt_props_tokens = evt_props.props_tokens();
-
-            quote!(
-                Ok(ok) => {
-                    #span_guard.complete_with(|span| {
-                        emit::__private::__private_complete_span(
-                            #rt_tokens,
-                            span,
-                            #template_tokens,
-                            #span_evt_props_tokens,
-                        )
-                    });
-
-                    Ok(ok)
-                }
-            )
-        };
-
-        let err_branch = {
-            let err_ident = Ident::new(emit_core::well_known::KEY_ERR, Span::call_site());
-
-            let mut evt_props = Props::new();
-            push_evt_props(
-                &mut evt_props,
-                err_lvl_tokens.or_else(|| default_lvl_tokens.clone()),
-            )?;
-            evt_props.push(&syn::parse2::<FieldValue>(quote!(#err_ident))?, false, true)?;
-            let span_evt_props_tokens = evt_props.props_tokens();
-
-            quote!(
-                Err(#err_ident) => {
-                    #span_guard.complete_with(|span| {
-                        emit::__private::__private_complete_span(
-                            #rt_tokens,
-                            span,
-                            #template_tokens,
-                            #span_evt_props_tokens,
-                        )
-                    });
-
-                    Err(#err_ident)
-                }
-            )
-        };
-
-        quote!(match #body_tokens {
-            #ok_branch,
-            #err_branch,
-        })
-    } else {
-        body_tokens
+    // If the span is applied to a Result-returning function then wrap the body
+    // We'll attach the error to the span if the call fails and set the appropriate level
+
+    let ok_branch = {
+        let mut evt_props = Props::new();
+        push_evt_props(
+            &mut evt_props,
+            ok_lvl_tokens.or_else(|| default_lvl_tokens.clone()),
+        )?;
+        let span_evt_props_tokens = evt_props.props_tokens();
+
+        quote!(
+            Ok(ok) => {
+                #span_guard.complete_with(|span| {
+                    emit::__private::__private_complete_span(
+                        #rt_tokens,
+                        span,
+                        #template_tokens,
+                        #span_evt_props_tokens,
+                    )
+                });
+
+                Ok(ok)
+            }
+        )
+    };
+
+    let err_branch = {
+        let err_ident = Ident::new(emit_core::well_known::KEY_ERR, Span::call_site());
+
+        let mut evt_props = Props::new();
+        push_evt_props(
+            &mut evt_props,
+            err_lvl_tokens.or_else(|| default_lvl_tokens.clone()),
+        )?;
+        evt_props.push(&syn::parse2::<FieldValue>(quote!(#err_ident))?, false, true)?;
+        let span_evt_props_tokens = evt_props.props_tokens();
+
+        quote!(
+            Err(#err_ident) => {
+                #span_guard.complete_with(|span| {
+                    emit::__private::__private_complete_span(
+                        #rt_tokens,
+                        span,
+                        #template_tokens,
+                        #span_evt_props_tokens,
+                    )
+                });
+
+                Err(#err_ident)
+            }
+        )
     };
 
+    let body_tokens = quote!(match #body_tokens {
+        #ok_branch,
+        #err_branch,
+    });
+
+    completion(body_tokens, default_lvl_tokens, rt_tokens, template_tokens)
+}
+
+fn completion(
+    body_tokens: TokenStream,
+    default_lvl_tokens: Option<TokenStream>,
+    rt_tokens: &TokenStream,
+    template_tokens: &TokenStream,
+) -> Result<Completion, syn::Error> {
     let mut evt_props = Props::new();
     push_evt_props(&mut evt_props, default_lvl_tokens)?;
     let span_evt_props_tokens = evt_props.props_tokens();


Test patch:
diff --git a/test/ui/Cargo.toml b/test/ui/Cargo.toml
index 839ae1e..e8461ca 100644
--- a/test/ui/Cargo.toml
+++ b/test/ui/Cargo.toml
@@ -42,5 +42,8 @@ version = "1"
 version = "1"
 features = ["full"]
 
+[dependencies.rustversion]
+version = "1"
+
 [lints.rust]
 unexpected_cfgs = { level = "allow", check-cfg = ['cfg(emit_disabled)'] }
diff --git a/test/ui/src/compile_fail/std/span_ok_lvl_impl_trait_return.rs b/test/ui/src/compile_fail/std/span_ok_lvl_impl_trait_return.rs
new file mode 100644
index 0000000..9f84341
--- /dev/null
+++ b/test/ui/src/compile_fail/std/span_ok_lvl_impl_trait_return.rs
@@ -0,0 +1,8 @@
+fn main() {
+
+}
+
+#[emit::span(ok_lvl: "info", err_lvl: "warn", "try_work")]
+fn try_work() -> Result<(), impl std::error::Error> {
+    Ok::<(), std::io::Error>(())
+}
diff --git a/test/ui/src/compile_fail/std/span_ok_lvl_impl_trait_return.stderr b/test/ui/src/compile_fail/std/span_ok_lvl_impl_trait_return.stderr
new file mode 100644
index 0000000..3479f82
--- /dev/null
+++ b/test/ui/src/compile_fail/std/span_ok_lvl_impl_trait_return.stderr
@@ -0,0 +1,7 @@
+error[E0562]: `impl Trait` is not allowed in the type of variable bindings
+ --> src/compile_fail/std/span_ok_lvl_impl_trait_return.rs:6:29
+  |
+6 | fn try_work() -> Result<(), impl std::error::Error> {
+  |                             ^^^^^^^^^^^^^^^^^^^^^^
+  |
+  = note: `impl Trait` is only allowed in arguments and return types of functions and methods
diff --git a/test/ui/src/lib.rs b/test/ui/src/lib.rs
index 7128ec1..f0f88ba 100644
--- a/test/ui/src/lib.rs
+++ b/test/ui/src/lib.rs
@@ -27,6 +27,7 @@ mod format;
 
 #[test]
 #[cfg(feature = "compile_fail")]
+#[rustversion::nightly]
 fn compile_fail_std() {
     let t = trybuild::TestCases::new();
     t.compile_fail("src/compile_fail/std/*.rs");
diff --git a/test/ui/src/span.rs b/test/ui/src/span.rs
index e619228..c0f09d3 100644
--- a/test/ui/src/span.rs
+++ b/test/ui/src/span.rs
@@ -750,3 +750,51 @@ fn span_props_precedence() {
         let _ = exec();
     });
 }
+
+#[test]
+fn span_impl_trait_return() {
+    static CALLED: StaticCalled = StaticCalled::new();
+    static RT: StaticRuntime = static_runtime(
+        |_| {
+            CALLED.record();
+        },
+        |_| true,
+    );
+
+    #[emit::span(rt: RT, "greet {user}")]
+    fn exec(user: &str) -> impl std::fmt::Display {
+        let _ = user;
+
+        "done"
+    }
+
+    let _ = exec("Rust");
+
+    RT.emitter().blocking_flush(Duration::from_secs(1));
+
+    assert!(CALLED.was_called());
+}
+
+#[tokio::test]
+async fn span_impl_trait_return_async() {
+    static CALLED: StaticCalled = StaticCalled::new();
+    static RT: StaticRuntime = static_runtime(
+        |_| {
+            CALLED.record();
+        },
+        |_| true,
+    );
+
+    #[emit::span(rt: RT, "greet {user}")]
+    async fn exec(user: &str) -> impl std::fmt::Display {
+        let _ = user;
+
+        "done"
+    }
+
+    let _ = exec("Rust").await;
+
+    RT.emitter().blocking_flush(Duration::from_secs(1));
+
+    assert!(CALLED.was_called());
+}