Rollup merge of #107981 - lcnr:canonicalization-uwu, r=compiler-errors

new solver: implement canonicalization and region constraints

see the corresponding rustc-dev-guide chapter: https://rustc-dev-guide.rust-lang.org/solve/canonicalization.html

r? ``@compiler-errors``

compiler/rustc_infer/src/infer/at.rs

@@ -369,6 +369,34 @@ impl<'tcx> ToTrace<'tcx> for Const<'tcx> {
     }
 }
 
+impl<'tcx> ToTrace<'tcx> for ty::GenericArg<'tcx> {
+    fn to_trace(
+        _: TyCtxt<'tcx>,
+        cause: &ObligationCause<'tcx>,
+        a_is_expected: bool,
+        a: Self,
+        b: Self,
+    ) -> TypeTrace<'tcx> {
+        use GenericArgKind::*;
+        TypeTrace {
+            cause: cause.clone(),
+            values: match (a.unpack(), b.unpack()) {
+                (Lifetime(a), Lifetime(b)) => Regions(ExpectedFound::new(a_is_expected, a, b)),
+                (Type(a), Type(b)) => Terms(ExpectedFound::new(a_is_expected, a.into(), b.into())),
+                (Const(a), Const(b)) => {
+                    Terms(ExpectedFound::new(a_is_expected, a.into(), b.into()))
+                }
+
+                (Lifetime(_), Type(_) | Const(_))
+                | (Type(_), Lifetime(_) | Const(_))
+                | (Const(_), Lifetime(_) | Type(_)) => {
+                    bug!("relating different kinds: {a:?} {b:?}")
+                }
+            },
+        }
+    }
+}
+
 impl<'tcx> ToTrace<'tcx> for ty::Term<'tcx> {
     fn to_trace(
         _: TyCtxt<'tcx>,

compiler/rustc_infer/src/infer/canonical/mod.rs

@@ -30,7 +30,7 @@ use rustc_middle::ty::{self, List, TyCtxt};
 use rustc_span::source_map::Span;
 
 pub use rustc_middle::infer::canonical::*;
-use substitute::CanonicalExt;
+pub use substitute::CanonicalExt;
 
 mod canonicalizer;
 pub mod query_response;
@@ -100,7 +100,11 @@ impl<'tcx> InferCtxt<'tcx> {
     /// variable for it. If this is an existentially quantified
     /// variable, then you'll get a new inference variable; if it is a
     /// universally quantified variable, you get a placeholder.
-    fn instantiate_canonical_var(
+    ///
+    /// FIXME(-Ztrait-solver=next): This is public because it's used by the
+    /// new trait solver which has a different canonicalization routine.
+    /// We should somehow deduplicate all of this.
+    pub fn instantiate_canonical_var(
         &self,
         span: Span,
         cv_info: CanonicalVarInfo<'tcx>,

compiler/rustc_infer/src/infer/canonical/query_response.rs

@@ -151,11 +151,21 @@ impl<'tcx> InferCtxt<'tcx> {
         })
     }
 
-    /// FIXME: This method should only be used for canonical queries and therefore be private.
-    ///
-    /// As the new solver does canonicalization slightly differently, this is also used there
-    /// for now. This should hopefully change fairly soon.
-    pub fn take_opaque_types_for_query_response(&self) -> Vec<(Ty<'tcx>, Ty<'tcx>)> {
+    /// Used by the new solver as that one takes the opaque types at the end of a probe
+    /// to deal with multiple candidates without having to recompute them.
+    pub fn clone_opaque_types_for_query_response(&self) -> Vec<(Ty<'tcx>, Ty<'tcx>)> {
+        self.inner
+            .borrow()
+            .opaque_type_storage
+            .opaque_types
+            .iter()
+            .map(|&(k, ref v)| {
+                (self.tcx.mk_opaque(k.def_id.to_def_id(), k.substs), v.hidden_type.ty)
+            })
+            .collect()
+    }
+
+    fn take_opaque_types_for_query_response(&self) -> Vec<(Ty<'tcx>, Ty<'tcx>)> {
         std::mem::take(&mut self.inner.borrow_mut().opaque_type_storage.opaque_types)
             .into_iter()
             .map(|(k, v)| (self.tcx.mk_opaque(k.def_id.to_def_id(), k.substs), v.hidden_type.ty))

compiler/rustc_infer/src/infer/canonical/substitute.rs

@@ -11,7 +11,9 @@ use rustc_middle::ty::fold::{FnMutDelegate, TypeFoldable};
 use rustc_middle::ty::subst::GenericArgKind;
 use rustc_middle::ty::{self, TyCtxt};
 
-pub(super) trait CanonicalExt<'tcx, V> {
+/// FIXME(-Ztrait-solver=next): This or public because it is shared with the
+/// new trait solver implementation. We should deduplicate canonicalization.
+pub trait CanonicalExt<'tcx, V> {
     /// Instantiate the wrapped value, replacing each canonical value
     /// with the value given in `var_values`.
     fn substitute(&self, tcx: TyCtxt<'tcx>, var_values: &CanonicalVarValues<'tcx>) -> V

compiler/rustc_middle/src/infer/canonical.rs

@@ -123,6 +123,11 @@ impl<'tcx> CanonicalVarInfo<'tcx> {
         self.kind.universe()
     }
 
+    #[must_use]
+    pub fn with_updated_universe(self, ui: ty::UniverseIndex) -> CanonicalVarInfo<'tcx> {
+        CanonicalVarInfo { kind: self.kind.with_updated_universe(ui) }
+    }
+
     pub fn is_existential(&self) -> bool {
         match self.kind {
             CanonicalVarKind::Ty(_) => true,
@@ -133,6 +138,28 @@ impl<'tcx> CanonicalVarInfo<'tcx> {
             CanonicalVarKind::PlaceholderConst(_, _) => false,
         }
     }
+
+    pub fn is_region(&self) -> bool {
+        match self.kind {
+            CanonicalVarKind::Region(_) | CanonicalVarKind::PlaceholderRegion(_) => true,
+            CanonicalVarKind::Ty(_)
+            | CanonicalVarKind::PlaceholderTy(_)
+            | CanonicalVarKind::Const(_, _)
+            | CanonicalVarKind::PlaceholderConst(_, _) => false,
+        }
+    }
+
+    pub fn expect_anon_placeholder(self) -> u32 {
+        match self.kind {
+            CanonicalVarKind::Ty(_)
+            | CanonicalVarKind::Region(_)
+            | CanonicalVarKind::Const(_, _) => bug!("expected placeholder: {self:?}"),
+
+            CanonicalVarKind::PlaceholderRegion(placeholder) => placeholder.name.expect_anon(),
+            CanonicalVarKind::PlaceholderTy(placeholder) => placeholder.name.expect_anon(),
+            CanonicalVarKind::PlaceholderConst(placeholder, _) => placeholder.name.as_u32(),
+        }
+    }
 }
 
 /// Describes the "kind" of the canonical variable. This is a "kind"
@@ -177,6 +204,38 @@ impl<'tcx> CanonicalVarKind<'tcx> {
             CanonicalVarKind::PlaceholderConst(placeholder, _) => placeholder.universe,
         }
     }
+
+    /// Replaces the universe of this canonical variable with `ui`.
+    ///
+    /// In case this is a float or int variable, this causes an ICE if
+    /// the updated universe is not the root.
+    pub fn with_updated_universe(self, ui: ty::UniverseIndex) -> CanonicalVarKind<'tcx> {
+        match self {
+            CanonicalVarKind::Ty(kind) => match kind {
+                CanonicalTyVarKind::General(_) => {
+                    CanonicalVarKind::Ty(CanonicalTyVarKind::General(ui))
+                }
+                CanonicalTyVarKind::Int | CanonicalTyVarKind::Float => {
+                    assert_eq!(ui, ty::UniverseIndex::ROOT);
+                    CanonicalVarKind::Ty(kind)
+                }
+            },
+            CanonicalVarKind::PlaceholderTy(placeholder) => {
+                CanonicalVarKind::PlaceholderTy(ty::Placeholder { universe: ui, ..placeholder })
+            }
+            CanonicalVarKind::Region(_) => CanonicalVarKind::Region(ui),
+            CanonicalVarKind::PlaceholderRegion(placeholder) => {
+                CanonicalVarKind::PlaceholderRegion(ty::Placeholder { universe: ui, ..placeholder })
+            }
+            CanonicalVarKind::Const(_, ty) => CanonicalVarKind::Const(ui, ty),
+            CanonicalVarKind::PlaceholderConst(placeholder, ty) => {
+                CanonicalVarKind::PlaceholderConst(
+                    ty::Placeholder { universe: ui, ..placeholder },
+                    ty,
+                )
+            }
+        }
+    }
 }
 
 /// Rust actually has more than one category of type variables;
