add `param_env` cache to canonicalization

1 files changed, 60 insertions, 14 deletions

diff --git a/compiler/rustc_next_trait_solver/src/canonicalizer.rs b/compiler/rustc_next_trait_solver/src/canonicalizer.rs
index 4b48ae7c417..cea77533178 100644
--- a/compiler/rustc_next_trait_solver/src/canonicalizer.rs
+++ b/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;
@@ -109,20 +110,65 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
             return (param_env, Default::default(), Vec::new());
         }
 
-        let mut env_canonicalizer = Canonicalizer {
-            delegate,
-            canonicalize_mode: CanonicalizeMode::Input { keep_static: true },
+        // 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,
+                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)
+                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`