setup canonical normalization for const norm

Author: BoxyUwU

compiler/rustc_middle/src/arena.rs

@@ -44,7 +44,7 @@ macro_rules! arena_types {
                         rustc_middle::traits::query::DropckOutlivesResult<'tcx>
                     >
                 >,
-            [] normalize_canonicalized_projection_ty:
+            [] normalize_canonicalized_projection:
                 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

@@ -2411,7 +2411,7 @@ rustc_queries! {
     /// Do not call this query directly: Invoke `normalize` instead.
     ///
     /// </div>
-    query normalize_canonicalized_projection_ty(
+    query normalize_canonicalized_projection(
         goal: CanonicalAliasGoal<'tcx>
     ) -> Result<
         &'tcx Canonical<'tcx, canonical::QueryResponse<'tcx, NormalizationResult<'tcx>>>,
@@ -2439,7 +2439,7 @@ rustc_queries! {
     /// Do not call this query directly: Invoke `normalize` instead.
     ///
     /// </div>
-    query normalize_canonicalized_inherent_projection_ty(
+    query normalize_canonicalized_inherent_projection(
         goal: CanonicalAliasGoal<'tcx>
     ) -> Result<
         &'tcx Canonical<'tcx, canonical::QueryResponse<'tcx, NormalizationResult<'tcx>>>,

compiler/rustc_middle/src/traits/query.rs

@@ -66,7 +66,7 @@ pub mod type_op {
 }
 
 pub type CanonicalAliasGoal<'tcx> =
-    CanonicalQueryInput<'tcx, ty::ParamEnvAnd<'tcx, ty::AliasTy<'tcx>>>;
+    CanonicalQueryInput<'tcx, ty::ParamEnvAnd<'tcx, ty::AliasTerm<'tcx>>>;
 
 pub type CanonicalMethodAutoderefStepsGoal<'tcx> =
     CanonicalQueryInput<'tcx, ty::ParamEnvAnd<'tcx, MethodAutoderefSteps<'tcx>>>;
@@ -192,11 +192,11 @@ pub struct MethodAutoderefBadTy<'tcx> {
     pub ty: Canonical<'tcx, QueryResponse<'tcx, Ty<'tcx>>>,
 }
 
-/// Result of the `normalize_canonicalized_{{,inherent_}projection,free}_ty` queries.
+/// Result of the `normalize_canonicalized_{{,inherent_}projection,free}` queries.
 #[derive(Clone, Debug, HashStable, TypeFoldable, TypeVisitable)]
 pub struct NormalizationResult<'tcx> {
     /// Result of the normalization.
-    pub normalized_ty: Ty<'tcx>,
+    pub normalized_term: ty::Term<'tcx>,
 }
 
 /// Outlives bounds are relationships between generic parameters,

compiler/rustc_trait_selection/src/traits/query/normalize.rs

@@ -1,9 +1,10 @@
 //! Code for the 'normalization' query. This consists of a wrapper
 //! which folds deeply, invoking the underlying
-//! `normalize_canonicalized_projection_ty` query when it encounters projections.
+//! `normalize_canonicalized_projection` query when it encounters projections.
 
 use rustc_data_structures::sso::SsoHashMap;
 use rustc_data_structures::stack::ensure_sufficient_stack;
+use rustc_hir::def::DefKind;
 use rustc_infer::traits::PredicateObligations;
 use rustc_macros::extension;
 pub use rustc_middle::traits::query::NormalizationResult;
@@ -253,76 +254,9 @@ impl<'a, 'tcx> FallibleTypeFolder<TyCtxt<'tcx>> for QueryNormalizer<'a, 'tcx> {
                 }
             }
 
-            ty::Projection | ty::Inherent | ty::Free => {
-                // See note in `rustc_trait_selection::traits::project`
-
-                let infcx = self.infcx;
-                let tcx = infcx.tcx;
-                // Just an optimization: When we don't have escaping bound vars,
-                // we don't need to replace them with placeholders.
-                let (data, maps) = if data.has_escaping_bound_vars() {
-                    let (data, mapped_regions, mapped_types, mapped_consts) =
-                        BoundVarReplacer::replace_bound_vars(infcx, &mut self.universes, data);
-                    (data, Some((mapped_regions, mapped_types, mapped_consts)))
-                } else {
-                    (data, None)
-                };
-                let data = data.try_fold_with(self)?;
-
-                let mut orig_values = OriginalQueryValues::default();
-                let c_data = infcx.canonicalize_query(self.param_env.and(data), &mut orig_values);
-                debug!("QueryNormalizer: c_data = {:#?}", c_data);
-                debug!("QueryNormalizer: orig_values = {:#?}", orig_values);
-                let result = match kind {
-                    ty::Projection => tcx.normalize_canonicalized_projection_ty(c_data),
-                    ty::Free => tcx.normalize_canonicalized_free_alias(c_data),
-                    ty::Inherent => tcx.normalize_canonicalized_inherent_projection_ty(c_data),
-                    kind => unreachable!("did not expect {kind:?} due to match arm above"),
-                }?;
-                // We don't expect ambiguity.
-                if !result.value.is_proven() {
-                    // Rustdoc normalizes possibly not well-formed types, so only
-                    // treat this as a bug if we're not in rustdoc.
-                    if !tcx.sess.opts.actually_rustdoc {
-                        tcx.dcx()
-                            .delayed_bug(format!("unexpected ambiguity: {c_data:?} {result:?}"));
-                    }
-                    return Err(NoSolution);
-                }
-                let InferOk { value: result, obligations } = infcx
-                    .instantiate_query_response_and_region_obligations(
-                        self.cause,
-                        self.param_env,
-                        &orig_values,
-                        result,
-                    )?;
-                debug!("QueryNormalizer: result = {:#?}", result);
-                debug!("QueryNormalizer: obligations = {:#?}", obligations);
-                self.obligations.extend(obligations);
-                let res = if let Some((mapped_regions, mapped_types, mapped_consts)) = maps {
-                    PlaceholderReplacer::replace_placeholders(
-                        infcx,
-                        mapped_regions,
-                        mapped_types,
-                        mapped_consts,
-                        &self.universes,
-                        result.normalized_ty,
-                    )
-                } else {
-                    result.normalized_ty
-                };
-                // `tcx.normalize_canonicalized_projection_ty` may normalize to a type that
-                // still has unevaluated consts, so keep normalizing here if that's the case.
-                // Similarly, `tcx.normalize_canonicalized_free_alias` will only unwrap one layer
-                // of type and we need to continue folding it to reveal the TAIT behind it.
-                if res != ty
-                    && (res.has_type_flags(ty::TypeFlags::HAS_CT_PROJECTION) || kind == ty::Free)
-                {
-                    res.try_fold_with(self)?
-                } else {
-                    res
-                }
-            }
+            ty::Projection | ty::Inherent | ty::Free => self
+                .try_fold_free_or_assoc(ty::AliasTerm::new(self.cx(), data.def_id, data.args))?
+                .expect_type(),
         };
 
         self.cache.insert(ty, res);
