[PERF] Collect higher-ranked coroutines outlives assumptions

compiler/rustc_macros/src/query.rs

@@ -51,6 +51,7 @@ impl Parse for Query {
         let key = Pat::parse_single(&arg_content)?;
         arg_content.parse::<Token![:]>()?;
         let arg = arg_content.parse()?;
+        let _ = arg_content.parse::<Option<Token![,]>>()?;
         let result = input.parse()?;
 
         // Parse the query modifiers

compiler/rustc_middle/src/query/erase.rs

@@ -6,7 +6,7 @@ use rustc_span::ErrorGuaranteed;
 
 use crate::query::CyclePlaceholder;
 use crate::ty::adjustment::CoerceUnsizedInfo;
-use crate::ty::{self, Ty};
+use crate::ty::{self, Ty, TyCtxt};
 use crate::{mir, traits};
 
 #[derive(Copy, Clone)]
@@ -207,6 +207,11 @@ impl EraseType for ty::Binder<'_, ty::FnSig<'_>> {
     type Result = [u8; size_of::<ty::Binder<'static, ty::FnSig<'static>>>()];
 }
 
+impl EraseType for ty::Binder<'_, ty::CoroutineWitnessTypes<TyCtxt<'_>>> {
+    type Result =
+        [u8; size_of::<ty::Binder<'static, ty::CoroutineWitnessTypes<TyCtxt<'static>>>>()];
+}
+
 impl EraseType for ty::Binder<'_, &'_ ty::List<Ty<'_>>> {
     type Result = [u8; size_of::<ty::Binder<'static, &'static ty::List<Ty<'static>>>>()];
 }

compiler/rustc_middle/src/query/mod.rs

@@ -922,6 +922,12 @@ rustc_queries! {
         separate_provide_extern
     }
 
+    query coroutine_hidden_types(
+        def_id: DefId,
+    ) -> ty::EarlyBinder<'tcx, ty::Binder<'tcx, ty::CoroutineWitnessTypes<TyCtxt<'tcx>>>> {
+        desc { "looking up the hidden types stored across await points in a coroutine" }
+    }
+
     /// Gets a map with the variances of every item in the local crate.
     ///
     /// <div class="warning">

compiler/rustc_middle/src/ty/context.rs

@@ -154,6 +154,8 @@ impl<'tcx> Interner for TyCtxt<'tcx> {
     type BoundRegion = ty::BoundRegion;
     type PlaceholderRegion = ty::PlaceholderRegion;
 
+    type RegionAssumptions = &'tcx ty::List<ty::OutlivesPredicate<'tcx, ty::GenericArg<'tcx>>>;
+
     type ParamEnv = ty::ParamEnv<'tcx>;
     type Predicate = Predicate<'tcx>;
 
@@ -340,7 +342,7 @@ impl<'tcx> Interner for TyCtxt<'tcx> {
     fn coroutine_hidden_types(
         self,
         def_id: DefId,
-    ) -> ty::EarlyBinder<'tcx, ty::Binder<'tcx, &'tcx ty::List<Ty<'tcx>>>> {
+    ) -> ty::EarlyBinder<'tcx, ty::Binder<'tcx, ty::CoroutineWitnessTypes<TyCtxt<'tcx>>>> {
         self.coroutine_hidden_types(def_id)
     }
 
@@ -854,6 +856,7 @@ pub struct CtxtInterners<'tcx> {
     offset_of: InternedSet<'tcx, List<(VariantIdx, FieldIdx)>>,
     valtree: InternedSet<'tcx, ty::ValTreeKind<'tcx>>,
     patterns: InternedSet<'tcx, List<ty::Pattern<'tcx>>>,
+    outlives: InternedSet<'tcx, List<ty::OutlivesPredicate<'tcx, ty::GenericArg<'tcx>>>>,
 }
 
 impl<'tcx> CtxtInterners<'tcx> {
@@ -891,6 +894,7 @@ impl<'tcx> CtxtInterners<'tcx> {
             offset_of: InternedSet::with_capacity(N),
             valtree: InternedSet::with_capacity(N),
             patterns: InternedSet::with_capacity(N),
+            outlives: InternedSet::with_capacity(N),
         }
     }
 
