More comments, final tweaks

Author: compiler-errors

compiler/rustc_ast_lowering/src/expr.rs

@@ -1075,6 +1075,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
             body,
             fn_decl_span: self.lower_span(fn_decl_span),
             fn_arg_span: Some(self.lower_span(fn_arg_span)),
+            // Lower this as a `CoroutineClosure`. That will ensure that HIR typeck
+            // knows that a `FnDecl` output type like `-> &str` actually means
+            // "coroutine that returns &str", rather than directly returning a `&str`.
             kind: hir::ClosureKind::CoroutineClosure(hir::CoroutineDesugaring::Async),
             constness: hir::Constness::NotConst,
         });

compiler/rustc_borrowck/src/type_check/input_output.rs

@@ -61,6 +61,8 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                 )),
                 "this needs to be modified if we're lowering non-async closures"
             );
+            // Make sure to use the args from `DefiningTy` so the right NLL region vids are prepopulated
+            // into the type.
             let args = args.as_coroutine_closure();
             let tupled_upvars_ty = ty::CoroutineClosureSignature::tupled_upvars_by_closure_kind(
                 self.tcx(),
@@ -87,11 +89,13 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
                     ty::CoroutineArgsParts {
                         parent_args: args.parent_args(),
                         kind_ty: Ty::from_closure_kind(self.tcx(), args.kind()),
+                        return_ty: user_provided_sig.output(),
+                        tupled_upvars_ty,
+                        // For async closures, none of these can be annotated, so just fill
+                        // them with fresh ty vars.
                         resume_ty: next_ty_var(),
                         yield_ty: next_ty_var(),
                         witness: next_ty_var(),
-                        return_ty: user_provided_sig.output(),
-                        tupled_upvars_ty: tupled_upvars_ty,
                     },
                 )
                 .args,

compiler/rustc_borrowck/src/universal_regions.rs

@@ -720,6 +720,14 @@ impl<'cx, 'tcx> UniversalRegionsBuilder<'cx, 'tcx> {
                 ty::Binder::dummy(inputs_and_output)
             }
 