@@ -337,12 +271,22 @@ impl<'a, 'tcx> FallibleTypeFolder<TyCtxt<'tcx>> for QueryNormalizer<'a, 'tcx> {
             return Ok(constant);
         }
 
-        let constant = crate::traits::with_replaced_escaping_bound_vars(
-            self.infcx,
-            &mut self.universes,
-            constant,
-            |constant| crate::traits::evaluate_const(&self.infcx, constant, self.param_env),
-        );
+        let uv = match constant.kind() {
+            ty::ConstKind::Unevaluated(uv) => uv,
+            _ => return constant.try_super_fold_with(self),
+        };
+
+        let constant = match self.cx().def_kind(uv.def) {
+            DefKind::AnonConst => crate::traits::with_replaced_escaping_bound_vars(
+                self.infcx,
+                &mut self.universes,
+                constant,
+                |constant| crate::traits::evaluate_const(&self.infcx, constant, self.param_env),
+            ),
+            _ => self
+                .try_fold_free_or_assoc(ty::AliasTerm::new(self.cx(), uv.def, uv.args))?
+                .expect_const(),
+        };
         debug!(?constant, ?self.param_env);
         constant.try_super_fold_with(self)
     }
@@ -359,3 +303,85 @@ impl<'a, 'tcx> FallibleTypeFolder<TyCtxt<'tcx>> for QueryNormalizer<'a, 'tcx> {
         }
     }
 }
