Rollup merge of #121344 - fmease:lta-constr-by-input, r=oli-obk

Expand weak alias types before collecting constrained/referenced late bound regions + refactorings

Fixes #114220.
Follow-up to #120780.

r? `@oli-obk`

compiler/rustc_middle/src/arena.rs

@@ -47,7 +47,7 @@ macro_rules! arena_types {
                         rustc_middle::traits::query::DropckOutlivesResult<'tcx>
                     >
                 >,
-            [] normalize_projection_ty:
+            [] normalize_canonicalized_projection_ty:
                 rustc_middle::infer::canonical::Canonical<'tcx,
                     rustc_middle::infer::canonical::QueryResponse<'tcx,
                         rustc_middle::traits::query::NormalizationResult<'tcx>

compiler/rustc_middle/src/query/mod.rs

@@ -31,7 +31,7 @@ use crate::query::plumbing::{
 };
 use crate::thir;
 use crate::traits::query::{
-    CanonicalPredicateGoal, CanonicalProjectionGoal, CanonicalTyGoal,
+    CanonicalAliasGoal, CanonicalPredicateGoal, CanonicalTyGoal,
     CanonicalTypeOpAscribeUserTypeGoal, CanonicalTypeOpEqGoal, CanonicalTypeOpNormalizeGoal,
     CanonicalTypeOpProvePredicateGoal, CanonicalTypeOpSubtypeGoal, NoSolution,
 };
@@ -1938,9 +1938,13 @@ rustc_queries! {
         arena_cache
     }
 
-    /// Do not call this query directly: invoke `normalize` instead.
-    query normalize_projection_ty(
-        goal: CanonicalProjectionGoal<'tcx>
+    /// <div class="warning">
+    ///
+    /// Do not call this query directly: Invoke `normalize` instead.
+    ///
+    /// </div>
+    query normalize_canonicalized_projection_ty(
+        goal: CanonicalAliasGoal<'tcx>
     ) -> Result<
         &'tcx Canonical<'tcx, canonical::QueryResponse<'tcx, NormalizationResult<'tcx>>>,
         NoSolution,
@@ -1948,9 +1952,13 @@ rustc_queries! {
         desc { "normalizing `{}`", goal.value.value }
     }
 
-    /// Do not call this query directly: invoke `normalize` instead.
-    query normalize_weak_ty(
-        goal: CanonicalProjectionGoal<'tcx>
+    /// <div class="warning">
+    ///
+    /// Do not call this query directly: Invoke `normalize` instead.
+    ///
+    /// </div>
+    query normalize_canonicalized_weak_ty(
+        goal: CanonicalAliasGoal<'tcx>
     ) -> Result<
         &'tcx Canonical<'tcx, canonical::QueryResponse<'tcx, NormalizationResult<'tcx>>>,
         NoSolution,
@@ -1958,9 +1966,13 @@ rustc_queries! {
         desc { "normalizing `{}`", goal.value.value }
     }
 
-    /// Do not call this query directly: invoke `normalize` instead.
-    query normalize_inherent_projection_ty(
-        goal: CanonicalProjectionGoal<'tcx>
+    /// <div class="warning">
+    ///
+    /// Do not call this query directly: Invoke `normalize` instead.
+    ///
+    /// </div>
+    query normalize_canonicalized_inherent_projection_ty(
+        goal: CanonicalAliasGoal<'tcx>
     ) -> Result<
         &'tcx Canonical<'tcx, canonical::QueryResponse<'tcx, NormalizationResult<'tcx>>>,
         NoSolution,

compiler/rustc_middle/src/traits/query.rs

@@ -67,7 +67,7 @@ pub mod type_op {
     }
 }
 
-pub type CanonicalProjectionGoal<'tcx> = Canonical<'tcx, ty::ParamEnvAnd<'tcx, ty::AliasTy<'tcx>>>;
+pub type CanonicalAliasGoal<'tcx> = Canonical<'tcx, ty::ParamEnvAnd<'tcx, ty::AliasTy<'tcx>>>;
 
 pub type CanonicalTyGoal<'tcx> = Canonical<'tcx, ty::ParamEnvAnd<'tcx, Ty<'tcx>>>;
 
@@ -177,10 +177,10 @@ pub struct MethodAutoderefBadTy<'tcx> {
     pub ty: Canonical<'tcx, QueryResponse<'tcx, Ty<'tcx>>>,
 }
 
-/// Result from the `normalize_projection_ty` query.
+/// Result of the `normalize_canonicalized_{{,inherent_}projection,weak}_ty` queries.
 #[derive(Clone, Debug, HashStable, TypeFoldable, TypeVisitable)]
 pub struct NormalizationResult<'tcx> {
-    /// Result of normalization.
+    /// Result of the normalization.
     pub normalized_ty: Ty<'tcx>,
 }
 

compiler/rustc_middle/src/ty/flags.rs

@@ -181,9 +181,10 @@ impl FlagComputation {
 
             &ty::Alias(kind, data) => {
                 self.add_flags(match kind {
-                    ty::Weak | ty::Projection => TypeFlags::HAS_TY_PROJECTION,
-                    ty::Inherent => TypeFlags::HAS_TY_INHERENT,
+                    ty::Projection => TypeFlags::HAS_TY_PROJECTION,
+                    ty::Weak => TypeFlags::HAS_TY_WEAK,
                     ty::Opaque => TypeFlags::HAS_TY_OPAQUE,
+                    ty::Inherent => TypeFlags::HAS_TY_INHERENT,
                 });
 
                 self.add_alias_ty(data);

compiler/rustc_middle/src/ty/util.rs

@@ -11,6 +11,7 @@ use crate::ty::{GenericArgKind, GenericArgsRef};
 use rustc_apfloat::Float as _;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::stable_hasher::{Hash128, HashStable, StableHasher};
+use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_errors::ErrorGuaranteed;
 use rustc_hir as hir;
 use rustc_hir::def::{CtorOf, DefKind, Res};
@@ -867,6 +868,63 @@ impl<'tcx> TyCtxt<'tcx> {
 
         self.mk_args_from_iter(args.into_iter().map(|arg| arg.into()).chain(opt_const_param))
     }
+
+    /// Expand any [weak alias types][weak] contained within the given `value`.
+    ///
+    /// This should be used over other normalization routines in situations where
+    /// it's important not to normalize other alias types and where the predicates
+    /// on the corresponding type alias shouldn't be taken into consideration.
+    ///
+    /// Whenever possible **prefer not to use this function**! Instead, use standard
+    /// normalization routines or if feasible don't normalize at all.
+    ///
+    /// This function comes in handy if you want to mimic the behavior of eager
+    /// type alias expansion in a localized manner.
+    ///
+    /// <div class="warning">
+    /// This delays a bug on overflow! Therefore you need to be certain that the
+    /// contained types get fully normalized at a later stage. Note that even on
+    /// overflow all well-behaved weak alias types get expanded correctly, so the
+    /// result is still useful.
+    /// </div>
+    ///
+    /// [weak]: ty::Weak
+    pub fn expand_weak_alias_tys<T: TypeFoldable<TyCtxt<'tcx>>>(self, value: T) -> T {
+        value.fold_with(&mut WeakAliasTypeExpander { tcx: self, depth: 0 })
+    }
+
+    /// Peel off all [weak alias types] in this type until there are none left.
+    ///
+    /// This only expands weak alias types in “head” / outermost positions. It can
+    /// be used over [expand_weak_alias_tys] as an optimization in situations where
+    /// one only really cares about the *kind* of the final aliased type but not
+    /// the types the other constituent types alias.
+    ///
+    /// <div class="warning">
+    /// This delays a bug on overflow! Therefore you need to be certain that the
+    /// type gets fully normalized at a later stage.
+    /// </div>
+    ///
+    /// [weak]: ty::Weak
+    /// [expand_weak_alias_tys]: Self::expand_weak_alias_tys
+    pub fn peel_off_weak_alias_tys(self, mut ty: Ty<'tcx>) -> Ty<'tcx> {
+        let ty::Alias(ty::Weak, _) = ty.kind() else { return ty };
+
+        let limit = self.recursion_limit();
+        let mut depth = 0;
+
+        while let ty::Alias(ty::Weak, alias) = ty.kind() {
+            if !limit.value_within_limit(depth) {
+                let guar = self.dcx().delayed_bug("overflow expanding weak alias type");
+                return Ty::new_error(self, guar);
+            }
+
+            ty = self.type_of(alias.def_id).instantiate(self, alias.args);
+            depth += 1;
+        }
+
+        ty
+    }
 }
 
 struct OpaqueTypeExpander<'tcx> {
@@ -1002,6 +1060,42 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for OpaqueTypeExpander<'tcx> {
     }
 }
 
