Auto merge of #130821 - lcnr:nalgebra-hang-2, r=compiler-errors

add caching to most type folders, rm region uniquification

Fixes the new minimization of the hang in nalgebra and nalgebra itself :3

this is a bit iffy, especially the cache in `TypeRelating`. I believe all the caches are correct, but it definitely adds some non-local complexity in places. The first commit removes region uniquification, reintroducing the ICE from https://github.com/rust-lang/trait-system-refactor-initiative/issues/27. This does not affect coherence and I would like to fix this by introducing OR-region constraints

r? `@compiler-errors`

compiler/rustc_next_trait_solver/src/canonicalizer.rs

@@ -1,5 +1,6 @@
 use std::cmp::Ordering;
 
+use rustc_type_ir::data_structures::HashMap;
 use rustc_type_ir::fold::{TypeFoldable, TypeFolder, TypeSuperFoldable};
 use rustc_type_ir::inherent::*;
 use rustc_type_ir::visit::TypeVisitableExt;
@@ -41,11 +42,20 @@ pub enum CanonicalizeMode {
 
 pub struct Canonicalizer<'a, D: SolverDelegate<Interner = I>, I: Interner> {
     delegate: &'a D,
+
+    // Immutable field.
     canonicalize_mode: CanonicalizeMode,
 
+    // Mutable fields.
     variables: &'a mut Vec<I::GenericArg>,
     primitive_var_infos: Vec<CanonicalVarInfo<I>>,
+    variable_lookup_table: HashMap<I::GenericArg, usize>,
     binder_index: ty::DebruijnIndex,
+
+    /// We only use the debruijn index during lookup. We don't need to
+    /// track the `variables` as each generic arg only results in a single
+    /// bound variable regardless of how many times it is encountered.
+    cache: HashMap<(ty::DebruijnIndex, I::Ty), I::Ty>,
 }
 
 impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
@@ -60,12 +70,14 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
             canonicalize_mode,
 
             variables,
+            variable_lookup_table: Default::default(),
             primitive_var_infos: Vec::new(),
             binder_index: ty::INNERMOST,
+
+            cache: Default::default(),
         };
 
         let value = value.fold_with(&mut canonicalizer);
-        // FIXME: Restore these assertions. Should we uplift type flags?
         assert!(!value.has_infer(), "unexpected infer in {value:?}");
         assert!(!value.has_placeholders(), "unexpected placeholders in {value:?}");
 
@@ -75,6 +87,37 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
         Canonical { defining_opaque_types, max_universe, variables, value }
     }
 
+    fn get_or_insert_bound_var(
+        &mut self,
+        arg: impl Into<I::GenericArg>,
+        canonical_var_info: CanonicalVarInfo<I>,
+    ) -> ty::BoundVar {
+        // FIXME: 16 is made up and arbitrary. We should look at some
+        // perf data here.
+        let arg = arg.into();
+        let idx = if self.variables.len() > 16 {
+            if self.variable_lookup_table.is_empty() {
+                self.variable_lookup_table.extend(self.variables.iter().copied().zip(0..));
+            }
+
+            *self.variable_lookup_table.entry(arg).or_insert_with(|| {
+                let var = self.variables.len();
+                self.variables.push(arg);
+                self.primitive_var_infos.push(canonical_var_info);
+                var
+            })
+        } else {
+            self.variables.iter().position(|&v| v == arg).unwrap_or_else(|| {
+                let var = self.variables.len();
+                self.variables.push(arg);
+                self.primitive_var_infos.push(canonical_var_info);
+                var
+            })
+        };
+
+        ty::BoundVar::from(idx)
+    }
+
     fn finalize(self) -> (ty::UniverseIndex, I::CanonicalVars) {
         let mut var_infos = self.primitive_var_infos;
         // See the rustc-dev-guide section about how we deal with universes
@@ -124,8 +167,8 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
         // - var_infos: [E0, U1, E2, U1, E1, E6, U6], curr_compressed_uv: 2, next_orig_uv: 6
         // - var_infos: [E0, U1, E1, U1, E1, E3, U3], curr_compressed_uv: 2, next_orig_uv: -
         //
-        // This algorithm runs in `O(n²)` where `n` is the number of different universe
-        // indices in the input. This should be fine as `n` is expected to be small.
+        // This algorithm runs in `O(mn)` where `n` is the number of different universes and
+        // `m` the number of variables. This should be fine as both are expected to be small.
         let mut curr_compressed_uv = ty::UniverseIndex::ROOT;
         let mut existential_in_new_uv = None;
         let mut next_orig_uv = Some(ty::UniverseIndex::ROOT);
@@ -185,14 +228,16 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
                 for var in var_infos.iter_mut() {
                     // We simply put all regions from the input into the highest
                     // compressed universe, so we only deal with them at the end.
-                    if !var.is_region() && is_existential == var.is_existential() {
-                        update_uv(var, orig_uv, is_existential)
+                    if !var.is_region() {
+                        if is_existential == var.is_existential() {
+                            update_uv(var, orig_uv, is_existential)
+                        }
                     }
                 }
             }
         }
 