@@ -2685,6 +2689,7 @@ slice_interners!(
     captures: intern_captures(&'tcx ty::CapturedPlace<'tcx>),
     offset_of: pub mk_offset_of((VariantIdx, FieldIdx)),
     patterns: pub mk_patterns(Pattern<'tcx>),
+    outlives: pub mk_outlives(ty::OutlivesPredicate<'tcx, ty::GenericArg<'tcx>>),
 );
 
 impl<'tcx> TyCtxt<'tcx> {
@@ -3100,6 +3105,17 @@ impl<'tcx> TyCtxt<'tcx> {
         T::collect_and_apply(iter, |xs| self.mk_bound_variable_kinds(xs))
     }
 
+    pub fn mk_outlives_from_iter<I, T>(self, iter: I) -> T::Output
+    where
+        I: Iterator<Item = T>,
+        T: CollectAndApply<
+                ty::OutlivesPredicate<'tcx, ty::GenericArg<'tcx>>,
+                &'tcx ty::List<ty::OutlivesPredicate<'tcx, ty::GenericArg<'tcx>>>,
+            >,
+    {
+        T::collect_and_apply(iter, |xs| self.mk_outlives(xs))
+    }
+
     /// Emit a lint at `span` from a lint struct (some type that implements `LintDiagnostic`,
     /// typically generated by `#[derive(LintDiagnostic)]`).
     #[track_caller]

compiler/rustc_middle/src/ty/structural_impls.rs

@@ -779,4 +779,5 @@ list_fold! {
     &'tcx ty::List<ty::PolyExistentialPredicate<'tcx>> : mk_poly_existential_predicates,
     &'tcx ty::List<PlaceElem<'tcx>> : mk_place_elems,
     &'tcx ty::List<ty::Pattern<'tcx>> : mk_patterns,
+    &'tcx ty::List<ty::OutlivesPredicate<'tcx, ty::GenericArg<'tcx>>> : mk_outlives,
 }

compiler/rustc_middle/src/ty/util.rs

@@ -26,7 +26,7 @@ use crate::query::Providers;
 use crate::ty::layout::{FloatExt, IntegerExt};
 use crate::ty::{
     self, Asyncness, FallibleTypeFolder, GenericArgKind, GenericArgsRef, Ty, TyCtxt, TypeFoldable,
-    TypeFolder, TypeSuperFoldable, TypeVisitableExt, Upcast, fold_regions,
+    TypeFolder, TypeSuperFoldable, TypeVisitableExt, Upcast,
 };
 
 #[derive(Copy, Clone, Debug)]
@@ -737,40 +737,6 @@ impl<'tcx> TyCtxt<'tcx> {
         }
     }
 
-    /// Return the set of types that should be taken into account when checking
-    /// trait bounds on a coroutine's internal state. This properly replaces
-    /// `ReErased` with new existential bound lifetimes.
-    pub fn coroutine_hidden_types(
-        self,
-        def_id: DefId,
-    ) -> ty::EarlyBinder<'tcx, ty::Binder<'tcx, &'tcx ty::List<Ty<'tcx>>>> {
-        let coroutine_layout = self.mir_coroutine_witnesses(def_id);
-        let mut vars = vec![];
-        let bound_tys = self.mk_type_list_from_iter(
-            coroutine_layout
-                .as_ref()
-                .map_or_else(|| [].iter(), |l| l.field_tys.iter())
-                .filter(|decl| !decl.ignore_for_traits)
-                .map(|decl| {
-                    let ty = fold_regions(self, decl.ty, |re, debruijn| {
-                        assert_eq!(re, self.lifetimes.re_erased);
-                        let var = ty::BoundVar::from_usize(vars.len());
-                        vars.push(ty::BoundVariableKind::Region(ty::BoundRegionKind::Anon));
-                        ty::Region::new_bound(
-                            self,
-                            debruijn,
-                            ty::BoundRegion { var, kind: ty::BoundRegionKind::Anon },
-                        )
-                    });
-                    ty
-                }),
-        );
-        ty::EarlyBinder::bind(ty::Binder::bind_with_vars(
-            bound_tys,
-            self.mk_bound_variable_kinds(&vars),
-        ))
-    }
-
     /// Expands the given impl trait type, stopping if the type is recursive.
     #[instrument(skip(self), level = "debug", ret)]
     pub fn try_expand_impl_trait_type(

compiler/rustc_next_trait_solver/src/solve/assembly/structural_traits.rs

@@ -83,7 +83,7 @@ where
             .cx()
             .coroutine_hidden_types(def_id)
             .instantiate(cx, args)
-            .map_bound(|tys| tys.to_vec())),
+            .map_bound(|bound| bound.types.to_vec())),
 
         ty::UnsafeBinder(bound_ty) => Ok(bound_ty.map_bound(|ty| vec![ty])),
 
