From 2bf613206a29da5a5d262e0b44bb7fac1d40fa31 Mon Sep 17 00:00:00 2001
From: Taiki Endo <te316e89@gmail.com>
Date: Tue, 19 Oct 2021 12:49:57 +0900
Subject: [PATCH] macros: fix type resolution error in #[tokio::main] (#4176)

---
 tests-build/tests/fail/macros_type_mismatch.rs |  8 ++++++++
 .../tests/fail/macros_type_mismatch.stderr     | 15 +++++++++++----
 tokio-macros/src/entry.rs                      | 18 ++++++++++--------
 tokio/tests/macros_test.rs                     | 16 ++++++++++++++++
 4 files changed, 45 insertions(+), 12 deletions(-)

diff --git a/tests-build/tests/fail/macros_type_mismatch.rs b/tests-build/tests/fail/macros_type_mismatch.rs
index 47e34938ea0..0a5b9c4c727 100644
--- a/tests-build/tests/fail/macros_type_mismatch.rs
+++ b/tests-build/tests/fail/macros_type_mismatch.rs
@@ -12,6 +12,14 @@ async fn missing_return_type() {
 
 #[tokio::main]
 async fn extra_semicolon() -> Result<(), ()> {
+    /* TODO(taiki-e): help message still wrong
+    help: try using a variant of the expected enum
+       |
+    23 |     Ok(Ok(());)
+       |
+    23 |     Err(Ok(());)
+       |
+    */
     Ok(());
 }
 
diff --git a/tests-build/tests/fail/macros_type_mismatch.stderr b/tests-build/tests/fail/macros_type_mismatch.stderr
index 978dfc13ae6..4122ffb56ce 100644
--- a/tests-build/tests/fail/macros_type_mismatch.stderr
+++ b/tests-build/tests/fail/macros_type_mismatch.stderr
@@ -27,12 +27,19 @@ error[E0308]: mismatched types
                    found enum `Result<(), _>`
 
 error[E0308]: mismatched types
-  --> $DIR/macros_type_mismatch.rs:14:31
+  --> $DIR/macros_type_mismatch.rs:23:5
    |
 14 | async fn extra_semicolon() -> Result<(), ()> {
-   |          ---------------      ^^^^^^^^^^^^^^ expected enum `Result`, found `()`
-   |          |
-   |          implicitly returns `()` as its body has no tail or `return` expression
+   |                               -------------- expected `Result<(), ()>` because of return type
+...
+23 |     Ok(());
+   |     ^^^^^^^ expected enum `Result`, found `()`
    |
    = note:   expected enum `Result<(), ()>`
            found unit type `()`
+help: try using a variant of the expected enum
+   |
+23 |     Ok(Ok(());)
+   |
+23 |     Err(Ok(());)
+   |
diff --git a/tokio-macros/src/entry.rs b/tokio-macros/src/entry.rs
index 317b2a9c00d..01f8ee4c1eb 100644
--- a/tokio-macros/src/entry.rs
+++ b/tokio-macros/src/entry.rs
@@ -339,15 +339,17 @@ fn parse_knobs(mut input: syn::ItemFn, is_test: bool, config: FinalConfig) -> To
     let body = &input.block;
     let brace_token = input.block.brace_token;
     let (tail_return, tail_semicolon) = match body.stmts.last() {
-        Some(syn::Stmt::Semi(expr, _)) => (
-            match expr {
-                syn::Expr::Return(_) => quote! { return },
-                _ => quote! {},
-            },
-            quote! {
-                ;
+        Some(syn::Stmt::Semi(expr, _)) => match expr {
+            syn::Expr::Return(_) => (quote! { return }, quote! { ; }),
+            _ => match &input.sig.output {
+                syn::ReturnType::Type(_, ty) if matches!(&**ty, syn::Type::Tuple(ty) if ty.elems.is_empty()) =>
+                {
+                    (quote! {}, quote! { ; }) // unit
+                }
+                syn::ReturnType::Default => (quote! {}, quote! { ; }), // unit
+                syn::ReturnType::Type(..) => (quote! {}, quote! {}),   // ! or another
             },
-        ),
+        },
         _ => (quote! {}, quote! {}),
     };
     input.block = syn::parse2(quote_spanned! {last_stmt_end_span=>
diff --git a/tokio/tests/macros_test.rs b/tokio/tests/macros_test.rs
index 7212c7ba183..bca2c9198a0 100644
--- a/tokio/tests/macros_test.rs
+++ b/tokio/tests/macros_test.rs
@@ -30,3 +30,19 @@ fn trait_method() {
     }
     ().f()
 }
+
+// https://github.com/tokio-rs/tokio/issues/4175
+#[tokio::main]
+pub async fn issue_4175_main_1() -> ! {
+    panic!();
+}
+#[tokio::main]
+pub async fn issue_4175_main_2() -> std::io::Result<()> {
+    panic!();
+}
+#[allow(unreachable_code)]
+#[tokio::test]
+pub async fn issue_4175_test() -> std::io::Result<()> {
+    return Ok(());
+    panic!();
+}