Compute yield and return types outside of check_fn

compiler/rustc_hir_typeck/src/check.rs

@@ -28,10 +28,10 @@
 pub(super) fn check_fn<'a, 'tcx>(
     fcx: &mut FnCtxt<'a, 'tcx>,
     fn_sig: ty::FnSig<'tcx>,
+    coroutine_types: Option<CoroutineTypes<'tcx>>,
     decl: &'tcx hir::FnDecl<'tcx>,
     fn_def_id: LocalDefId,
     body: &'tcx hir::Body<'tcx>,
-    closure_kind: Option<hir::ClosureKind>,
     params_can_be_unsized: bool,
 ) -> Option<CoroutineTypes<'tcx>> {
     let fn_id = fcx.tcx.local_def_id_to_hir_id(fn_def_id);
@@ -49,54 +49,13 @@ pub(super) fn check_fn<'a, 'tcx>(
             fcx.param_env,
         ));
 
+    fcx.coroutine_types = coroutine_types;
     fcx.ret_coercion = Some(RefCell::new(CoerceMany::new(ret_ty)));
 
     let span = body.value.span;
 
     forbid_intrinsic_abi(tcx, span, fn_sig.abi);
 
-    if let Some(hir::ClosureKind::Coroutine(kind)) = closure_kind {
-        let yield_ty = match kind {
-            hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::Gen, _)
-            | hir::CoroutineKind::Coroutine(_) => {
-                let yield_ty = fcx.next_ty_var(TypeVariableOrigin {
-                    kind: TypeVariableOriginKind::TypeInference,
-                    span,
-                });
-                fcx.require_type_is_sized(yield_ty, span, traits::SizedYieldType);
-                yield_ty
-            }
-            // HACK(-Ztrait-solver=next): In the *old* trait solver, we must eagerly
-            // guide inference on the yield type so that we can handle `AsyncIterator`
-            // in this block in projection correctly. In the new trait solver, it is
-            // not a problem.
-            hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::AsyncGen, _) => {
-                let yield_ty = fcx.next_ty_var(TypeVariableOrigin {
-                    kind: TypeVariableOriginKind::TypeInference,
-                    span,
-                });
-                fcx.require_type_is_sized(yield_ty, span, traits::SizedYieldType);
-
-                Ty::new_adt(
-                    tcx,
-                    tcx.adt_def(tcx.require_lang_item(hir::LangItem::Poll, Some(span))),
-                    tcx.mk_args(&[Ty::new_adt(
-                        tcx,
-                        tcx.adt_def(tcx.require_lang_item(hir::LangItem::Option, Some(span))),
-                        tcx.mk_args(&[yield_ty.into()]),
-                    )
-                    .into()]),
-                )
-            }
-            hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::Async, _) => Ty::new_unit(tcx),
-        };
-
-        // Resume type defaults to `()` if the coroutine has no argument.
-        let resume_ty = fn_sig.inputs().get(0).copied().unwrap_or_else(|| Ty::new_unit(tcx));
-
-        fcx.coroutine_types = Some(CoroutineTypes { resume_ty, yield_ty });
-    }
-
     GatherLocalsVisitor::new(fcx).visit_body(body);
 
     // C-variadic fns also have a `VaList` input that's not listed in `fn_sig`

compiler/rustc_hir_typeck/src/closure.rs

@@ -72,7 +72,8 @@ fn check_closure(
         opt_kind: Option<ty::ClosureKind>,
         expected_sig: Option<ExpectedSig<'tcx>>,
     ) -> Ty<'tcx> {
-        let body = self.tcx.hir().body(closure.body);
+        let tcx = self.tcx;
+        let body = tcx.hir().body(closure.body);
 
         trace!("decl = {:#?}", closure.fn_decl);
         let expr_def_id = closure.def_id;