@@ -249,7 +249,7 @@ where
             .cx()
             .coroutine_hidden_types(def_id)
             .instantiate(ecx.cx(), args)
-            .map_bound(|tys| tys.to_vec())),
+            .map_bound(|bound| bound.types.to_vec())),
     }
 }
 

compiler/rustc_trait_selection/src/traits/select/mod.rs

@@ -2866,7 +2866,7 @@ fn rebind_coroutine_witness_types<'tcx>(
     let shifted_coroutine_types =
         tcx.shift_bound_var_indices(bound_vars.len(), bound_coroutine_types.skip_binder());
     ty::Binder::bind_with_vars(
-        ty::EarlyBinder::bind(shifted_coroutine_types.to_vec()).instantiate(tcx, args),
+        ty::EarlyBinder::bind(shifted_coroutine_types.types.to_vec()).instantiate(tcx, args),
         tcx.mk_bound_variable_kinds_from_iter(
             bound_vars.iter().chain(bound_coroutine_types.bound_vars()),
         ),

compiler/rustc_traits/src/coroutine_witnesses.rs

@@ -0,0 +1,84 @@
+use rustc_hir::def_id::DefId;
+use rustc_infer::infer::TyCtxtInferExt;
+use rustc_infer::infer::canonical::query_response::make_query_region_constraints;
+use rustc_infer::traits::{Obligation, ObligationCause};
+use rustc_middle::ty::{self, Ty, TyCtxt, fold_regions};
+use rustc_trait_selection::traits::{ObligationCtxt, with_replaced_escaping_bound_vars};
+
+/// Return the set of types that should be taken into account when checking
+/// trait bounds on a coroutine's internal state. This properly replaces
+/// `ReErased` with new existential bound lifetimes.
+pub(crate) fn coroutine_hidden_types<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def_id: DefId,
+) -> ty::EarlyBinder<'tcx, ty::Binder<'tcx, ty::CoroutineWitnessTypes<TyCtxt<'tcx>>>> {
+    let coroutine_layout = tcx.mir_coroutine_witnesses(def_id);
+    let mut vars = vec![];
+    let bound_tys = tcx.mk_type_list_from_iter(
+        coroutine_layout
+            .as_ref()
+            .map_or_else(|| [].iter(), |l| l.field_tys.iter())
+            .filter(|decl| !decl.ignore_for_traits)
+            .map(|decl| {
+                let ty = fold_regions(tcx, decl.ty, |re, debruijn| {
+                    assert_eq!(re, tcx.lifetimes.re_erased);
+                    let var = ty::BoundVar::from_usize(vars.len());
+                    vars.push(ty::BoundVariableKind::Region(ty::BoundRegionKind::Anon));
+                    ty::Region::new_bound(
+                        tcx,
+                        debruijn,
+                        ty::BoundRegion { var, kind: ty::BoundRegionKind::Anon },
+                    )
+                });
+                ty
+            }),
+    );
+
+    let assumptions = compute_assumptions(tcx, def_id, bound_tys);
+
+    ty::EarlyBinder::bind(ty::Binder::bind_with_vars(
+        ty::CoroutineWitnessTypes { types: bound_tys, assumptions },
+        tcx.mk_bound_variable_kinds(&vars),
+    ))
+}
+
+fn compute_assumptions<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def_id: DefId,
+    bound_tys: &'tcx ty::List<Ty<'tcx>>,
+) -> &'tcx ty::List<ty::OutlivesPredicate<'tcx, ty::GenericArg<'tcx>>> {
+    let infcx = tcx.infer_ctxt().build(ty::TypingMode::Analysis {
+        defining_opaque_types_and_generators: ty::List::empty(),
+    });
+    with_replaced_escaping_bound_vars(&infcx, &mut vec![None], bound_tys, |bound_tys| {
+        let param_env = tcx.param_env(def_id);
+        let ocx = ObligationCtxt::new(&infcx);
+
+        ocx.register_obligations(bound_tys.iter().map(|ty| {
+            Obligation::new(
+                tcx,
+                ObligationCause::dummy(),
+                param_env,
+                ty::ClauseKind::WellFormed(ty.into()),
+            )
+        }));
+        let _errors = ocx.select_all_or_error();
+
+        // Cannot use `take_registered_region_obligations` as we may compute the response
+        // inside of a `probe` whenever we have multiple choices inside of the solver.
+        let region_obligations = infcx.take_registered_region_obligations();
+        let region_constraints = infcx.take_and_reset_region_constraints();
+        tcx.mk_outlives_from_iter(
+            make_query_region_constraints(
+                tcx,
+                region_obligations
+                    .iter()
+                    .map(|r_o| (r_o.sup_type, r_o.sub_region, r_o.origin.to_constraint_category())),
+                &region_constraints,
+            )
+            .outlives
+            .into_iter()
+            .map(|(o, _)| o),
+        )
+    })
+}

