From bb33200047722e321480f7e5004a74395804958d Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Mon, 18 Dec 2023 01:45:09 +0000
Subject: [PATCH] Make sure all kinds of generators only return unit

---
 compiler/rustc_hir_typeck/src/closure.rs |  7 +++---
 tests/ui/coroutine/return-types.rs       | 21 ++++++++++++++++
 tests/ui/coroutine/return-types.stderr   | 31 ++++++++++++++++++++++++
 3 files changed, 56 insertions(+), 3 deletions(-)
 create mode 100644 tests/ui/coroutine/return-types.rs
 create mode 100644 tests/ui/coroutine/return-types.stderr

diff --git a/compiler/rustc_hir_typeck/src/closure.rs b/compiler/rustc_hir_typeck/src/closure.rs
index 7e43d67587b..d19d304128a 100644
--- a/compiler/rustc_hir_typeck/src/closure.rs
+++ b/compiler/rustc_hir_typeck/src/closure.rs
@@ -650,9 +650,10 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                         },
                     )
                 }
-                // For a `gen {}` block created as a `gen fn` body, we need the return type to be
-                // ().
-                Some(hir::CoroutineKind::Gen(hir::CoroutineSource::Fn)) => self.tcx.types.unit,
+                // All `gen {}` and `async gen {}` must return unit.
+                Some(hir::CoroutineKind::Gen(_) | hir::CoroutineKind::AsyncGen(_)) => {
+                    self.tcx.types.unit
+                }
 
                 _ => astconv.ty_infer(None, decl.output.span()),
             },
diff --git a/tests/ui/coroutine/return-types.rs b/tests/ui/coroutine/return-types.rs
new file mode 100644
index 00000000000..3543d6293f7
--- /dev/null
+++ b/tests/ui/coroutine/return-types.rs
@@ -0,0 +1,21 @@
+// compile-flags: --edition 2024 -Zunstable-options
+
+#![feature(gen_blocks)]
+
+async gen fn async_gen_fn() -> i32 { 0 }
+//~^ ERROR mismatched types
+
+gen fn gen_fn() -> i32 { 0 }
+//~^ ERROR mismatched types
+
+fn async_gen_block() {
+    async gen { yield (); 1 };
+    //~^ ERROR mismatched types
+}
+
+fn gen_block() {
+    gen { yield (); 1 };
+    //~^ ERROR mismatched types
+}
+
+fn main() {}
diff --git a/tests/ui/coroutine/return-types.stderr b/tests/ui/coroutine/return-types.stderr
new file mode 100644
index 00000000000..7be96e538d9
--- /dev/null
+++ b/tests/ui/coroutine/return-types.stderr
@@ -0,0 +1,31 @@
+error[E0308]: mismatched types
+  --> $DIR/return-types.rs:5:38
+   |
+LL | async gen fn async_gen_fn() -> i32 { 0 }
+   |                                ---   ^ expected `()`, found integer
+   |                                |
+   |                                expected `()` because of return type
+
+error[E0308]: mismatched types
+  --> $DIR/return-types.rs:8:26
+   |
+LL | gen fn gen_fn() -> i32 { 0 }
+   |                    ---   ^ expected `()`, found integer
+   |                    |
+   |                    expected `()` because of return type
+
+error[E0308]: mismatched types
+  --> $DIR/return-types.rs:12:27
+   |
+LL |     async gen { yield (); 1 };
+   |                           ^ expected `()`, found integer
+
+error[E0308]: mismatched types
+  --> $DIR/return-types.rs:17:21
+   |
+LL |     gen { yield (); 1 };
+   |                     ^ expected `()`, found integer
+
+error: aborting due to 4 previous errors
+
+For more information about this error, try `rustc --explain E0308`.