@@ -213,7 +272,8 @@ pub struct QueryResponse<'tcx, R> {
     pub value: R,
 }
 
-#[derive(Clone, Debug, Default, HashStable, TypeFoldable, TypeVisitable, Lift)]
+#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
+#[derive(HashStable, TypeFoldable, TypeVisitable, Lift)]
 pub struct QueryRegionConstraints<'tcx> {
     pub outlives: Vec<QueryOutlivesConstraint<'tcx>>,
     pub member_constraints: Vec<MemberConstraint<'tcx>>,

compiler/rustc_middle/src/infer/mod.rs

@@ -12,7 +12,8 @@ use rustc_span::Span;
 /// ```text
 /// R0 member of [O1..On]
 /// ```
-#[derive(Debug, Clone, HashStable, TypeFoldable, TypeVisitable, Lift)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+#[derive(HashStable, TypeFoldable, TypeVisitable, Lift)]
 pub struct MemberConstraint<'tcx> {
     /// The `DefId` and substs of the opaque type causing this constraint.
     /// Used for error reporting.

compiler/rustc_middle/src/traits/solve.rs

@@ -2,6 +2,7 @@ use std::ops::ControlFlow;
 
 use rustc_data_structures::intern::Interned;
 
+use crate::infer::canonical::QueryRegionConstraints;
 use crate::ty::{
     FallibleTypeFolder, Ty, TyCtxt, TypeFoldable, TypeFolder, TypeVisitable, TypeVisitor,
 };
