Auto merge of #141451 - lcnr:canonicalize-env-cache, r=<try>

add more `TypeFlags` fast paths, cache `param_env` canonicalization

BLocked on #141581

Author: rust-bors[bot]

compiler/rustc_middle/src/ty/context.rs

@@ -179,6 +179,17 @@ impl<'tcx> Interner for TyCtxt<'tcx> {
         f(&mut *self.new_solver_evaluation_cache.lock())
     }
 
+    fn canonical_param_env_cache_get_or_insert<R>(
+        self,
+        param_env: ty::ParamEnv<'tcx>,
+        f: impl FnOnce() -> ty::CanonicalParamEnvCacheEntry<Self>,
+        from_entry: impl FnOnce(&ty::CanonicalParamEnvCacheEntry<Self>) -> R,
+    ) -> R {
+        let mut cache = self.new_solver_canonical_param_env_cache.lock();
+        let entry = cache.entry(param_env).or_insert_with(f);
+        from_entry(entry)
+    }
+
     fn evaluation_is_concurrent(&self) -> bool {
         self.sess.threads() > 1
     }
@@ -1444,6 +1455,8 @@ pub struct GlobalCtxt<'tcx> {
 
     /// Caches the results of goal evaluation in the new solver.
     pub new_solver_evaluation_cache: Lock<search_graph::GlobalCache<TyCtxt<'tcx>>>,
+    pub new_solver_canonical_param_env_cache:
+        Lock<FxHashMap<ty::ParamEnv<'tcx>, ty::CanonicalParamEnvCacheEntry<TyCtxt<'tcx>>>>,
 
     pub canonical_param_env_cache: CanonicalParamEnvCache<'tcx>,
 
@@ -1692,6 +1705,7 @@ impl<'tcx> TyCtxt<'tcx> {
             selection_cache: Default::default(),
             evaluation_cache: Default::default(),
             new_solver_evaluation_cache: Default::default(),
+            new_solver_canonical_param_env_cache: Default::default(),
             canonical_param_env_cache: Default::default(),
             data_layout,
             alloc_map: interpret::AllocMap::new(),

compiler/rustc_next_trait_solver/src/canonicalizer.rs

@@ -4,8 +4,9 @@ use rustc_type_ir::data_structures::{HashMap, ensure_sufficient_stack};
 use rustc_type_ir::inherent::*;
 use rustc_type_ir::solve::{Goal, QueryInput};
 use rustc_type_ir::{
-    self as ty, Canonical, CanonicalTyVarKind, CanonicalVarKind, Flags, InferCtxtLike, Interner,
-    TypeFlags, TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitableExt,
+    self as ty, Canonical, CanonicalParamEnvCacheEntry, CanonicalTyVarKind, CanonicalVarKind,
+    Flags, InferCtxtLike, Interner, TypeFlags, TypeFoldable, TypeFolder, TypeSuperFoldable,
+    TypeVisitableExt,
 };
 
 use crate::delegate::SolverDelegate;
@@ -100,6 +101,76 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
         Canonical { max_universe, variables, value }
     }
 
+    fn canonicalize_param_env(
+        delegate: &'a D,
+        variables: &'a mut Vec<I::GenericArg>,
+        param_env: I::ParamEnv,
+    ) -> (I::ParamEnv, HashMap<I::GenericArg, usize>, Vec<CanonicalVarKind<I>>) {
+        if !param_env.has_type_flags(NEEDS_CANONICAL) {
+            return (param_env, Default::default(), Vec::new());
+        }
+
+        // Check whether we can use the global cache for this param_env. As we only use
+        // the `param_env` itself as the cache key, considering any additional information
+        // durnig its canonicalization would be incorrect. We always canonicalize region
+        // inference variables in a separate universe, so these are fine. However, we do
+        // track the universe of type and const inference variables so these must not be
+        // globally cached. We don't rely on any additional information when canonicalizing
+        // placeholders.
+        if !param_env.has_non_region_infer() {
+            delegate.cx().canonical_param_env_cache_get_or_insert(
+                param_env,
+                || {
+                    let mut variables = Vec::new();
+                    let mut env_canonicalizer = Canonicalizer {
+                        delegate,
+                        canonicalize_mode: CanonicalizeMode::Input { keep_static: true },
+
+                        variables: &mut variables,
+                        variable_lookup_table: Default::default(),
+                        var_kinds: Vec::new(),
+                        binder_index: ty::INNERMOST,
+
+                        cache: Default::default(),
+                    };
+                    let param_env = param_env.fold_with(&mut env_canonicalizer);
+                    debug_assert_eq!(env_canonicalizer.binder_index, ty::INNERMOST);
+                    CanonicalParamEnvCacheEntry {
+                        param_env,
+                        variable_lookup_table: env_canonicalizer.variable_lookup_table,
+                        var_kinds: env_canonicalizer.var_kinds,
+                        variables,
+                    }
+                },
+                |&CanonicalParamEnvCacheEntry {
+                     param_env,
+                     variables: ref cache_variables,
+                     ref variable_lookup_table,
+                     ref var_kinds,
+                 }| {
+                    debug_assert!(variables.is_empty());
+                    variables.extend(cache_variables.iter().copied());
+                    (param_env, variable_lookup_table.clone(), var_kinds.clone())
+                },
+            )
+        } else {
+            let mut env_canonicalizer = Canonicalizer {
+                delegate,
+                canonicalize_mode: CanonicalizeMode::Input { keep_static: true },
+
+                variables,
+                variable_lookup_table: Default::default(),
+                var_kinds: Vec::new(),
+                binder_index: ty::INNERMOST,
+
+                cache: Default::default(),
+            };
+            let param_env = param_env.fold_with(&mut env_canonicalizer);
+            debug_assert_eq!(env_canonicalizer.binder_index, ty::INNERMOST);
+            (param_env, env_canonicalizer.variable_lookup_table, env_canonicalizer.var_kinds)
+        }
+    }
+
     /// When canonicalizing query inputs, we keep `'static` in the `param_env`
     /// but erase it everywhere else. We generally don't want to depend on region
     /// identity, so while it should not matter whether `'static` is kept in the
