eagerly normalize during generalization

compiler/rustc_borrowck/src/type_check/relate_tys.rs

@@ -616,4 +616,9 @@ fn register_alias_relate_predicate(&mut self, a: Ty<'tcx>, b: Ty<'tcx>) {
             }
         })]);
     }
+
+    fn try_eagerly_normalize_alias(&mut self, _alias: ty::AliasTy<'tcx>) -> Ty<'tcx> {
+        // Past hir typeck, so we don't have to worry about type inference anymore.
+        self.type_checker.infcx.next_ty_var(self.span())
+    }
 }

compiler/rustc_infer/src/infer/at.rs

@@ -140,6 +140,8 @@ pub fn sup<T>(
                 ty::Contravariant,
                 actual,
                 self.cause.span,
+                // TODO: should normalize
+                &mut |_alias| self.infcx.next_ty_var(self.cause.span),
             )
             .map(|goals| self.goals_to_obligations(goals))
         } else {
@@ -173,6 +175,8 @@ pub fn sub<T>(
                 ty::Covariant,
                 actual,
                 self.cause.span,
+                // TODO: should normalize
+                &mut |_alias| self.infcx.next_ty_var(self.cause.span),
             )
             .map(|goals| self.goals_to_obligations(goals))
         } else {
@@ -225,6 +229,8 @@ pub fn eq_trace<T>(
                 ty::Invariant,
                 actual,
                 self.cause.span,
+                // TODO: should normalize
+                &mut |_alias| self.infcx.next_ty_var(self.cause.span),
             )
             .map(|goals| self.goals_to_obligations(goals))
         } else {

compiler/rustc_infer/src/infer/relate/generalize.rs

@@ -76,6 +76,7 @@ pub fn instantiate_ty_var<R: PredicateEmittingRelation<InferCtxt<'tcx>>>(
             target_vid,
             instantiation_variance,
             source_ty,
+            &mut |alias| relation.try_eagerly_normalize_alias(alias),
         )?;
 
         // Constrain `b_vid` to the generalized type `generalized_ty`.
@@ -210,6 +211,7 @@ pub(crate) fn instantiate_const_var<R: PredicateEmittingRelation<InferCtxt<'tcx>
             target_vid,
             ty::Invariant,
             source_ct,
+            &mut |alias| relation.try_eagerly_normalize_alias(alias),
         )?;
 
         debug_assert!(!generalized_ct.is_ct_infer());
