WIP

tokio-rs · Mar 28, 2024 · d13ab28 · d13ab28
1 parent 7e7f42a
commit d13ab28
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 14 deletions.
diff --git a/tokio-macros/Cargo.toml b/tokio-macros/Cargo.toml
@@ -24,7 +24,7 @@ proc-macro = true
 [dependencies]
 proc-macro2 = "1.0.60"
 quote = "1"
-syn = { version = "2.0", features = ["full"] }
+syn = { version = "2.0", features = ["full", "extra-traits"] }
 
 [dev-dependencies]
 tokio = { version = "1.0.0", path = "../tokio", features = ["full"] }

diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs
@@ -1,7 +1,7 @@
 use proc_macro2::{Span, TokenStream, TokenTree};
 use quote::{quote, quote_spanned, ToTokens};
 use syn::parse::{Parse, ParseStream, Parser};
-use syn::{braced, Attribute, Ident, Path, Signature, Visibility};
+use syn::{braced, parse_quote, Attribute, Ident, Path, Signature, Visibility};
 
 // syn::AttributeArgs does not implement syn::Parse
 type AttributeArgs = syn::punctuated::Punctuated<syn::Meta, syn::Token![,]>;
@@ -380,24 +380,39 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
         }
     };
 
+    let unit_type = Box::new(parse_quote! { () });
+    let never_type = Box::new(parse_quote! { ! });
+
     let output_type = match &input.sig.output {
         // For functions with no return value syn doesn't print anything,
         // but that doesn't work as `Output` for our boxed `Future`, so
         // default to `()` (the same type as the function output).
-        syn::ReturnType::Default => quote! { () },
-        syn::ReturnType::Type(_, ret_type) => quote! { #ret_type },
+        syn::ReturnType::Default => &unit_type,
+        // match on never type
+        syn::ReturnType::Type(_, ret_type) if ret_type == &never_type => &unit_type,
+        syn::ReturnType::Type(_, ret_type) => ret_type,
     };
 
-    input.stmts.last_mut().map(|stmt| {
-        *stmt = quote! {
-            let _lst_stmt: #output_type = {
-                #stmt
-            };
-            _lst_stmt
-        };
-    });
+    // input.stmts.last_mut().map(|stmt| {
+    //     if output_type != &unit_type {
+    //         *stmt = quote! {
+    //             let _lst_stmt: #output_type = {
+    //                 #stmt
+    //             };
+    //             _lst_stmt
+    //         };
+    //     }
+    // });
 
     let body = input.body();
+    let body = if output_type == &unit_type {
+        quote! { #body }
+    } else {
+        quote! {
+            let main_body: #output_type = #body;
+            main_body
+        }
+    };
 
     // For test functions pin the body to the stack and use `Pin<&mut dyn
     // Future>` to reduce the amount of `Runtime::block_on` (and related
@@ -410,13 +425,13 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
     // there will be no benefit.
     let body = if is_test {
         quote! {
-            let body = async #body;
+            let body = async { #body };
             #crate_path::pin!(body);
             let body: ::core::pin::Pin<&mut dyn ::core::future::Future<Output = #output_type>> = body;
         }
     } else {
         quote! {
-            let body = async #body;
+            let body = async { #body };
         }
     };