+
+impl<'a, 'tcx> QueryNormalizer<'a, 'tcx> {
+    fn try_fold_free_or_assoc(
+        &mut self,
+        term: ty::AliasTerm<'tcx>,
+    ) -> Result<ty::Term<'tcx>, NoSolution> {
+        let infcx = self.infcx;
+        let tcx = infcx.tcx;
+        // Just an optimization: When we don't have escaping bound vars,
+        // we don't need to replace them with placeholders.
+        let (term, maps) = if term.has_escaping_bound_vars() {
+            let (term, mapped_regions, mapped_types, mapped_consts) =
+                BoundVarReplacer::replace_bound_vars(infcx, &mut self.universes, term);
+            (term, Some((mapped_regions, mapped_types, mapped_consts)))
+        } else {
+            (term, None)
+        };
+        let term = term.try_fold_with(self)?;
+
+        let mut orig_values = OriginalQueryValues::default();
+        let c_term = infcx.canonicalize_query(self.param_env.and(term), &mut orig_values);
+        debug!("QueryNormalizer: c_term = {:#?}", c_term);
+        debug!("QueryNormalizer: orig_values = {:#?}", orig_values);
+        let result = match term.kind(tcx) {
+            ty::AliasTermKind::ProjectionTy | ty::AliasTermKind::ProjectionConst => {
+                tcx.normalize_canonicalized_projection(c_term)
+            }
+            ty::AliasTermKind::FreeTy | ty::AliasTermKind::FreeConst => {
+                tcx.normalize_canonicalized_free_alias(c_term)
+            }
+            ty::AliasTermKind::InherentTy | ty::AliasTermKind::InherentConst => {
+                tcx.normalize_canonicalized_inherent_projection(c_term)
+            }
+            kind @ (ty::AliasTermKind::OpaqueTy | ty::AliasTermKind::UnevaluatedConst) => {
+                unreachable!("did not expect {kind:?} due to match arm above")
+            }
+        }?;
+        // We don't expect ambiguity.
+        if !result.value.is_proven() {
+            // Rustdoc normalizes possibly not well-formed types, so only
+            // treat this as a bug if we're not in rustdoc.
+            if !tcx.sess.opts.actually_rustdoc {
+                tcx.dcx().delayed_bug(format!("unexpected ambiguity: {c_term:?} {result:?}"));
+            }
+            return Err(NoSolution);
+        }
+        let InferOk { value: result, obligations } = infcx
+            .instantiate_query_response_and_region_obligations(
+                self.cause,
+                self.param_env,
+                &orig_values,
+                result,
+            )?;
+        debug!("QueryNormalizer: result = {:#?}", result);
+        debug!("QueryNormalizer: obligations = {:#?}", obligations);
+        self.obligations.extend(obligations);
+        let res = if let Some((mapped_regions, mapped_types, mapped_consts)) = maps {
+            PlaceholderReplacer::replace_placeholders(
+                infcx,
+                mapped_regions,
+                mapped_types,
+                mapped_consts,
+                &self.universes,
+                result.normalized_term,
+            )
+        } else {
+            result.normalized_term
+        };
+        // `tcx.normalize_canonicalized_projection` may normalize to a type that
+        // still has unevaluated consts, so keep normalizing here if that's the case.
+        // Similarly, `tcx.normalize_canonicalized_free_alias` will only unwrap one layer
+        // of type and we need to continue folding it to reveal the TAIT behind it.
+        if res != term.to_term(tcx)
+            && (res.as_type().map_or(false, |t| t.has_type_flags(ty::TypeFlags::HAS_CT_PROJECTION))
+                || term.kind(tcx) == ty::AliasTermKind::FreeTy)
+        {
+            res.try_fold_with(self)
+        } else {
+            Ok(res)
+        }
+    }
+}