@@ -249,6 +251,7 @@ fn generalize<T: Into<Term<'tcx>> + Relate<TyCtxt<'tcx>>>(
         target_vid: impl Into<TermVid>,
         ambient_variance: ty::Variance,
         source_term: T,
+        normalize: &mut dyn FnMut(ty::AliasTy<'tcx>) -> Ty<'tcx>,
     ) -> RelateResult<'tcx, Generalization<T>> {
         assert!(!source_term.has_escaping_bound_vars());
         let (for_universe, root_vid) = match target_vid.into() {
@@ -269,8 +272,9 @@ fn generalize<T: Into<Term<'tcx>> + Relate<TyCtxt<'tcx>>>(
             for_universe,
             root_term: source_term.into(),
             ambient_variance,
-            in_alias: false,
+            state: GeneralizationState::Default,
             cache: Default::default(),
+            normalize,
         };
 
         let value_may_be_infer = generalizer.relate(source_term, source_term)?;
@@ -317,6 +321,13 @@ fn visit_region(&mut self, r: ty::Region<'tcx>) {
     }
 }
 
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+enum GeneralizationState {
+    Default,
+    InAlias,
+    InNormalizedAlias,
+}
+
 /// The "generalizer" is used when handling inference variables.
 ///
 /// The basic strategy for handling a constraint like `?A <: B` is to
@@ -361,9 +372,14 @@ struct Generalizer<'me, 'tcx> {
     /// This is necessary to correctly handle
     /// `<T as Bar<<?0 as Foo>::Assoc>::Assoc == ?0`. This equality can
     /// hold by either normalizing the outer or the inner associated type.
-    in_alias: bool,
+    // TODO: update doc comment
+    state: GeneralizationState,
 
-    cache: SsoHashMap<(Ty<'tcx>, ty::Variance, bool), Ty<'tcx>>,
+    cache: SsoHashMap<(Ty<'tcx>, ty::Variance, GeneralizationState), Ty<'tcx>>,
+
+    /// Normalize an alias in the trait solver.
+    /// If normalization fails, a fresh infer var is returned.
+    normalize: &'me mut dyn FnMut(ty::AliasTy<'tcx>) -> Ty<'tcx>,
 }
 
 impl<'tcx> Generalizer<'_, 'tcx> {
@@ -409,17 +425,34 @@ fn generalize_alias_ty(
         // with inference variables can cause incorrect ambiguity.
         //
         // cc trait-system-refactor-initiative#110
-        if self.infcx.next_trait_solver() && !alias.has_escaping_bound_vars() && !self.in_alias {
-            return Ok(self.next_ty_var_for_alias());
+        if self.infcx.next_trait_solver()
+            && !alias.has_escaping_bound_vars()
+            && match self.state {
+                GeneralizationState::Default => true,
+                GeneralizationState::InAlias => false,
+                // When generalizing an alias after normalizing,
+                // the outer alias should be treated as rigid and we shouldn't try generalizing it again.
+                // If we recursively find more aliases, the state should have been set back to InAlias.
+                GeneralizationState::InNormalizedAlias => unreachable!(),
+            }
+        {
+            let normalized_alias = (self.normalize)(alias);
+
+            self.state = GeneralizationState::InNormalizedAlias;
+            // recursively generalize, treat the outer alias as rigid to avoid infinite recursion
+            let res = self.relate(normalized_alias, normalized_alias);
+
+            // only one way to get here
+            self.state = GeneralizationState::Default;
+
+            return res;
         }
 
-        let is_nested_alias = mem::replace(&mut self.in_alias, true);
+        let previous_state = mem::replace(&mut self.state, GeneralizationState::InAlias);
         let result = match self.relate(alias, alias) {
             Ok(alias) => Ok(alias.to_ty(self.cx())),
-            Err(e) => {
-                if is_nested_alias {
-                    return Err(e);
-                } else {
+            Err(e) => match previous_state {
+                GeneralizationState::Default => {
                     let mut visitor = MaxUniverse::new();
                     alias.visit_with(&mut visitor);
                     let infer_replacement_is_complete =
@@ -432,9 +465,14 @@ fn generalize_alias_ty(
                     debug!("generalization failure in alias");
                     Ok(self.next_ty_var_for_alias())
                 }
-            }
+                GeneralizationState::InAlias => return Err(e),
+                // When generalizing an alias after normalizing,
+                // the outer alias should be treated as rigid and we shouldn't try generalizing it again.
+                // If we recursively find more aliases, the state should have been set back to InAlias.
+                GeneralizationState::InNormalizedAlias => unreachable!(),
+            },
         };
-        self.in_alias = is_nested_alias;
+        self.state = previous_state;
         result
     }
 }
@@ -488,7 +526,7 @@ fn relate_with_variance<T: Relate<TyCtxt<'tcx>>>(
     fn tys(&mut self, t: Ty<'tcx>, t2: Ty<'tcx>) -> RelateResult<'tcx, Ty<'tcx>> {
         assert_eq!(t, t2); // we are misusing TypeRelation here; both LHS and RHS ought to be ==
 
-        if let Some(&result) = self.cache.get(&(t, self.ambient_variance, self.in_alias)) {
+        if let Some(&result) = self.cache.get(&(t, self.ambient_variance, self.state)) {
             return Ok(result);
         }
 
@@ -553,9 +591,17 @@ fn tys(&mut self, t: Ty<'tcx>, t2: Ty<'tcx>) -> RelateResult<'tcx, Ty<'tcx>> {
                             // cc trait-system-refactor-initiative#108
                             if self.infcx.next_trait_solver()
                                 && !matches!(self.infcx.typing_mode(), TypingMode::Coherence)
-                                && self.in_alias
                             {
-                                inner.type_variables().equate(vid, new_var_id);
+                                match self.state {
+                                    GeneralizationState::InAlias => {
+                                        inner.type_variables().equate(vid, new_var_id);
+                                    }
+                                    GeneralizationState::Default
+                                    | GeneralizationState::InNormalizedAlias => {}
+                                }
+                                GeneralizerState::Default
+                                | GeneralizerState::ShallowStructurallyRelateAliases
+                                | GeneralizerState::StructurallyRelateAliases => {}
                             }
 
                             debug!("replacing original vid={:?} with new={:?}", vid, new_var_id);
@@ -585,14 +631,25 @@ fn tys(&mut self, t: Ty<'tcx>, t2: Ty<'tcx>) -> RelateResult<'tcx, Ty<'tcx>> {
             }
 
             ty::Alias(_, data) => match self.structurally_relate_aliases {
-                StructurallyRelateAliases::No => self.generalize_alias_ty(data),
+                StructurallyRelateAliases::No => match self.state {
+                    GeneralizationState::Default | GeneralizationState::InAlias => {
+                        self.generalize_alias_ty(data)
+                    }
+                    GeneralizationState::InNormalizedAlias => {
+                        // We can switch back to default, we've treated one layer as rigid by doing this operation.
+                        self.state = GeneralizationState::Default;
+                        let res = relate::structurally_relate_tys(self, t, t);
+                        self.state = GeneralizationState::InNormalizedAlias;
+                        res
+                    }
+                },
                 StructurallyRelateAliases::Yes => relate::structurally_relate_tys(self, t, t),
             },
 
             _ => relate::structurally_relate_tys(self, t, t),
         }?;
 
-        self.cache.insert((t, self.ambient_variance, self.in_alias), g);
+        self.cache.insert((t, self.ambient_variance, self.state), g);
         Ok(g)
     }
 
@@ -683,9 +740,17 @@ fn consts(
                             // for more details.
                             if self.infcx.next_trait_solver()
                                 && !matches!(self.infcx.typing_mode(), TypingMode::Coherence)
-                                && self.in_alias
                             {
-                                variable_table.union(vid, new_var_id);
+                                match self.state {
+                                    GeneralizationState::InAlias => {
+                                        variable_table.union(vid, new_var_id);
+                                    }
+                                    GeneralizationState::Default
+                                    | GeneralizationState::InNormalizedAlias => {}
+                                }
+                                GeneralizerState::Default
+                                | GeneralizerState::ShallowStructurallyRelateAliases
+                                | GeneralizerState::StructurallyRelateAliases => {}
                             }
                             Ok(ty::Const::new_var(self.cx(), new_var_id))
                         }

compiler/rustc_infer/src/infer/relate/lattice.rs

@@ -299,4 +299,9 @@ fn register_alias_relate_predicate(&mut self, a: Ty<'tcx>, b: Ty<'tcx>) {
             ty::AliasRelationDirection::Equate,
         ))]);
     }
+
+    fn try_eagerly_normalize_alias(&mut self, _alias: ty::AliasTy<'tcx>) -> Ty<'tcx> {
+        // TODO: this should actually normalize
+        self.infcx.next_ty_var(self.span())
+    }
 }

compiler/rustc_infer/src/infer/relate/type_relating.rs

@@ -396,4 +396,13 @@ fn register_alias_relate_predicate(&mut self, a: Ty<'tcx>, b: Ty<'tcx>) {
             }
         })]);
     }
