fix: #[trace] compilation error when using trait objects (#130)

Co-authored-by: Andy Lok <andylokandy@hotmail.com>

Author: ArniDagur

Cargo.lock

@@ -19,6 +19,8 @@ use syn::punctuated::Punctuated;
 use syn::spanned::Spanned;
 use syn::*;
 
+use crate::visit_mut::VisitMut;
+
 /// An attribute macro designed to eliminate boilerplate code.
 ///
 /// This macro automatically creates a span for the annotated function. The span name defaults to
@@ -124,14 +126,13 @@ pub fn trace(
 
     let args = parse_macro_input!(args as Args);
     let input = parse_macro_input!(item as ItemFn);
-
     let func_name = input.sig.ident.to_string();
-    // check for async_trait-like patterns in the block, and instrument
-    // the future instead of the wrapper
+
+    // Check for async_trait-like patterns in the block, and instrument
+    // the future instead of the wrapper.
     let func_body = if let Some(internal_fun) =
         get_async_trait_info(&input.block, input.sig.asyncness.is_some())
     {
-        // let's rewrite some statements!
         match internal_fun.kind {
             // async-trait <= 0.1.43
             AsyncTraitKind::Function => {
@@ -141,23 +142,26 @@ pub fn trace(
             }
             // async-trait >= 0.1.44
             AsyncTraitKind::Async(async_expr) => {
-                // fallback if we couldn't find the '__async_trait' binding, might be
-                // useful for crates exhibiting the same behaviors as async-trait
                 let instrumented_block =
-                    gen_block(&func_name, &async_expr.block, true, false, &args);
+                    gen_block(&func_name, &async_expr.block, true, false, &args, None);
                 let async_attrs = &async_expr.attrs;
                 quote::quote! {
                     Box::pin(#(#async_attrs) * #instrumented_block)
                 }
             }
         }
     } else {
+        let output_ty = match input.sig.output {
+            ReturnType::Type(_, ref ty) => (**ty).clone(),
+            ReturnType::Default => parse_quote! { () },
+        };
         gen_block(
             &func_name,
             &input.block,
             input.sig.asyncness.is_some(),
             input.sig.asyncness.is_some(),
             &args,
+            Some(output_ty),
         )
     };
 
@@ -371,10 +375,14 @@ fn gen_block(
     async_context: bool,
     async_keyword: bool,
     args: &Args,
+    output_ty: Option<Type>,
 ) -> proc_macro2::TokenStream {
     let name = gen_name(block.span(), func_name, args);
     let properties = gen_properties(block.span(), args);
     let crate_path = &args.crate_path;
+    let output_ty_hint = output_ty
+        .map(erase_impl_trait)
+        .unwrap_or_else(|| parse_quote! { _ });
 
     // Generate the instrumented function body.
     // If the function is an `async fn`, this will wrap it in an async block.
@@ -392,7 +400,10 @@ fn gen_block(
                 {
                     let __span__ = #crate_path::Span::enter_with_local_parent( #name ) #properties;
                     #crate_path::future::FutureExt::in_span(
-                        async move { #block },
+                        async move {
+                            let __ret__: #output_ty_hint = #block;
+                            __ret__
+                        },
                         __span__,
                     )
                 }
@@ -547,3 +558,25 @@ fn path_to_string(path: &Path) -> String {
     }
     res
 }
+
+/// Replaces any `impl Trait` with `_` so it can be used as the type in
+/// a `let` statement's LHS.
+struct ImplTraitEraser;
+
+impl VisitMut for ImplTraitEraser {
+    fn visit_type_mut(&mut self, t: &mut Type) {
+        if let Type::ImplTrait(..) = t {
+            *t = syn::TypeInfer {
+                underscore_token: Token![_](t.span()),
+            }
+            .into();
+        } else {
+            syn::visit_mut::visit_type_mut(self, t);
+        }
+    }
+}
+
+fn erase_impl_trait(mut ty: Type) -> Type {
+    ImplTraitEraser.visit_type_mut(&mut ty);
+    ty
+}

fastrace-macro/src/lib.rs

@@ -17,6 +17,8 @@ log = { workspace = true }
 logcall = { version = "0.1" }
 tokio = { workspace = true }
 trybuild = { version = "1.0" }
+futures = { version = "0.3" }
+async-stream = { version = "0.3" }
 
 # The procedural macro `trace` only supports async-trait higher than or equal to 0.1.52
 async-trait = { version = "0.1.52" }

tests/macros/Cargo.toml

@@ -0,0 +1,37 @@
+use async_trait::async_trait;
+use std::future::Future;
+use std::pin::Pin;
+
+use fastrace::trace;
+
+#[derive(Debug)]
+pub struct InnerError;
+
+#[derive(Debug)]
+pub struct OuterError(InnerError);
+
+pub type MyFuture = Pin<Box<dyn Future<Output = Result<u32, OuterError>> + Send>>;
+
+#[async_trait]
+pub trait MyTrait {
+    async fn f() -> Result<MyFuture, OuterError>;
+}
+
+pub struct MyStruct;
+
+#[async_trait]
+impl MyTrait for MyStruct {
+    #[trace]
+    async fn f() -> Result<MyFuture, OuterError> {
+        let inner = async { Err(InnerError) };
+
+        let mapped = async move { inner.await.map_err(OuterError) };
+
+        Ok(Box::pin(mapped))
+    }
+}
+
+#[tokio::main]
+async fn main() {
+    let _ = MyStruct::f().await;
+}

tests/macros/tests/ui/ok/async-trait-dyn-future.rs

@@ -0,0 +1,26 @@
+use std::future::Future;
+use std::pin::Pin;
+
+use fastrace::trace;
+
+#[derive(Debug)]
+pub struct InnerError;
+
+#[derive(Debug)]
+pub struct OuterError(InnerError);
+
+pub type MyFuture = Pin<Box<dyn Future<Output = Result<u32, OuterError>> + Send>>;
+
+#[trace]
+pub async fn f() -> Result<MyFuture, OuterError> {
+    let inner = async { Err(InnerError) };
+
+    let mapped = async move { inner.await.map_err(OuterError) };
+
+    Ok(Box::pin(mapped))
+}
+
+#[tokio::main]
+async fn main() {
+    let _ = f().await;
+}

tests/macros/tests/ui/ok/dyn-future.rs

@@ -0,0 +1,16 @@
+use fastrace::trace;
+use futures::Stream;
+
+#[trace]
+async fn stream() -> impl Stream<Item = i64> {
+    async_stream::stream! {
+        for i in 0..100 {
+            yield i;
+        }
+    }
+}
+
+#[tokio::main]
+async fn main() {
+    let _ = stream().await;
+}