Auto merge of #118356 - bvanjoi:merge_coroutinue_into_closure, r=<try>

reduce the calls to `tcx.is_coroutine`

This commit includes two changes:

1. It merges two arms to address the comment raised in #118311 (comment)
2. It attempts to solve the performance regression discussed in #118319 (comment) by reducing the number of calls to `tcx.is_coroutine`.

r? `@cjgillot`

Author: bors

Diff for: compiler/rustc_hir_analysis/src/check/check.rs

@@ -1388,6 +1388,7 @@ fn opaque_type_cycle_error(
                 struct OpaqueTypeCollector {
                     opaques: Vec<DefId>,
                     closures: Vec<DefId>,
+                    coroutine: Vec<DefId>,
                 }
                 impl<'tcx> ty::visit::TypeVisitor<TyCtxt<'tcx>> for OpaqueTypeCollector {
                     fn visit_ty(&mut self, t: Ty<'tcx>) -> ControlFlow<Self::BreakTy> {
@@ -1396,10 +1397,14 @@ fn opaque_type_cycle_error(
                                 self.opaques.push(def);
                                 ControlFlow::Continue(())
                             }
-                            ty::Closure(def_id, ..) | ty::Coroutine(def_id, ..) => {
+                            ty::Closure(def_id, ..) => {
                                 self.closures.push(def_id);
                                 t.super_visit_with(self)
                             }
+                            ty::Coroutine(def_id, ..) => {
+                                self.coroutine.push(def_id);
+                                t.super_visit_with(self)
+                            }
                             _ => t.super_visit_with(self),
                         }
                     }
@@ -1417,43 +1422,48 @@ fn opaque_type_cycle_error(
                     err.span_label(sp, format!("returning here with type `{ty}`"));
                 }
 
-                for closure_def_id in visitor.closures {
-                    let Some(closure_local_did) = closure_def_id.as_local() else {
-                        continue;
-                    };
-                    let typeck_results = tcx.typeck(closure_local_did);
-
-                    let mut label_match = |ty: Ty<'_>, span| {
-                        for arg in ty.walk() {
-                            if let ty::GenericArgKind::Type(ty) = arg.unpack()
-                                && let ty::Alias(
-                                    ty::Opaque,
-                                    ty::AliasTy { def_id: captured_def_id, .. },
-                                ) = *ty.kind()
-                                && captured_def_id == opaque_def_id.to_def_id()
-                            {
-                                err.span_label(
-                                    span,
-                                    format!(
-                                        "{} captures itself here",
-                                        tcx.def_descr(closure_def_id)
-                                    ),
-                                );
-                            }
+                let label_match = |err: &mut DiagnosticBuilder<'_, _>, ty: Ty<'_>, span, def_id| {
+                    for arg in ty.walk() {
+                        if let ty::GenericArgKind::Type(ty) = arg.unpack()
+                            && let ty::Alias(
+                                ty::Opaque,
+                                ty::AliasTy { def_id: captured_def_id, .. },
+                            ) = *ty.kind()
+                            && captured_def_id == opaque_def_id.to_def_id()
+                        {
+                            err.span_label(
+                                span,
+                                format!("{} captures itself here", tcx.def_descr(def_id)),
+                            );
                         }
-                    };
+                    }
+                };
 
+                let capture = |err: &mut DiagnosticBuilder<'_, _>, def_id: DefId| {
+                    let Some(local_id) = def_id.as_local() else {
+                        return;
+                    };
+                    let typeck_results = tcx.typeck(local_id);
                     // Label any closure upvars that capture the opaque
-                    for capture in typeck_results.closure_min_captures_flattened(closure_local_did)
-                    {
-                        label_match(capture.place.ty(), capture.get_path_span(tcx));
+                    for capture in typeck_results.closure_min_captures_flattened(local_id) {
+                        label_match(err, capture.place.ty(), capture.get_path_span(tcx), def_id);
                     }
-                    // Label any coroutine locals that capture the opaque
-                    if tcx.is_coroutine(closure_def_id)
-                        && let Some(coroutine_layout) = tcx.mir_coroutine_witnesses(closure_def_id)
-                    {
+                };
+
+                for closure_def_id in visitor.closures {
+                    capture(&mut err, closure_def_id);
+                }
+
+                for coroutine_def_id in visitor.coroutine {
+                    capture(&mut err, coroutine_def_id);
+                    if let Some(coroutine_layout) = tcx.mir_coroutine_witnesses(coroutine_def_id) {
                         for interior_ty in &coroutine_layout.field_tys {
-                            label_match(interior_ty.ty, interior_ty.source_info.span);
+                            label_match(
+                                &mut err,
+                                interior_ty.ty,
+                                interior_ty.source_info.span,
+                                coroutine_def_id,
+                            );
                         }
                     }
                 }

Diff for: compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -1617,6 +1617,8 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
         for (def_id, encode_const, encode_opt) in keys_and_jobs {
             debug_assert!(encode_const || encode_opt);
 
+            let def_kind = tcx.def_kind(def_id);
+
             debug!("EntryBuilder::encode_mir({:?})", def_id);
             if encode_opt {
                 record!(self.tables.optimized_mir[def_id.to_def_id()] <- tcx.optimized_mir(def_id));
@@ -1626,7 +1628,8 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
                 record!(self.tables.closure_saved_names_of_captured_variables[def_id.to_def_id()]
                     <- tcx.closure_saved_names_of_captured_variables(def_id));
 
-                if self.tcx.is_coroutine(def_id.to_def_id())
+                if DefKind::Closure == def_kind
+                    && tcx.is_coroutine(def_id.to_def_id())
                     && let Some(witnesses) = tcx.mir_coroutine_witnesses(def_id)
                 {
                     record!(self.tables.mir_coroutine_witnesses[def_id.to_def_id()] <- witnesses);
@@ -1653,7 +1656,8 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
             }
             record!(self.tables.promoted_mir[def_id.to_def_id()] <- tcx.promoted_mir(def_id));
 
-            if self.tcx.is_coroutine(def_id.to_def_id())
+            if DefKind::Closure == def_kind
+                && tcx.is_coroutine(def_id.to_def_id())
                 && let Some(witnesses) = tcx.mir_coroutine_witnesses(def_id)
             {
                 record!(self.tables.mir_coroutine_witnesses[def_id.to_def_id()] <- witnesses);

Diff for: compiler/rustc_mir_build/src/build/mod.rs

@@ -638,25 +638,25 @@ fn construct_error(tcx: TyCtxt<'_>, def_id: LocalDefId, guar: ErrorGuaranteed) -
             );
             (sig.inputs().to_vec(), sig.output(), None)
         }
-        DefKind::Closure if coroutine_kind.is_some() => {
-            let coroutine_ty = tcx.type_of(def_id).instantiate_identity();
-            let ty::Coroutine(_, args, _) = coroutine_ty.kind() else { bug!() };
-            let args = args.as_coroutine();
-            let yield_ty = args.yield_ty();
-            let return_ty = args.return_ty();
-            (vec![coroutine_ty, args.resume_ty()], return_ty, Some(yield_ty))
-        }
         DefKind::Closure => {
-            let closure_ty = tcx.type_of(def_id).instantiate_identity();
-            let ty::Closure(_, args) = closure_ty.kind() else { bug!() };
-            let args = args.as_closure();
-            let sig = tcx.liberate_late_bound_regions(def_id.to_def_id(), args.sig());
-            let self_ty = match args.kind() {
-                ty::ClosureKind::Fn => Ty::new_imm_ref(tcx, tcx.lifetimes.re_erased, closure_ty),
-                ty::ClosureKind::FnMut => Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, closure_ty),
-                ty::ClosureKind::FnOnce => closure_ty,
-            };
-            ([self_ty].into_iter().chain(sig.inputs().to_vec()).collect(), sig.output(), None)
+            let ty = tcx.type_of(def_id).instantiate_identity();
+            if let ty::Coroutine(_, args, _) = ty.kind() {
+                let args = args.as_coroutine();
+                let yield_ty = args.yield_ty();
+                let return_ty = args.return_ty();
+                (vec![ty, args.resume_ty()], return_ty, Some(yield_ty))
+            } else if let ty::Closure(_, args) = ty.kind() {
+                let args = args.as_closure();
+                let sig = tcx.liberate_late_bound_regions(def_id.to_def_id(), args.sig());
+                let self_ty = match args.kind() {
+                    ty::ClosureKind::Fn => Ty::new_imm_ref(tcx, tcx.lifetimes.re_erased, ty),
+                    ty::ClosureKind::FnMut => Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, ty),
+                    ty::ClosureKind::FnOnce => ty,
+                };
+                ([self_ty].into_iter().chain(sig.inputs().to_vec()).collect(), sig.output(), None)
+            } else {
+                bug!()
+            }
         }
         dk => bug!("{:?} is not a body: {:?}", def_id, dk),
     };

Diff for: compiler/rustc_mir_build/src/thir/cx/mod.rs

@@ -37,7 +37,10 @@ pub(crate) fn thir_body(
 
         // The resume argument may be missing, in that case we need to provide it here.
         // It will always be `()` in this case.
-        if tcx.is_coroutine(owner_def.to_def_id()) && body.params.is_empty() {
+        if body.params.is_empty()
+            && tcx.def_kind(owner_def) == DefKind::Closure
+            && tcx.is_coroutine(owner_def.to_def_id())
+        {
             cx.thir.params.push(Param {
                 ty: Ty::new_unit(tcx),
                 pat: None,
@@ -119,45 +122,30 @@ impl<'tcx> Cx<'tcx> {
 
     fn closure_env_param(&self, owner_def: LocalDefId, owner_id: HirId) -> Option<Param<'tcx>> {
         match self.tcx.def_kind(owner_def) {
-            DefKind::Closure if self.tcx.is_coroutine(owner_def.to_def_id()) => {
-                let coroutine_ty = self.typeck_results.node_type(owner_id);
-                let coroutine_param = Param {
-                    ty: coroutine_ty,
-                    pat: None,
-                    ty_span: None,
-                    self_kind: None,
-                    hir_id: None,
-                };
-                Some(coroutine_param)
-            }
             DefKind::Closure => {
-                let closure_ty = self.typeck_results.node_type(owner_id);
-
-                let ty::Closure(closure_def_id, closure_args) = *closure_ty.kind() else {
-                    bug!("closure expr does not have closure type: {:?}", closure_ty);
+                let ty = self.typeck_results.node_type(owner_id);
+                let ty = if ty.is_coroutine() {
+                    ty
+                } else if let ty::Closure(closure_def_id, closure_args) = *ty.kind() {
+                    let bound_vars = self
+                        .tcx
+                        .mk_bound_variable_kinds(&[ty::BoundVariableKind::Region(ty::BrEnv)]);
+                    let br = ty::BoundRegion {
+                        var: ty::BoundVar::from_usize(bound_vars.len() - 1),
+                        kind: ty::BrEnv,
+                    };
+                    let env_region = ty::Region::new_bound(self.tcx, ty::INNERMOST, br);
+                    let closure_env_ty =
+                        self.tcx.closure_env_ty(closure_def_id, closure_args, env_region).unwrap();
+                    self.tcx.instantiate_bound_regions_with_erased(ty::Binder::bind_with_vars(
+                        closure_env_ty,
+                        bound_vars,
+                    ))
+                } else {
+                    bug!("closure expr does not have closure type: {:?}", ty);
                 };
-
-                let bound_vars =
-                    self.tcx.mk_bound_variable_kinds(&[ty::BoundVariableKind::Region(ty::BrEnv)]);
-                let br = ty::BoundRegion {
-                    var: ty::BoundVar::from_usize(bound_vars.len() - 1),
-                    kind: ty::BrEnv,
-                };
-                let env_region = ty::Region::new_bound(self.tcx, ty::INNERMOST, br);
-                let closure_env_ty =
-                    self.tcx.closure_env_ty(closure_def_id, closure_args, env_region).unwrap();
-                let liberated_closure_env_ty = self.tcx.instantiate_bound_regions_with_erased(
-                    ty::Binder::bind_with_vars(closure_env_ty, bound_vars),
-                );
-                let env_param = Param {
-                    ty: liberated_closure_env_ty,
-                    pat: None,
-                    ty_span: None,
-                    self_kind: None,
-                    hir_id: None,
-                };
-
-                Some(env_param)
+                let param = Param { ty, pat: None, ty_span: None, self_kind: None, hir_id: None };
+                Some(param)
             }
             _ => None,
         }

Diff for: compiler/rustc_mir_transform/src/const_prop.rs

@@ -84,7 +84,7 @@ impl<'tcx> MirPass<'tcx> for ConstProp {
 
         // FIXME(welseywiser) const prop doesn't work on coroutines because of query cycles
         // computing their layout.
-        if tcx.is_coroutine(def_id.to_def_id()) {
+        if def_kind == DefKind::Closure && tcx.is_coroutine(def_id.to_def_id()) {
             trace!("ConstProp skipped for coroutine {:?}", def_id);
             return;
         }

Diff for: compiler/rustc_mir_transform/src/const_prop_lint.rs

@@ -61,7 +61,7 @@ impl<'tcx> MirLint<'tcx> for ConstPropLint {
 
         // FIXME(welseywiser) const prop doesn't work on coroutines because of query cycles
         // computing their layout.
-        if tcx.is_coroutine(def_id.to_def_id()) {
+        if def_kind == DefKind::Closure && tcx.is_coroutine(def_id.to_def_id()) {
             trace!("ConstPropLint skipped for coroutine {:?}", def_id);
             return;
         }

Diff for: compiler/rustc_mir_transform/src/lib.rs

@@ -395,13 +395,13 @@ fn inner_mir_for_ctfe(tcx: TyCtxt<'_>, def: LocalDefId) -> Body<'_> {
 /// mir borrowck *before* doing so in order to ensure that borrowck can be run and doesn't
 /// end up missing the source MIR due to stealing happening.
 fn mir_drops_elaborated_and_const_checked(tcx: TyCtxt<'_>, def: LocalDefId) -> &Steal<Body<'_>> {
-    if tcx.is_coroutine(def.to_def_id()) {
+    let def_kind = tcx.def_kind(def);
+    if def_kind == DefKind::Closure && tcx.is_coroutine(def.to_def_id()) {
         tcx.ensure_with_value().mir_coroutine_witnesses(def);
     }
     let mir_borrowck = tcx.mir_borrowck(def);
 
-    let is_fn_like = tcx.def_kind(def).is_fn_like();
-    if is_fn_like {
+    if def_kind.is_fn_like() {
         // Do not compute the mir call graph without said call graph actually being used.
         if pm::should_run_pass(tcx, &inline::Inline) {
             tcx.ensure_with_value().mir_inliner_callees(ty::InstanceDef::Item(def.to_def_id()));