@@ -83,26 +84,79 @@ fn check_closure(
 
         debug!(?bound_sig, ?liberated_sig);
 
+        // FIXME: We could probably actually just unify this further --
+        // instead of having a `FnSig` and a `Option<CoroutineTypes>`,
+        // we can have a `ClosureSignature { Coroutine { .. }, Closure { .. } }`,
+        // similar to how `ty::GenSig` is a distinct data structure.
+        let coroutine_types = match closure.kind {
+            hir::ClosureKind::Closure => None,
+            hir::ClosureKind::Coroutine(kind) => {
+                let yield_ty = match kind {
+                    hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::Gen, _)
+                    | hir::CoroutineKind::Coroutine(_) => {
+                        let yield_ty = self.next_ty_var(TypeVariableOrigin {
+                            kind: TypeVariableOriginKind::TypeInference,
+                            span: expr_span,
+                        });
+                        self.require_type_is_sized(yield_ty, expr_span, traits::SizedYieldType);
+                        yield_ty
+                    }
+                    // HACK(-Ztrait-solver=next): In the *old* trait solver, we must eagerly
+                    // guide inference on the yield type so that we can handle `AsyncIterator`
+                    // in this block in projection correctly. In the new trait solver, it is
+                    // not a problem.
+                    hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::AsyncGen, _) => {
+                        let yield_ty = self.next_ty_var(TypeVariableOrigin {
+                            kind: TypeVariableOriginKind::TypeInference,
+                            span: expr_span,
+                        });
+                        self.require_type_is_sized(yield_ty, expr_span, traits::SizedYieldType);
+
+                        Ty::new_adt(
+                            tcx,
+                            tcx.adt_def(
+                                tcx.require_lang_item(hir::LangItem::Poll, Some(expr_span)),
+                            ),
+                            tcx.mk_args(&[Ty::new_adt(
+                                tcx,
+                                tcx.adt_def(
+                                    tcx.require_lang_item(hir::LangItem::Option, Some(expr_span)),
+                                ),
+                                tcx.mk_args(&[yield_ty.into()]),
+                            )
+                            .into()]),
+                        )
+                    }
+                    hir::CoroutineKind::Desugared(hir::CoroutineDesugaring::Async, _) => {
+                        tcx.types.unit
+                    }
+                };
+
+                // Resume type defaults to `()` if the coroutine has no argument.
+                let resume_ty = liberated_sig.inputs().get(0).copied().unwrap_or(tcx.types.unit);
+
+                Some(CoroutineTypes { resume_ty, yield_ty })
+            }
+        };
+
         let mut fcx = FnCtxt::new(self, self.param_env, closure.def_id);