compiler/rustc_traits/src/normalize_projection_ty.rs

@@ -12,18 +12,18 @@ use tracing::debug;
 
 pub(crate) fn provide(p: &mut Providers) {
     *p = Providers {
-        normalize_canonicalized_projection_ty,
+        normalize_canonicalized_projection,
         normalize_canonicalized_free_alias,
-        normalize_canonicalized_inherent_projection_ty,
+        normalize_canonicalized_inherent_projection,
         ..*p
     };
 }
 
-fn normalize_canonicalized_projection_ty<'tcx>(
+fn normalize_canonicalized_projection<'tcx>(
     tcx: TyCtxt<'tcx>,
     goal: CanonicalAliasGoal<'tcx>,
 ) -> Result<&'tcx Canonical<'tcx, QueryResponse<'tcx, NormalizationResult<'tcx>>>, NoSolution> {
-    debug!("normalize_canonicalized_projection_ty(goal={:#?})", goal);
+    debug!("normalize_canonicalized_projection(goal={:#?})", goal);
 
     tcx.infer_ctxt().enter_canonical_trait_query(
         &goal,
@@ -32,7 +32,7 @@ fn normalize_canonicalized_projection_ty<'tcx>(
             let selcx = &mut SelectionContext::new(ocx.infcx);
             let cause = ObligationCause::dummy();
             let mut obligations = PredicateObligations::new();
-            let answer = traits::normalize_projection_term(
+            let normalized_term = traits::normalize_projection_term(
                 selcx,
                 param_env,
                 goal.into(),
@@ -61,10 +61,7 @@ fn normalize_canonicalized_projection_ty<'tcx>(
                 return Err(NoSolution);
             }
 
-            // FIXME(associated_const_equality): All users of normalize_canonicalized_projection_ty
-            // expected a type, but there is the possibility it could've been a const now.
-            // Maybe change it to a Term later?
-            Ok(NormalizationResult { normalized_ty: answer.expect_type() })
+            Ok(NormalizationResult { normalized_term })
         },
     )
 }
@@ -89,25 +86,29 @@ fn normalize_canonicalized_free_alias<'tcx>(
                 },
             );
             ocx.register_obligations(obligations);
-            let normalized_ty = tcx.type_of(goal.def_id).instantiate(tcx, goal.args);
-            Ok(NormalizationResult { normalized_ty })
+            let normalized_term = if goal.kind(tcx).is_type() {
+                tcx.type_of(goal.def_id).instantiate(tcx, goal.args).into()
+            } else {
+                todo!()
+            };
+            Ok(NormalizationResult { normalized_term })
         },
     )
 }
 
-fn normalize_canonicalized_inherent_projection_ty<'tcx>(
+fn normalize_canonicalized_inherent_projection<'tcx>(
     tcx: TyCtxt<'tcx>,
     goal: CanonicalAliasGoal<'tcx>,
 ) -> Result<&'tcx Canonical<'tcx, QueryResponse<'tcx, NormalizationResult<'tcx>>>, NoSolution> {
-    debug!("normalize_canonicalized_inherent_projection_ty(goal={:#?})", goal);
+    debug!("normalize_canonicalized_inherent_projection(goal={:#?})", goal);
 
     tcx.infer_ctxt().enter_canonical_trait_query(
         &goal,
         |ocx, ParamEnvAnd { param_env, value: goal }| {
             let selcx = &mut SelectionContext::new(ocx.infcx);
             let cause = ObligationCause::dummy();
             let mut obligations = PredicateObligations::new();
-            let answer = traits::normalize_inherent_projection(
+            let normalized_term = traits::normalize_inherent_projection(
                 selcx,
                 param_env,
                 goal.into(),
@@ -117,7 +118,7 @@ fn normalize_canonicalized_inherent_projection_ty<'tcx>(
             );
             ocx.register_obligations(obligations);
 
-            Ok(NormalizationResult { normalized_ty: answer.expect_type() })
+            Ok(NormalizationResult { normalized_term })
         },
     )
 }