diff --git a/tokio-macros/Cargo.toml b/tokio-macros/Cargo.toml index ea9839c6d06..b40dc184050 100644 --- 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 index 544130473af..97dc9997973 100644 --- 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; @@ -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> = body; } } else { quote! { - let body = async #body; + let body = async { #body }; } };