-        // We uniquify regions and always put them into their own universe
+        // We put all regions into a separate universe.
         let mut first_region = true;
         for var in var_infos.iter_mut() {
             if var.is_region() {
@@ -208,93 +253,8 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
         let var_infos = self.delegate.cx().mk_canonical_var_infos(&var_infos);
         (curr_compressed_uv, var_infos)
     }
-}
 
-impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicalizer<'_, D, I> {
-    fn cx(&self) -> I {
-        self.delegate.cx()
-    }
-
-    fn fold_binder<T>(&mut self, t: ty::Binder<I, T>) -> ty::Binder<I, T>
-    where
-        T: TypeFoldable<I>,
-    {
-        self.binder_index.shift_in(1);
-        let t = t.super_fold_with(self);
-        self.binder_index.shift_out(1);
-        t
-    }
-
-    fn fold_region(&mut self, r: I::Region) -> I::Region {
-        let kind = match r.kind() {
-            ty::ReBound(..) => return r,
-
-            // We may encounter `ReStatic` in item signatures or the hidden type
-            // of an opaque. `ReErased` should only be encountered in the hidden
-            // type of an opaque for regions that are ignored for the purposes of
-            // captures.
-            //
-            // FIXME: We should investigate the perf implications of not uniquifying
-            // `ReErased`. We may be able to short-circuit registering region
-            // obligations if we encounter a `ReErased` on one side, for example.
-            ty::ReStatic | ty::ReErased | ty::ReError(_) => match self.canonicalize_mode {
-                CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
-                CanonicalizeMode::Response { .. } => return r,
-            },
-
-            ty::ReEarlyParam(_) | ty::ReLateParam(_) => match self.canonicalize_mode {
-                CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
-                CanonicalizeMode::Response { .. } => {
-                    panic!("unexpected region in response: {r:?}")
-                }
-            },
-
-            ty::RePlaceholder(placeholder) => match self.canonicalize_mode {
-                // We canonicalize placeholder regions as existentials in query inputs.
-                CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
-                CanonicalizeMode::Response { max_input_universe } => {
-                    // If we have a placeholder region inside of a query, it must be from
-                    // a new universe.
-                    if max_input_universe.can_name(placeholder.universe()) {
-                        panic!("new placeholder in universe {max_input_universe:?}: {r:?}");
-                    }
-                    CanonicalVarKind::PlaceholderRegion(placeholder)
-                }
-            },
-
-            ty::ReVar(vid) => {
-                assert_eq!(
-                    self.delegate.opportunistic_resolve_lt_var(vid),
-                    r,
-                    "region vid should have been resolved fully before canonicalization"
-                );
-                match self.canonicalize_mode {
-                    CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
-                    CanonicalizeMode::Response { .. } => {
-                        CanonicalVarKind::Region(self.delegate.universe_of_lt(vid).unwrap())
-                    }
-                }
-            }
-        };
-
-        let existing_bound_var = match self.canonicalize_mode {
-            CanonicalizeMode::Input => None,
-            CanonicalizeMode::Response { .. } => {
-                self.variables.iter().position(|&v| v == r.into()).map(ty::BoundVar::from)
-            }
-        };
-
-        let var = existing_bound_var.unwrap_or_else(|| {
-            let var = ty::BoundVar::from(self.variables.len());
-            self.variables.push(r.into());
-            self.primitive_var_infos.push(CanonicalVarInfo { kind });
-            var
-        });
-
-        Region::new_anon_bound(self.cx(), self.binder_index, var)
-    }
-
-    fn fold_ty(&mut self, t: I::Ty) -> I::Ty {
+    fn cached_fold_ty(&mut self, t: I::Ty) -> I::Ty {
         let kind = match t.kind() {
             ty::Infer(i) => match i {
                 ty::TyVar(vid) => {
@@ -368,20 +328,98 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
             | ty::Tuple(_)
             | ty::Alias(_, _)
             | ty::Bound(_, _)
-            | ty::Error(_) => return t.super_fold_with(self),
+            | ty::Error(_) => {
+                return t.super_fold_with(self);
+            }
         };
 
-        let var = ty::BoundVar::from(
-            self.variables.iter().position(|&v| v == t.into()).unwrap_or_else(|| {
-                let var = self.variables.len();
-                self.variables.push(t.into());
-                self.primitive_var_infos.push(CanonicalVarInfo { kind });
-                var
-            }),
-        );
+        let var = self.get_or_insert_bound_var(t, CanonicalVarInfo { kind });
 
         Ty::new_anon_bound(self.cx(), self.binder_index, var)
     }
+}
+
+impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicalizer<'_, D, I> {
+    fn cx(&self) -> I {
+        self.delegate.cx()
+    }
+
+    fn fold_binder<T>(&mut self, t: ty::Binder<I, T>) -> ty::Binder<I, T>
+    where
+        T: TypeFoldable<I>,
+    {
+        self.binder_index.shift_in(1);
+        let t = t.super_fold_with(self);
+        self.binder_index.shift_out(1);
+        t
+    }
+
+    fn fold_region(&mut self, r: I::Region) -> I::Region {
+        let kind = match r.kind() {
+            ty::ReBound(..) => return r,
+
+            // We may encounter `ReStatic` in item signatures or the hidden type
+            // of an opaque. `ReErased` should only be encountered in the hidden
+            // type of an opaque for regions that are ignored for the purposes of
+            // captures.
+            //
+            // FIXME: We should investigate the perf implications of not uniquifying
+            // `ReErased`. We may be able to short-circuit registering region
+            // obligations if we encounter a `ReErased` on one side, for example.
+            ty::ReStatic | ty::ReErased | ty::ReError(_) => match self.canonicalize_mode {
+                CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
+                CanonicalizeMode::Response { .. } => return r,
+            },
+
+            ty::ReEarlyParam(_) | ty::ReLateParam(_) => match self.canonicalize_mode {
+                CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
+                CanonicalizeMode::Response { .. } => {
+                    panic!("unexpected region in response: {r:?}")
+                }
+            },
+
+            ty::RePlaceholder(placeholder) => match self.canonicalize_mode {
+                // We canonicalize placeholder regions as existentials in query inputs.
+                CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
+                CanonicalizeMode::Response { max_input_universe } => {
+                    // If we have a placeholder region inside of a query, it must be from
+                    // a new universe.
+                    if max_input_universe.can_name(placeholder.universe()) {
+                        panic!("new placeholder in universe {max_input_universe:?}: {r:?}");
+                    }
+                    CanonicalVarKind::PlaceholderRegion(placeholder)
+                }
+            },
+
+            ty::ReVar(vid) => {
+                assert_eq!(
+                    self.delegate.opportunistic_resolve_lt_var(vid),
+                    r,
+                    "region vid should have been resolved fully before canonicalization"
+                );
+                match self.canonicalize_mode {
+                    CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
+                    CanonicalizeMode::Response { .. } => {
+                        CanonicalVarKind::Region(self.delegate.universe_of_lt(vid).unwrap())
+                    }
+                }
+            }
+        };
+
+        let var = self.get_or_insert_bound_var(r, CanonicalVarInfo { kind });
+
+        Region::new_anon_bound(self.cx(), self.binder_index, var)
+    }
+
+    fn fold_ty(&mut self, t: I::Ty) -> I::Ty {
+        if let Some(&ty) = self.cache.get(&(self.binder_index, t)) {
+            ty
+        } else {
+            let res = self.cached_fold_ty(t);
+            assert!(self.cache.insert((self.binder_index, t), res).is_none());
+            res
+        }
+    }
 
     fn fold_const(&mut self, c: I::Const) -> I::Const {
         let kind = match c.kind() {
@@ -419,14 +457,7 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
             | ty::ConstKind::Expr(_) => return c.super_fold_with(self),
         };
 
-        let var = ty::BoundVar::from(
-            self.variables.iter().position(|&v| v == c.into()).unwrap_or_else(|| {
-                let var = self.variables.len();
-                self.variables.push(c.into());
-                self.primitive_var_infos.push(CanonicalVarInfo { kind });
-                var
-            }),
-        );
+        let var = self.get_or_insert_bound_var(c, CanonicalVarInfo { kind });
 
         Const::new_anon_bound(self.cx(), self.binder_index, var)
     }

compiler/rustc_next_trait_solver/src/resolve.rs

@@ -1,3 +1,4 @@
+use rustc_type_ir::data_structures::DelayedMap;
 use rustc_type_ir::fold::{TypeFoldable, TypeFolder, TypeSuperFoldable};
 use rustc_type_ir::inherent::*;
 use rustc_type_ir::visit::TypeVisitableExt;
@@ -15,11 +16,14 @@ where
     I: Interner,
 {
     delegate: &'a D,
+    /// We're able to use a cache here as the folder does not have any
+    /// mutable state.
+    cache: DelayedMap<I::Ty, I::Ty>,
 }
 
 impl<'a, D: SolverDelegate> EagerResolver<'a, D> {
     pub fn new(delegate: &'a D) -> Self {
-        EagerResolver { delegate }
+        EagerResolver { delegate, cache: Default::default() }
     }
 }
 
@@ -42,7 +46,12 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for EagerResolv
             ty::Infer(ty::FloatVar(vid)) => self.delegate.opportunistic_resolve_float_var(vid),
             _ => {
                 if t.has_infer() {
-                    t.super_fold_with(self)
+                    if let Some(&ty) = self.cache.get(&t) {
+                        return ty;
+                    }
+                    let res = t.super_fold_with(self);
+                    assert!(self.cache.insert(t, res));
+                    res
                 } else {
                     t
                 }

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

@@ -3,7 +3,7 @@ use std::ops::ControlFlow;
 use derive_where::derive_where;
 #[cfg(feature = "nightly")]
 use rustc_macros::{HashStable_NoContext, TyDecodable, TyEncodable};
-use rustc_type_ir::data_structures::ensure_sufficient_stack;
+use rustc_type_ir::data_structures::{HashMap, HashSet, ensure_sufficient_stack};
 use rustc_type_ir::fold::{TypeFoldable, TypeFolder, TypeSuperFoldable};
 use rustc_type_ir::inherent::*;
 use rustc_type_ir::relate::Relate;
@@ -579,18 +579,16 @@ where
 
     #[instrument(level = "trace", skip(self))]
     pub(super) fn add_normalizes_to_goal(&mut self, mut goal: Goal<I, ty::NormalizesTo<I>>) {
-        goal.predicate = goal
-            .predicate
-            .fold_with(&mut ReplaceAliasWithInfer { ecx: self, param_env: goal.param_env });
+        goal.predicate =
+            goal.predicate.fold_with(&mut ReplaceAliasWithInfer::new(self, goal.param_env));
         self.inspect.add_normalizes_to_goal(self.delegate, self.max_input_universe, goal);
         self.nested_goals.normalizes_to_goals.push(goal);
     }
 
     #[instrument(level = "debug", skip(self))]
     pub(super) fn add_goal(&mut self, source: GoalSource, mut goal: Goal<I, I::Predicate>) {
-        goal.predicate = goal
-            .predicate
-            .fold_with(&mut ReplaceAliasWithInfer { ecx: self, param_env: goal.param_env });
+        goal.predicate =
+            goal.predicate.fold_with(&mut ReplaceAliasWithInfer::new(self, goal.param_env));
         self.inspect.add_goal(self.delegate, self.max_input_universe, source, goal);
         self.nested_goals.goals.push((source, goal));
     }
@@ -654,6 +652,7 @@ where
             term: I::Term,
             universe_of_term: ty::UniverseIndex,
             delegate: &'a D,
+            cache: HashSet<I::Ty>,
         }
 
         impl<D: SolverDelegate<Interner = I>, I: Interner> ContainsTermOrNotNameable<'_, D, I> {
@@ -671,6 +670,10 @@ where
         {
             type Result = ControlFlow<()>;
             fn visit_ty(&mut self, t: I::Ty) -> Self::Result {
+                if self.cache.contains(&t) {
+                    return ControlFlow::Continue(());
+                }
+
                 match t.kind() {
                     ty::Infer(ty::TyVar(vid)) => {
                         if let ty::TermKind::Ty(term) = self.term.kind() {
@@ -683,17 +686,18 @@ where
                             }
                         }
 
-                        self.check_nameable(self.delegate.universe_of_ty(vid).unwrap())
+                        self.check_nameable(self.delegate.universe_of_ty(vid).unwrap())?;
                     }
-                    ty::Placeholder(p) => self.check_nameable(p.universe()),
+                    ty::Placeholder(p) => self.check_nameable(p.universe())?,
                     _ => {
                         if t.has_non_region_infer() || t.has_placeholders() {
-                            t.super_visit_with(self)
-                        } else {
-                            ControlFlow::Continue(())
+                            t.super_visit_with(self)?
                         }
                     }
                 }
+
+                assert!(self.cache.insert(t));
+                ControlFlow::Continue(())
             }
 
             fn visit_const(&mut self, c: I::Const) -> Self::Result {
@@ -728,6 +732,7 @@ where
             delegate: self.delegate,
             universe_of_term,
             term: goal.predicate.term,
+            cache: Default::default(),
         };
         goal.predicate.alias.visit_with(&mut visitor).is_continue()
             && goal.param_env.visit_with(&mut visitor).is_continue()
@@ -1017,6 +1022,17 @@ where
 {
     ecx: &'me mut EvalCtxt<'a, D>,
     param_env: I::ParamEnv,
+    cache: HashMap<I::Ty, I::Ty>,
+}
+
+impl<'me, 'a, D, I> ReplaceAliasWithInfer<'me, 'a, D, I>
+where
+    D: SolverDelegate<Interner = I>,
+    I: Interner,
+{
+    fn new(ecx: &'me mut EvalCtxt<'a, D>, param_env: I::ParamEnv) -> Self {
+        ReplaceAliasWithInfer { ecx, param_env, cache: Default::default() }
+    }
 }
 
 impl<D, I> TypeFolder<I> for ReplaceAliasWithInfer<'_, '_, D, I>
@@ -1043,7 +1059,17 @@ where
                 );
                 infer_ty
             }
-            _ => ty.super_fold_with(self),
+            _ => {
+                if !ty.has_aliases() {
+                    ty
+                } else if let Some(&entry) = self.cache.get(&ty) {
+                    return entry;
+                } else {
+                    let res = ty.super_fold_with(self);
+                    assert!(self.cache.insert(ty, res).is_none());
+                    res
+                }
+            }
         }
     }