@@ -114,37 +185,17 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
         input: QueryInput<I, P>,
     ) -> ty::Canonical<I, QueryInput<I, P>> {
         // First canonicalize the `param_env` while keeping `'static`
-        let mut env_canonicalizer = Canonicalizer {
-            delegate,
-            canonicalize_mode: CanonicalizeMode::Input { keep_static: true },
-
-            variables,
-            variable_lookup_table: Default::default(),
-            var_kinds: Vec::new(),
-            binder_index: ty::INNERMOST,
-
-            cache: Default::default(),
-        };
-
-        let param_env = input.goal.param_env;
-        let param_env = if param_env.has_type_flags(NEEDS_CANONICAL) {
-            param_env.fold_with(&mut env_canonicalizer)
-        } else {
-            param_env
-        };
-
-        debug_assert_eq!(env_canonicalizer.binder_index, ty::INNERMOST);
+        let (param_env, variable_lookup_table, var_kinds) =
+            Canonicalizer::canonicalize_param_env(delegate, variables, input.goal.param_env);
         // Then canonicalize the rest of the input without keeping `'static`
         // while *mostly* reusing the canonicalizer from above.
         let mut rest_canonicalizer = Canonicalizer {
             delegate,
             canonicalize_mode: CanonicalizeMode::Input { keep_static: false },
 
-            variables: env_canonicalizer.variables,
-            // We're able to reuse the `variable_lookup_table` as whether or not
-            // it already contains an entry for `'static` does not matter.
-            variable_lookup_table: env_canonicalizer.variable_lookup_table,
-            var_kinds: env_canonicalizer.var_kinds,
+            variables,
+            variable_lookup_table,
+            var_kinds,
             binder_index: ty::INNERMOST,
 
             // We do not reuse the cache as it may contain entries whose canonicalized

compiler/rustc_type_ir/src/canonical.rs

@@ -7,6 +7,7 @@ use derive_where::derive_where;
 use rustc_macros::{Decodable_NoContext, Encodable_NoContext, HashStable_NoContext};
 use rustc_type_ir_macros::{Lift_Generic, TypeFoldable_Generic, TypeVisitable_Generic};
 
+use crate::data_structures::HashMap;
 use crate::inherent::*;
 use crate::{self as ty, Interner, TypingMode, UniverseIndex};
 
@@ -333,3 +334,11 @@ impl<I: Interner> Index<ty::BoundVar> for CanonicalVarValues<I> {
         &self.var_values.as_slice()[value.as_usize()]
     }
 }
+
+#[derive_where(Clone, Debug; I: Interner)]
+pub struct CanonicalParamEnvCacheEntry<I: Interner> {
+    pub param_env: I::ParamEnv,
+    pub variables: Vec<I::GenericArg>,
+    pub variable_lookup_table: HashMap<I::GenericArg, usize>,
+    pub var_kinds: Vec<CanonicalVarKind<I>>,
+}

compiler/rustc_type_ir/src/interner.rs

@@ -12,7 +12,7 @@ use crate::lang_items::TraitSolverLangItem;
 use crate::relate::Relate;
 use crate::solve::{CanonicalInput, ExternalConstraintsData, PredefinedOpaquesData, QueryResult};
 use crate::visit::{Flags, TypeVisitable};
-use crate::{self as ty, search_graph};
+use crate::{self as ty, CanonicalParamEnvCacheEntry, search_graph};
 
 #[cfg_attr(feature = "nightly", rustc_diagnostic_item = "type_ir_interner")]
 pub trait Interner:
@@ -149,6 +149,13 @@ pub trait Interner:
 
     fn with_global_cache<R>(self, f: impl FnOnce(&mut search_graph::GlobalCache<Self>) -> R) -> R;
 
+    fn canonical_param_env_cache_get_or_insert<R>(
+        self,
+        param_env: Self::ParamEnv,
+        f: impl FnOnce() -> CanonicalParamEnvCacheEntry<Self>,
+        from_entry: impl FnOnce(&CanonicalParamEnvCacheEntry<Self>) -> R,
+    ) -> R;
+
     fn evaluation_is_concurrent(&self) -> bool;
 
     fn expand_abstract_consts<T: TypeFoldable<Self>>(self, t: T) -> T;