macros: improve type diagnostics

tests-build/tests/fail/macros_invalid_input.stderr

@@ -123,3 +123,17 @@ note: the lint level is defined here
    |
 1  | #![deny(duplicate_macro_attributes)]
    |         ^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error[E0277]: `()` is not a future
+ --> tests/fail/macros_invalid_input.rs:5:1
+  |
+5 | #[tokio::main]
+  | ^^^^^^^^^^^^^^ `()` is not a future
+  |
+  = help: the trait `Future` is not implemented for `()`
+note: required by a bound in `tests_build::tokio::runtime::Runtime::block_on`
+ --> $WORKSPACE/tokio/src/runtime/runtime.rs
+  |
+  |     pub fn block_on<F: Future>(&self, future: F) -> F::Output {
+  |                        ^^^^^^ required by this bound in `Runtime::block_on`
+  = note: this error originates in the attribute macro `tokio::main` (in Nightly builds, run with -Z macro-backtrace for more info)

tests-build/tests/fail/macros_type_mismatch.rs

@@ -15,11 +15,6 @@ async fn extra_semicolon() -> Result<(), ()> {
     Ok(());
 }
 
-#[tokio::main]
-async fn invalid_return_type() -> Option<()> {
-    ()
-}
-
 // https://github.com/tokio-rs/tokio/issues/4635
 #[allow(redundant_semicolons)]
 #[rustfmt::skip]

tests-build/tests/fail/macros_type_mismatch.stderr

@@ -6,68 +6,33 @@ error[E0308]: mismatched types
   |
   = note: expected unit type `()`
                   found enum `Result<(), _>`
-help: a return type might be missing here
-  |
-4 | async fn missing_semicolon_or_return_type() -> _ {
-  |                                             ++++
-help: consider using `Result::expect` to unwrap the `Result<(), _>` value, panicking if the value is a `Result::Err`
-  |
-5 |     Ok(()).expect("REASON")
-  |           +++++++++++++++++
 
 error[E0308]: mismatched types
-  --> tests/fail/macros_type_mismatch.rs:10:5
+  --> tests/fail/macros_type_mismatch.rs:10:12
    |
 10 |     return Ok(());
-   |     ^^^^^^^^^^^^^ expected `()`, found `Result<(), _>`
+   |            ^^^^^^ expected `()`, found `Result<(), _>`
    |
    = note: expected unit type `()`
                    found enum `Result<(), _>`
-help: a return type might be missing here
-   |
-9  | async fn missing_return_type() -> _ {
-   |                                ++++
-help: consider using `Result::expect` to unwrap the `Result<(), _>` value, panicking if the value is a `Result::Err`
-   |
-10 |     return Ok(()).expect("REASON");
-   |                  +++++++++++++++++
 
 error[E0308]: mismatched types
-  --> tests/fail/macros_type_mismatch.rs:15:5
+  --> tests/fail/macros_type_mismatch.rs:14:46
    |
-14 | async fn extra_semicolon() -> Result<(), ()> {
-   |                               -------------- expected `Result<(), ()>` because of return type
-15 |     Ok(());
-   |     ^^^^^^ expected `Result<(), ()>`, found `()`
+14 |   async fn extra_semicolon() -> Result<(), ()> {
+   |  ______________________________________________^
+15 | |     Ok(());
+   | |           - help: remove this semicolon to return this value
+16 | | }
+   | |_^ expected `Result<(), ()>`, found `()`
    |
    = note:   expected enum `Result<(), ()>`
            found unit type `()`
-help: try wrapping the expression in a variant of `Result`
-   |
-15 |     Ok(Ok(()));
-   |     +++      +
-15 |     Err(Ok(()));
-   |     ++++      +
-
-error[E0308]: mismatched types
-  --> tests/fail/macros_type_mismatch.rs:20:5
-   |
-19 | async fn invalid_return_type() -> Option<()> {
-   |                                   ---------- expected `Option<()>` because of return type
-20 |     ()
-   |     ^^ expected `Option<()>`, found `()`
-   |
-   = note:   expected enum `Option<()>`
-           found unit type `()`
-help: try wrapping the expression in `Some`
-   |
-20 |     Some(())
-   |     +++++  +
 
 error[E0308]: mismatched types
-  --> tests/fail/macros_type_mismatch.rs:28:5
+  --> tests/fail/macros_type_mismatch.rs:23:12
    |
-27 | async fn issue_4635() {
+22 | async fn issue_4635() {
    |                      - help: try adding a return type: `-> i32`
-28 |     return 1;
-   |     ^^^^^^^^ expected `()`, found integer
+23 |     return 1;
+   |            ^ expected `()`, found integer

tokio-macros/src/entry.rs

@@ -390,8 +390,83 @@ fn build_config(
     config.build()
 }
 
+fn fn_without_args(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenStream {
+    let async_keyword = input.sig.asyncness.take();
+    let fn_sig = &input.sig;
+    let fn_name = &input.sig.ident;
+
+    let crate_path = config
+        .crate_name
+        .map(ToTokens::into_token_stream)
+        .unwrap_or_else(|| Ident::new("tokio", Span::call_site()).into_token_stream());
+
+    let mut rt = match config.flavor {
+        RuntimeFlavor::CurrentThread => quote! {
+            #crate_path::runtime::Builder::new_current_thread()
+        },
+        RuntimeFlavor::Threaded => quote! {
+            #crate_path::runtime::Builder::new_multi_thread()
+        },
+    };
+    if let Some(v) = config.worker_threads {
+        rt = quote! { #rt.worker_threads(#v) };
+    }
+    if let Some(v) = config.start_paused {
+        rt = quote! { #rt.start_paused(#v) };
+    }
+    if let Some(v) = config.unhandled_panic {
+        let unhandled_panic = v.into_tokens(&crate_path);
+        rt = quote! { #rt.unhandled_panic(#unhandled_panic) };
+    }
+
+    let generated_attrs = if is_test {
+        quote! { #[::core::prelude::v1::test] }
+    } else {
+        quote! {}
+    };
+
+    // This explicit `return` is intentional. See tokio-rs/tokio#4636
+    let last_block = quote! {
+        #[allow(clippy::expect_used, clippy::diverging_sub_expression, clippy::needless_return)]
+        {
+            return #rt
+                .enable_all()
+                .build()
+                .expect("Failed building the Runtime")
+                .block_on(body);
+        }
+    };
+
+    let body = input.body;
+
+    input.body = if is_test {
+        let output_type = match &input.sig.output {
+            syn::ReturnType::Default => quote! { () },
+            syn::ReturnType::Type(_, ret_type) => quote! { #ret_type },
+        };
+        quote! {
+            #async_keyword #fn_sig #body
+            let body = #fn_name();
+
+            #crate_path::pin!(body);
+            let body: ::core::pin::Pin<&mut dyn ::core::future::Future<Output = #output_type>> = body;
+        }
+    } else {
+        quote! {
+            #async_keyword #fn_sig #body
+            let body = #fn_name();
+        }
+    };
+
+    input.into_tokens(generated_attrs, last_block)
+}
+
 fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenStream {
-    input.sig.asyncness = None;
+    if input.sig.inputs.is_empty() {
+        return fn_without_args(input, is_test, config);
+    }
+
+    let async_keyword = input.sig.asyncness.take();
 
     // If type mismatch occurs, the current rustc points to the last statement.
     let (last_stmt_start_span, last_stmt_end_span) = {
@@ -415,10 +490,7 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
 
                         for tt in stream.by_ref() {
                             match tt {
-                                TokenTree::Punct(p) if p.as_char() == ';' => {
-                                    // end = p.span();
-                                    break;
-                                }
+                                TokenTree::Punct(p) if p.as_char() == ';' => break,
                                 tt => end = tt.span(),
                             }
                         }
@@ -465,7 +537,6 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
         quote! {}
     };
 
-    let body_ident = quote! { body };
     // This explicit `return` is intentional. See tokio-rs/tokio#4636
     let last_block = quote_spanned! {last_stmt_end_span=>
         #[allow(clippy::expect_used, clippy::diverging_sub_expression, clippy::needless_return)]
@@ -474,7 +545,7 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
                 .enable_all()
                 .build()
                 .expect("Failed building the Runtime")
-                .block_on(#body_ident);
+                .block_on(body);
         }
     };
 
@@ -498,13 +569,13 @@ fn parse_knobs(mut input: ItemFn, is_test: bool, config: FinalConfig) -> TokenSt
             syn::ReturnType::Type(_, ret_type) => quote! { #ret_type },
         };
         quote! {
-            let body = async #body;
+            let body = #async_keyword #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_keyword #body;
         }
     };
     input.into_tokens(generated_attrs, last_block)