compiler/rustc_traits/src/lib.rs

@@ -5,6 +5,7 @@
 // tidy-alphabetical-end
 
 mod codegen;
+mod coroutine_witnesses;
 mod dropck_outlives;
 mod evaluate_obligation;
 mod implied_outlives_bounds;
@@ -24,4 +25,5 @@ pub fn provide(p: &mut Providers) {
     normalize_erasing_regions::provide(p);
     type_op::provide(p);
     p.codegen_select_candidate = codegen::codegen_select_candidate;
+    p.coroutine_hidden_types = coroutine_witnesses::coroutine_hidden_types;
 }

compiler/rustc_type_ir/src/interner.rs

@@ -139,6 +139,13 @@ pub trait Interner:
     type BoundRegion: Copy + Debug + Hash + Eq + BoundVarLike<Self>;
     type PlaceholderRegion: PlaceholderLike;
 
+    type RegionAssumptions: Copy
+        + Debug
+        + Hash
+        + Eq
+        + SliceLike<Item = ty::OutlivesPredicate<Self, Self::GenericArg>>
+        + TypeFoldable<Self>;
+
     // Predicates
     type ParamEnv: ParamEnv<Self>;
     type Predicate: Predicate<Self>;
@@ -210,7 +217,7 @@ pub trait Interner:
     fn coroutine_hidden_types(
         self,
         def_id: Self::DefId,
-    ) -> ty::EarlyBinder<Self, ty::Binder<Self, Self::Tys>>;
+    ) -> ty::EarlyBinder<Self, ty::Binder<Self, ty::CoroutineWitnessTypes<Self>>>;
 
     fn fn_sig(
         self,

compiler/rustc_type_ir/src/region_kind.rs

@@ -193,7 +193,7 @@ impl<I: Interner> fmt::Debug for RegionKind<I> {
 
             ReVar(vid) => write!(f, "{vid:?}"),
 
-            RePlaceholder(placeholder) => write!(f, "{placeholder:?}"),
+            RePlaceholder(placeholder) => write!(f, "'{placeholder:?}"),
 
             // Use `'{erased}` as the output instead of `'erased` so that its more obviously distinct from
             // a `ReEarlyParam` named `'erased`. Technically that would print as `'erased/#IDX` so this is

compiler/rustc_type_ir/src/ty_kind.rs

@@ -1163,3 +1163,14 @@ pub struct FnHeader<I: Interner> {
     pub safety: I::Safety,
     pub abi: I::Abi,
 }
+
+#[derive_where(Clone, Copy, Debug, PartialEq, Eq, Hash; I: Interner)]
+#[cfg_attr(
+    feature = "nightly",
+    derive(Encodable_NoContext, Decodable_NoContext, HashStable_NoContext)
+)]
+#[derive(TypeVisitable_Generic, TypeFoldable_Generic, Lift_Generic)]
+pub struct CoroutineWitnessTypes<I: Interner> {
+    pub types: I::Tys,
+    pub assumptions: I::RegionAssumptions,
+}