+struct WeakAliasTypeExpander<'tcx> {
+    tcx: TyCtxt<'tcx>,
+    depth: usize,
+}
+
+impl<'tcx> TypeFolder<TyCtxt<'tcx>> for WeakAliasTypeExpander<'tcx> {
+    fn interner(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
+
+    fn fold_ty(&mut self, ty: Ty<'tcx>) -> Ty<'tcx> {
+        if !ty.has_type_flags(ty::TypeFlags::HAS_TY_WEAK) {
+            return ty;
+        }
+        let ty::Alias(ty::Weak, alias) = ty.kind() else {
+            return ty.super_fold_with(self);
+        };
+        if !self.tcx.recursion_limit().value_within_limit(self.depth) {
+            let guar = self.tcx.dcx().delayed_bug("overflow expanding weak alias type");
+            return Ty::new_error(self.tcx, guar);
+        }
+
+        self.depth += 1;
+        ensure_sufficient_stack(|| {
+            self.tcx.type_of(alias.def_id).instantiate(self.tcx, alias.args).fold_with(self)
+        })
+    }
+
+    fn fold_const(&mut self, ct: ty::Const<'tcx>) -> ty::Const<'tcx> {
+        if !ct.ty().has_type_flags(ty::TypeFlags::HAS_TY_WEAK) {
+            return ct;
+        }
+        ct.super_fold_with(self)
+    }
+}
+
 impl<'tcx> Ty<'tcx> {
     /// Returns the `Size` for primitive types (bool, uint, int, char, float).
     pub fn primitive_size(self, tcx: TyCtxt<'tcx>) -> Size {

compiler/rustc_middle/src/ty/visit.rs

@@ -2,6 +2,7 @@ use crate::ty::{self, Binder, Ty, TyCtxt, TypeFlags};
 
 use rustc_data_structures::fx::FxHashSet;
 use rustc_data_structures::sso::SsoHashSet;
+use rustc_type_ir::fold::TypeFoldable;
 use std::ops::ControlFlow;
 
 pub use rustc_type_ir::visit::{TypeSuperVisitable, TypeVisitable, TypeVisitableExt, TypeVisitor};
@@ -109,10 +110,10 @@ impl<'tcx> TyCtxt<'tcx> {
     /// variables will also be equated.
     pub fn collect_constrained_late_bound_regions<T>(
         self,
-        value: &Binder<'tcx, T>,
+        value: Binder<'tcx, T>,
     ) -> FxHashSet<ty::BoundRegionKind>
     where
-        T: TypeVisitable<TyCtxt<'tcx>>,
+        T: TypeFoldable<TyCtxt<'tcx>>,
     {
         self.collect_late_bound_regions(value, true)
     }
@@ -120,24 +121,26 @@ impl<'tcx> TyCtxt<'tcx> {
     /// Returns a set of all late-bound regions that appear in `value` anywhere.
     pub fn collect_referenced_late_bound_regions<T>(
         self,
-        value: &Binder<'tcx, T>,
+        value: Binder<'tcx, T>,
     ) -> FxHashSet<ty::BoundRegionKind>
     where
-        T: TypeVisitable<TyCtxt<'tcx>>,
+        T: TypeFoldable<TyCtxt<'tcx>>,
     {
         self.collect_late_bound_regions(value, false)
     }
 
     fn collect_late_bound_regions<T>(
         self,
-        value: &Binder<'tcx, T>,
-        just_constraint: bool,
+        value: Binder<'tcx, T>,
+        just_constrained: bool,
     ) -> FxHashSet<ty::BoundRegionKind>
     where
-        T: TypeVisitable<TyCtxt<'tcx>>,
+        T: TypeFoldable<TyCtxt<'tcx>>,
     {
-        let mut collector = LateBoundRegionsCollector::new(just_constraint);
-        let result = value.as_ref().skip_binder().visit_with(&mut collector);
+        let mut collector = LateBoundRegionsCollector::new(just_constrained);
+        let value = value.skip_binder();
+        let value = if just_constrained { self.expand_weak_alias_tys(value) } else { value };
+        let result = value.visit_with(&mut collector);
         assert!(result.is_continue()); // should never have stopped early
         collector.regions
     }
@@ -258,11 +261,7 @@ struct LateBoundRegionsCollector {
 
 impl LateBoundRegionsCollector {
     fn new(just_constrained: bool) -> Self {
-        LateBoundRegionsCollector {
-            current_index: ty::INNERMOST,
-            regions: Default::default(),
-            just_constrained,
-        }
+        Self { current_index: ty::INNERMOST, regions: Default::default(), just_constrained }
     }
 }
 
@@ -278,12 +277,16 @@ impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for LateBoundRegionsCollector {
     }
 
     fn visit_ty(&mut self, t: Ty<'tcx>) -> ControlFlow<Self::BreakTy> {
-        // if we are only looking for "constrained" region, we have to
-        // ignore the inputs to a projection, as they may not appear
-        // in the normalized form
         if self.just_constrained {
-            if let ty::Alias(..) = t.kind() {
-                return ControlFlow::Continue(());
+            match t.kind() {
+                // If we are only looking for "constrained" regions, we have to ignore the
+                // inputs to a projection as they may not appear in the normalized form.
+                ty::Alias(ty::Projection | ty::Inherent | ty::Opaque, _) => {
+                    return ControlFlow::Continue(());
+                }
+                // All weak alias types should've been expanded beforehand.
+                ty::Alias(ty::Weak, _) => bug!("unexpected weak alias type"),
+                _ => {}
             }
         }