-        let coroutine_types = check_fn(
+        check_fn(
             &mut fcx,
             liberated_sig,
+            coroutine_types,
             closure.fn_decl,
             expr_def_id,
             body,
-            Some(closure.kind),
             // Closure "rust-call" ABI doesn't support unsized params
             false,
         );
 
-        let parent_args = GenericArgs::identity_for_item(
+        let parent_args =
-            self.tcx,
+            GenericArgs::identity_for_item(tcx, tcx.typeck_root_def_id(expr_def_id.to_def_id()));
-            self.tcx.typeck_root_def_id(expr_def_id.to_def_id()),
-        );
 
         let tupled_upvars_ty = self.next_root_ty_var(TypeVariableOrigin {
             kind: TypeVariableOriginKind::ClosureSynthetic,
-            span: self.tcx.def_span(expr_def_id),
+            span: expr_span,
         });
 
         match closure.kind {
@@ -111,8 +165,8 @@ fn check_closure(
                 // Tuple up the arguments and insert the resulting function type into
                 // the `closures` table.
                 let sig = bound_sig.map_bound(|sig| {
-                    self.tcx.mk_fn_sig(
+                    tcx.mk_fn_sig(
-                        [Ty::new_tup(self.tcx, sig.inputs())],
+                        [Ty::new_tup(tcx, sig.inputs())],
                         sig.output(),
                         sig.c_variadic,
                         sig.unsafety,
@@ -123,7 +177,7 @@ fn check_closure(
                 debug!(?sig, ?opt_kind);
 
                 let closure_kind_ty = match opt_kind {
-                    Some(kind) => Ty::from_closure_kind(self.tcx, kind),
+                    Some(kind) => Ty::from_closure_kind(tcx, kind),
 
                     // Create a type variable (for now) to represent the closure kind.
                     // It will be unified during the upvar inference phase (`upvar.rs`)
@@ -135,16 +189,16 @@ fn check_closure(
                 };
 
                 let closure_args = ty::ClosureArgs::new(
-                    self.tcx,
+                    tcx,
                     ty::ClosureArgsParts {
                         parent_args,
                         closure_kind_ty,
-                        closure_sig_as_fn_ptr_ty: Ty::new_fn_ptr(self.tcx, sig),
+                        closure_sig_as_fn_ptr_ty: Ty::new_fn_ptr(tcx, sig),
                         tupled_upvars_ty,
                     },
                 );
 
-                Ty::new_closure(self.tcx, expr_def_id.to_def_id(), closure_args.args)
+                Ty::new_closure(tcx, expr_def_id.to_def_id(), closure_args.args)
             }
             hir::ClosureKind::Coroutine(_) => {
                 let Some(CoroutineTypes { resume_ty, yield_ty }) = coroutine_types else {
@@ -161,7 +215,7 @@ fn check_closure(
                 ));
 
                 let coroutine_args = ty::CoroutineArgs::new(
-                    self.tcx,
+                    tcx,
                     ty::CoroutineArgsParts {
                         parent_args,
                         resume_ty,
@@ -172,7 +226,7 @@ fn check_closure(
                     },
                 );
 
-                Ty::new_coroutine(self.tcx, expr_def_id.to_def_id(), coroutine_args.args)
+                Ty::new_coroutine(tcx, expr_def_id.to_def_id(), coroutine_args.args)
             }
         }
     }

compiler/rustc_hir_typeck/src/lib.rs

@@ -193,7 +193,7 @@ fn typeck_with_fallback<'tcx>(
         let fn_sig = tcx.liberate_late_bound_regions(def_id.to_def_id(), fn_sig);
         let fn_sig = fcx.normalize(body.value.span, fn_sig);
 
-        check_fn(&mut fcx, fn_sig, decl, def_id, body, None, tcx.features().unsized_fn_params);
+        check_fn(&mut fcx, fn_sig, None, decl, def_id, body, tcx.features().unsized_fn_params);
     } else {
         let expected_type = if let Some(&hir::Ty { kind: hir::TyKind::Infer, span, .. }) = body_ty {
             Some(fcx.next_ty_var(TypeVariableOrigin {

tests/ui/coroutine/gen_block.e2024.stderr

@@ -8,10 +8,10 @@ LL |     let _ = || yield true;
    = help: add `#![feature(coroutines)]` to the crate attributes to enable
 
 error[E0282]: type annotations needed
-  --> $DIR/gen_block.rs:6:17
+  --> $DIR/gen_block.rs:6:13
    |
 LL |     let x = gen {};
-   |                 ^^ cannot infer type
+   |             ^^^^^^ cannot infer type
 
 error: aborting due to 2 previous errors
 

tests/ui/coroutine/sized-yield.stderr

@@ -1,8 +1,8 @@
 error[E0277]: the size for values of type `str` cannot be known at compilation time
-  --> $DIR/sized-yield.rs:8:27
+  --> $DIR/sized-yield.rs:8:19
    |
 LL |       let mut gen = move || {
-   |  ___________________________^
+   |  ___________________^
 LL | |
 LL | |         yield s[..];
 LL | |     };

tests/ui/feature-gates/feature-gate-gen_blocks.e2024.stderr

@@ -35,16 +35,16 @@ LL |     async gen {};
    = help: add `#![feature(gen_blocks)]` to the crate attributes to enable
 
 error[E0282]: type annotations needed
-  --> $DIR/feature-gate-gen_blocks.rs:5:9
+  --> $DIR/feature-gate-gen_blocks.rs:5:5
    |
 LL |     gen {};
-   |         ^^ cannot infer type
+   |     ^^^^^^ cannot infer type
 
 error[E0282]: type annotations needed
-  --> $DIR/feature-gate-gen_blocks.rs:12:15
+  --> $DIR/feature-gate-gen_blocks.rs:12:5
    |
 LL |     async gen {};
-   |               ^^ cannot infer type
+   |     ^^^^^^^^^^^^ cannot infer type
 
 error: aborting due to 6 previous errors