@@ -18,20 +19,25 @@ impl<'tcx> std::ops::Deref for ExternalConstraints<'tcx> {
 }
 
 /// Additional constraints returned on success.
-#[derive(Debug, PartialEq, Eq, Clone, Hash, Default)]
+#[derive(Debug, PartialEq, Eq, Clone, Hash, Default, TypeFoldable, TypeVisitable)]
 pub struct ExternalConstraintsData<'tcx> {
     // FIXME: implement this.
-    pub regions: (),
+    pub region_constraints: QueryRegionConstraints<'tcx>,
     pub opaque_types: Vec<(Ty<'tcx>, Ty<'tcx>)>,
 }
 
+// FIXME: Having to clone `region_constraints` for folding feels bad and
+// probably isn't great wrt performance.
+//
+// Not sure how to fix this, maybe we should also intern `opaque_types` and
+// `region_constraints` here or something.
 impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for ExternalConstraints<'tcx> {
     fn try_fold_with<F: FallibleTypeFolder<TyCtxt<'tcx>>>(
         self,
         folder: &mut F,
     ) -> Result<Self, F::Error> {
         Ok(FallibleTypeFolder::interner(folder).mk_external_constraints(ExternalConstraintsData {
-            regions: (),
+            region_constraints: self.region_constraints.clone().try_fold_with(folder)?,
             opaque_types: self
                 .opaque_types
                 .iter()
@@ -42,7 +48,7 @@ impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for ExternalConstraints<'tcx> {
 
     fn fold_with<F: TypeFolder<TyCtxt<'tcx>>>(self, folder: &mut F) -> Self {
         TypeFolder::interner(folder).mk_external_constraints(ExternalConstraintsData {
-            regions: (),
+            region_constraints: self.region_constraints.clone().fold_with(folder),
             opaque_types: self.opaque_types.iter().map(|opaque| opaque.fold_with(folder)).collect(),
         })
     }
@@ -53,7 +59,7 @@ impl<'tcx> TypeVisitable<TyCtxt<'tcx>> for ExternalConstraints<'tcx> {
         &self,
         visitor: &mut V,
     ) -> std::ops::ControlFlow<V::BreakTy> {
-        self.regions.visit_with(visitor)?;
+        self.region_constraints.visit_with(visitor)?;
         self.opaque_types.visit_with(visitor)?;
         ControlFlow::Continue(())
     }

compiler/rustc_middle/src/ty/sty.rs

@@ -107,6 +107,15 @@ impl BoundRegionKind {
             _ => None,
         }
     }
+
+    pub fn expect_anon(&self) -> u32 {
+        match *self {
+            BoundRegionKind::BrNamed(_, _) | BoundRegionKind::BrEnv => {
+                bug!("expected anon region: {self:?}")
+            }
+            BoundRegionKind::BrAnon(idx, _) => idx,
+        }
+    }
 }
 
 pub trait Article {