+
+    fn try_eagerly_normalize_alias(
+        &mut self,
+        _alias: rustc_type_ir::AliasTy<<InferCtxt<'tcx> as rustc_type_ir::InferCtxtLike>::Interner>,
+    ) -> <<InferCtxt<'tcx> as rustc_type_ir::InferCtxtLike>::Interner as rustc_type_ir::Interner>::Ty
+    {
+        // We only try to eagerly normalize aliases if we're using the new solver.
+        unreachable!()
+    }
 }

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

@@ -976,7 +976,12 @@ fn try_eagerly_replace_alias(
         let replacement = self.ecx.instantiate_binder_with_infer(*replacement);
         self.nested.extend(
             self.ecx
-                .eq_and_get_goals(self.param_env, alias_term, replacement.projection_term)
+                .relate_and_get_goals(
+                    self.param_env,
+                    alias_term,
+                    ty::Invariant,
+                    replacement.projection_term,
+                )
                 .expect("expected to be able to unify goal projection with dyn's projection"),
         );
 

compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs

@@ -408,7 +408,7 @@ pub(super) fn ignore_candidate_head_usages(&mut self, usages: CandidateHeadUsage
 
     /// Recursively evaluates `goal`, returning whether any inference vars have
     /// been constrained and the certainty of the result.
-    fn evaluate_goal(
+    pub(super) fn evaluate_goal(
         &mut self,
         source: GoalSource,
         goal: Goal<I, I::Predicate>,
@@ -1018,7 +1018,8 @@ pub(super) fn relate<T: Relate<I>>(
         variance: ty::Variance,
         rhs: T,
     ) -> Result<(), NoSolution> {
-        let goals = self.delegate.relate(param_env, lhs, variance, rhs, self.origin_span)?;
+        let goals = self.relate_and_get_goals(param_env, lhs, variance, rhs)?;
+
         for &goal in goals.iter() {
             let source = match goal.predicate.kind().skip_binder() {
                 ty::PredicateKind::Subtype { .. } | ty::PredicateKind::AliasRelate(..) => {
@@ -1039,13 +1040,37 @@ pub(super) fn relate<T: Relate<I>>(
     /// If possible, try using `eq` instead which automatically handles nested
     /// goals correctly.
     #[instrument(level = "trace", skip(self, param_env), ret)]
-    pub(super) fn eq_and_get_goals<T: Relate<I>>(
-        &self,
+    pub(super) fn relate_and_get_goals<T: Relate<I>>(
+        &mut self,
         param_env: I::ParamEnv,
         lhs: T,
+        variance: ty::Variance,
         rhs: T,
     ) -> Result<Vec<Goal<I, I::Predicate>>, NoSolution> {
-        Ok(self.delegate.relate(param_env, lhs, ty::Variance::Invariant, rhs, self.origin_span)?)
+        let cx = self.cx();
+        let delegate = self.delegate;
+        let origin_span = self.origin_span;
+
+        let mut normalize = |alias: ty::AliasTy<I>| {
+            let inference_var = self.next_ty_infer();
+
+            let goal = Goal::new(
+                cx,
+                param_env,
+                ty::PredicateKind::AliasRelate(
+                    alias.to_ty(cx).into(),
+                    inference_var.into(),
+                    ty::AliasRelationDirection::Equate,
+                ),
+            );
+
+            // Ignore the result. If we can't eagerly normalize, returning the inference variable is enough.
+            let _ = self.evaluate_goal(GoalSource::TypeRelating, goal, None);
+
+            self.resolve_vars_if_possible(inference_var)
+        };
+
+        Ok(delegate.relate(param_env, lhs, variance, rhs, origin_span, &mut normalize)?)
     }
 
     pub(super) fn instantiate_binder_with_infer<T: TypeFoldable<I> + Copy>(

compiler/rustc_type_ir/src/relate/combine.rs

@@ -40,6 +40,8 @@ fn register_predicates(
 
     /// Register `AliasRelate` obligation(s) that both types must be related to each other.
     fn register_alias_relate_predicate(&mut self, a: I::Ty, b: I::Ty);
+
+    fn try_eagerly_normalize_alias(&mut self, alias: ty::AliasTy<I>) -> I::Ty;
 }
 
 pub fn super_combine_tys<Infcx, I, R>(

compiler/rustc_type_ir/src/relate/solver_relating.rs

@@ -15,6 +15,7 @@ fn relate<T: Relate<Self::Interner>>(
         variance: ty::Variance,
         rhs: T,
         span: <Self::Interner as Interner>::Span,
+        normalize: &mut dyn FnMut(ty::AliasTy<Self::Interner>) -> <Self::Interner as Interner>::Ty,
     ) -> Result<
         Vec<Goal<Self::Interner, <Self::Interner as Interner>::Predicate>>,
         TypeError<Self::Interner>,
@@ -32,40 +33,46 @@ fn eq_structurally_relating_aliases<T: Relate<Self::Interner>>(
     >;
 }
 
-impl<Infcx: InferCtxtLike> RelateExt for Infcx {
-    fn relate<T: Relate<Self::Interner>>(
+impl<I: Interner, Infcx: InferCtxtLike<Interner = I>> RelateExt for Infcx {
+    fn relate<T: Relate<I>>(
         &self,
-        param_env: <Self::Interner as Interner>::ParamEnv,
+        param_env: I::ParamEnv,
         lhs: T,
         variance: ty::Variance,
         rhs: T,
-        span: <Self::Interner as Interner>::Span,
-    ) -> Result<
-        Vec<Goal<Self::Interner, <Self::Interner as Interner>::Predicate>>,
-        TypeError<Self::Interner>,
-    > {
-        let mut relate =
-            SolverRelating::new(self, StructurallyRelateAliases::No, variance, param_env, span);
+        span: I::Span,
+        normalize: &mut dyn FnMut(ty::AliasTy<I>) -> I::Ty,
+    ) -> Result<Vec<Goal<I, I::Predicate>>, TypeError<I>> {
+        let mut relate = SolverRelating::new(
+            self,
+            StructurallyRelateAliases::No,
+            variance,
+            param_env,
+            span,
+            normalize,
+        );
         relate.relate(lhs, rhs)?;
         Ok(relate.goals)
     }
 
-    fn eq_structurally_relating_aliases<T: Relate<Self::Interner>>(
+    fn eq_structurally_relating_aliases<T: Relate<I>>(
         &self,
-        param_env: <Self::Interner as Interner>::ParamEnv,
+        param_env: I::ParamEnv,
         lhs: T,
         rhs: T,
-        span: <Self::Interner as Interner>::Span,
-    ) -> Result<
-        Vec<Goal<Self::Interner, <Self::Interner as Interner>::Predicate>>,
-        TypeError<Self::Interner>,
-    > {
+        span: I::Span,
+    ) -> Result<Vec<Goal<I, I::Predicate>>, TypeError<I>> {
+        // Structurally relating, we treat aliases as rigid,
+        // so we shouldn't ever try to normalize them.
+        let mut normalize_unreachable = |_alias| unreachable!();
+
         let mut relate = SolverRelating::new(
             self,
             StructurallyRelateAliases::Yes,
             ty::Invariant,
             param_env,
             span,
+            &mut normalize_unreachable,
         );
         relate.relate(lhs, rhs)?;
         Ok(relate.goals)
@@ -75,12 +82,14 @@ fn eq_structurally_relating_aliases<T: Relate<Self::Interner>>(
 /// Enforce that `a` is equal to or a subtype of `b`.
 pub struct SolverRelating<'infcx, Infcx, I: Interner> {
     infcx: &'infcx Infcx,
+
     // Immutable fields.
     structurally_relate_aliases: StructurallyRelateAliases,
     param_env: I::ParamEnv,
     span: I::Span,
     // Mutable fields.
     ambient_variance: ty::Variance,
+    normalize: &'infcx mut dyn FnMut(ty::AliasTy<I>) -> I::Ty,
     goals: Vec<Goal<I, I::Predicate>>,
     /// The cache only tracks the `ambient_variance` as it's the
     /// only field which is mutable and which meaningfully changes
@@ -118,12 +127,14 @@ pub fn new(
         ambient_variance: ty::Variance,
         param_env: I::ParamEnv,
         span: I::Span,
+        normalize: &'infcx mut dyn FnMut(ty::AliasTy<I>) -> I::Ty,
     ) -> Self {
         SolverRelating {
             infcx,
             structurally_relate_aliases,
             span,
             ambient_variance,
+            normalize,
             param_env,
             goals: vec![],
             cache: Default::default(),
@@ -406,4 +417,8 @@ fn register_alias_relate_predicate(&mut self, a: I::Ty, b: I::Ty) {
             }
         })]);
     }
+
+    fn try_eagerly_normalize_alias(&mut self, alias: ty::AliasTy<I>) -> I::Ty {
+        (self.normalize)(alias)
+    }
 }