+            // Construct the signature of the CoroutineClosure for the purposes of borrowck.
+            // This is pretty straightforward -- we:
+            // 1. first grab the `coroutine_closure_sig`,
+            // 2. compute the self type (`&`/`&mut`/no borrow),
+            // 3. flatten the tupled_input_tys,
+            // 4. construct the correct generator type to return with
+            //    `CoroutineClosureSignature::to_coroutine_given_kind_and_upvars`.
+            // Then we wrap it all up into a list of inputs and output.
             DefiningTy::CoroutineClosure(def_id, args) => {
                 assert_eq!(self.mir_def.to_def_id(), def_id);
                 let closure_sig = args.as_coroutine_closure().coroutine_closure_sig();

compiler/rustc_hir_analysis/src/collect.rs

@@ -1533,15 +1533,8 @@ fn coroutine_kind(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option<hir::CoroutineK
 }
 
 fn coroutine_for_closure(tcx: TyCtxt<'_>, def_id: LocalDefId) -> DefId {
-    let Node::Expr(&hir::Expr {
-        kind:
-            hir::ExprKind::Closure(&rustc_hir::Closure {
-                kind: hir::ClosureKind::CoroutineClosure(_),
-                body,
-                ..
-            }),
-        ..
-    }) = tcx.hir_node_by_def_id(def_id)
+    let &rustc_hir::Closure { kind: hir::ClosureKind::CoroutineClosure(_), body, .. } =
+        tcx.hir_node_by_def_id(def_id).expect_closure()
     else {
         bug!()
     };

compiler/rustc_hir_typeck/src/callee.rs

@@ -165,6 +165,11 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 return Some(CallStep::DeferredClosure(def_id, closure_sig));
             }
 
+            // When calling a `CoroutineClosure` that is local to the body, we will
+            // not know what its `closure_kind` is yet. Instead, just fill in the
+            // signature with an infer var for the `tupled_upvars_ty` of the coroutine,
+            // and record a deferred call resolution which will constrain that var
+            // as part of `AsyncFn*` trait confirmation.
             ty::CoroutineClosure(def_id, args) if self.closure_kind(adjusted_ty).is_none() => {
                 let def_id = def_id.expect_local();
                 let closure_args = args.as_coroutine_closure();
@@ -182,6 +187,10 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     coroutine_closure_sig.to_coroutine(
                         self.tcx,
                         closure_args.parent_args(),
+                        // Inherit the kind ty of the closure, since we're calling this
+                        // coroutine with the most relaxed `AsyncFn*` trait that we can.
+                        // We don't necessarily need to do this here, but it saves us
+                        // computing one more infer var that will get constrained later.
                         closure_args.kind_ty(),
                         self.tcx.coroutine_for_closure(def_id),
                         tupled_upvars_ty,

compiler/rustc_hir_typeck/src/closure.rs

@@ -175,6 +175,10 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     interior,
                 ));
 
+                // Coroutines that come from coroutine closures have not yet determined
+                // their kind ty, so make a fresh infer var which will be constrained
+                // later during upvar analysis. Regular coroutines always have the kind
+                // ty of `().`
                 let kind_ty = match kind {
                     hir::CoroutineKind::Desugared(_, hir::CoroutineSource::Closure) => self
                         .next_ty_var(TypeVariableOrigin {
@@ -203,6 +207,8 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 )
             }
             hir::ClosureKind::CoroutineClosure(kind) => {
+                // async closures always return the type ascribed after the `->` (if present),
+                // and yield `()`.
                 let (bound_return_ty, bound_yield_ty) = match kind {
                     hir::CoroutineDesugaring::Async => {
                         (bound_sig.skip_binder().output(), tcx.types.unit)
@@ -211,6 +217,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                         todo!("`gen` and `async gen` closures not supported yet")
                     }
                 };
+                // Compute all of the variables that will be used to populate the coroutine.
                 let resume_ty = self.next_ty_var(TypeVariableOrigin {
                     kind: TypeVariableOriginKind::ClosureSynthetic,
                     span: expr_span,
@@ -258,20 +265,28 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     kind: TypeVariableOriginKind::ClosureSynthetic,
                     span: expr_span,
                 });
+
+                // We need to turn the liberated signature that we got from HIR, which
+                // looks something like `|Args...| -> T`, into a signature that is suitable
+                // for type checking the inner body of the closure, which always returns a
+                // coroutine. To do so, we use the `CoroutineClosureSignature` to compute
+                // the coroutine type, filling in the tupled_upvars_ty and kind_ty with infer
+                // vars which will get constrained during upvar analysis.
+                let coroutine_output_ty = tcx.liberate_late_bound_regions(
+                    expr_def_id.to_def_id(),
+                    closure_args.coroutine_closure_sig().map_bound(|sig| {
+                        sig.to_coroutine(
+                            tcx,
+                            parent_args,
+                            closure_kind_ty,
+                            tcx.coroutine_for_closure(expr_def_id),
+                            coroutine_upvars_ty,
+                        )
+                    }),
+                );
                 liberated_sig = tcx.mk_fn_sig(
                     liberated_sig.inputs().iter().copied(),
-                    tcx.liberate_late_bound_regions(
-                        expr_def_id.to_def_id(),
-                        closure_args.coroutine_closure_sig().map_bound(|sig| {
-                            sig.to_coroutine(
-                                tcx,
-                                parent_args,
-                                closure_kind_ty,
-                                tcx.coroutine_for_closure(expr_def_id),
-                                coroutine_upvars_ty,
-                            )
-                        }),
-                    ),
+                    coroutine_output_ty,
                     liberated_sig.c_variadic,
                     liberated_sig.unsafety,
                     liberated_sig.abi,

compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs

@@ -335,7 +335,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     continue;
                 }
 
-                // For this check, we do *not* want to treat async coroutine-closures (async blocks)
+                // For this check, we do *not* want to treat async coroutine closures (async blocks)
                 // as proper closures. Doing so would regress type inference when feeding
                 // the return value of an argument-position async block to an argument-position
                 // closure wrapped in a block.

compiler/rustc_hir_typeck/src/upvar.rs

@@ -336,6 +336,9 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             }
         }
 
+        // For coroutine-closures, we additionally must compute the
+        // `coroutine_captures_by_ref_ty` type, which is used to generate the by-ref
+        // version of the coroutine-closure's output coroutine.
         if let UpvarArgs::CoroutineClosure(args) = args {
             let closure_env_region: ty::Region<'_> = ty::Region::new_bound(
                 self.tcx,
@@ -353,7 +356,8 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                         self.tcx.coroutine_for_closure(closure_def_id).expect_local(),
                     )
                     // Skip the captures that are just moving the closure's args
-                    // into the coroutine. These are always by move.
+                    // into the coroutine. These are always by move, and we append
+                    // those later in the `CoroutineClosureSignature` helper functions.
                     .skip(
                         args.as_coroutine_closure()
                             .coroutine_closure_sig()
@@ -365,6 +369,9 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     .map(|captured_place| {
                         let upvar_ty = captured_place.place.ty();
                         let capture = captured_place.info.capture_kind;
+                        // Not all upvars are captured by ref, so use
+                        // `apply_capture_kind_on_capture_ty` to ensure that we
+                        // compute the right captured type.
                         apply_capture_kind_on_capture_ty(
                             self.tcx,
                             upvar_ty,
@@ -394,6 +401,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 coroutine_captures_by_ref_ty,
             );
 
+            // Additionally, we can now constrain the coroutine's kind type.
             let ty::Coroutine(_, coroutine_args) =
                 *self.typeck_results.borrow().expr_ty(body.value).kind()
             else {

compiler/rustc_middle/src/ty/flags.rs

@@ -145,11 +145,11 @@ impl FlagComputation {
                     self.flags -= TypeFlags::STILL_FURTHER_SPECIALIZABLE;
                 }
 
-                self.add_ty(args.signature_parts_ty());
-                self.add_ty(args.coroutine_witness_ty());
-                self.add_ty(args.coroutine_captures_by_ref_ty());
                 self.add_ty(args.kind_ty());
+                self.add_ty(args.signature_parts_ty());
                 self.add_ty(args.tupled_upvars_ty());
+                self.add_ty(args.coroutine_captures_by_ref_ty());
+                self.add_ty(args.coroutine_witness_ty());
             }
 
             &ty::Bound(debruijn, _) => {

compiler/rustc_middle/src/ty/instance.rs

@@ -646,9 +646,16 @@ impl<'tcx> Instance<'tcx> {
                 bug!()
             };
 
+            // If the closure's kind ty disagrees with the identity closure's kind ty,
+            // then this must be a coroutine generated by one of the `ConstructCoroutineInClosureShim`s.
             if args.as_coroutine().kind_ty() == id_args.as_coroutine().kind_ty() {
                 Some(Instance { def: ty::InstanceDef::Item(coroutine_def_id), args })
             } else {
+                assert_eq!(
+                    args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap(),
+                    ty::ClosureKind::FnOnce,
+                    "FIXME(async_closures): Generate a by-mut body here."
+                );
                 Some(Instance {
                     def: ty::InstanceDef::CoroutineByMoveShim { coroutine_def_id },
                     args,