Auto merge of #117749 - aliemjay:perf-canon-cache, r=<try>

[perf] canonicalize param env only once

Canonicalize ParamEnv only once and store it. Then whenever we try to canonicalize `ParamEnvAnd<'tcx, T>` we only have to canonicalize `T` and then merge the results.

Prelimiary results show ~3-4% savings in diesel and serde benchmarks.

r? `@ghost`

compiler/rustc_borrowck/src/type_check/canonical.rs

@@ -38,7 +38,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
         let old_universe = self.infcx.universe();
 
         let TypeOpOutput { output, constraints, error_info } =
-            op.fully_perform(self.infcx, locations.span(self.body))?;
+            op.fully_perform(self.infcx, locations.span(self.body), &self.canon_cache)?;
 
         debug!(?output, ?constraints);
 

compiler/rustc_borrowck/src/type_check/free_region_relations.rs

@@ -11,7 +11,7 @@ use rustc_middle::ty::{self, RegionVid, Ty};
 use rustc_span::{ErrorGuaranteed, Span, DUMMY_SP};
 use rustc_trait_selection::traits::query::type_op::{self, TypeOp};
 use std::rc::Rc;
-use type_op::TypeOpOutput;
+use type_op::{CanonCache, TypeOpOutput};
 
 use crate::{
     type_check::constraint_conversion,
@@ -50,13 +50,15 @@ pub(crate) struct CreateResult<'tcx> {
 pub(crate) fn create<'tcx>(
     infcx: &InferCtxt<'tcx>,
     param_env: ty::ParamEnv<'tcx>,
+    canon_cache: &CanonCache<'tcx>,
     implicit_region_bound: ty::Region<'tcx>,
     universal_regions: &Rc<UniversalRegions<'tcx>>,
     constraints: &mut MirTypeckRegionConstraints<'tcx>,
 ) -> CreateResult<'tcx> {
     UniversalRegionRelationsBuilder {
         infcx,
         param_env,
+        canon_cache,
         implicit_region_bound,
         constraints,
         universal_regions: universal_regions.clone(),
@@ -174,6 +176,7 @@ impl UniversalRegionRelations<'_> {
 struct UniversalRegionRelationsBuilder<'this, 'tcx> {
     infcx: &'this InferCtxt<'tcx>,
     param_env: ty::ParamEnv<'tcx>,
+    canon_cache: &'this CanonCache<'tcx>,
     universal_regions: Rc<UniversalRegions<'tcx>>,
     implicit_region_bound: ty::Region<'tcx>,
     constraints: &'this mut MirTypeckRegionConstraints<'tcx>,
@@ -243,7 +246,7 @@ impl<'tcx> UniversalRegionRelationsBuilder<'_, 'tcx> {
             let TypeOpOutput { output: norm_ty, constraints: constraints_normalize, .. } = self
                 .param_env
                 .and(type_op::normalize::Normalize::new(ty))
-                .fully_perform(self.infcx, span)
+                .fully_perform(self.infcx, span, self.canon_cache)
                 .unwrap_or_else(|guar| TypeOpOutput {
                     output: Ty::new_error(self.infcx.tcx, guar),
                     constraints: None,
@@ -317,7 +320,7 @@ impl<'tcx> UniversalRegionRelationsBuilder<'_, 'tcx> {
         let TypeOpOutput { output: bounds, constraints, .. } = self
             .param_env
             .and(type_op::implied_outlives_bounds::ImpliedOutlivesBounds { ty })
-            .fully_perform(self.infcx, DUMMY_SP)
+            .fully_perform(self.infcx, DUMMY_SP, self.canon_cache)
             .map_err(|_: ErrorGuaranteed| debug!("failed to compute implied bounds {:?}", ty))
             .ok()?;
         debug!(?bounds, ?constraints);

compiler/rustc_borrowck/src/type_check/liveness/trace.rs

@@ -642,11 +642,11 @@ impl<'tcx> LivenessContext<'_, '_, '_, 'tcx> {
     ) -> DropData<'tcx> {
         debug!("compute_drop_data(dropped_ty={:?})", dropped_ty,);
 
-        match typeck
-            .param_env
-            .and(DropckOutlives::new(dropped_ty))
-            .fully_perform(typeck.infcx, DUMMY_SP)
-        {
+        match typeck.param_env.and(DropckOutlives::new(dropped_ty)).fully_perform(
+            typeck.infcx,
+            DUMMY_SP,
+            &typeck.canon_cache,
+        ) {
             Ok(TypeOpOutput { output, constraints, .. }) => {
                 DropData { dropck_result: output, region_constraint_data: constraints }
             }

compiler/rustc_borrowck/src/type_check/mod.rs

@@ -16,7 +16,7 @@ use rustc_hir::def_id::LocalDefId;
 use rustc_hir::lang_items::LangItem;
 use rustc_index::bit_set::SparseBitMatrix;
 use rustc_index::{IndexSlice, IndexVec};
-use rustc_infer::infer::canonical::QueryRegionConstraints;
+use rustc_infer::infer::canonical::{OriginalQueryValues, QueryRegionConstraints};
 use rustc_infer::infer::outlives::env::RegionBoundPairs;
 use rustc_infer::infer::region_constraints::RegionConstraintData;
 use rustc_infer::infer::type_variable::{TypeVariableOrigin, TypeVariableOriginKind};
@@ -43,7 +43,7 @@ use rustc_span::{Span, DUMMY_SP};
 use rustc_target::abi::{FieldIdx, FIRST_VARIANT};
 use rustc_trait_selection::traits::query::type_op::custom::scrape_region_constraints;
 use rustc_trait_selection::traits::query::type_op::custom::CustomTypeOp;
-use rustc_trait_selection::traits::query::type_op::{TypeOp, TypeOpOutput};
+use rustc_trait_selection::traits::query::type_op::{CanonCache, TypeOp, TypeOpOutput};
 
 use rustc_trait_selection::traits::PredicateObligation;
 
@@ -152,13 +152,19 @@ pub(crate) fn type_check<'mir, 'tcx>(
         universe_causes: FxIndexMap::default(),
     };
 
+    let canon_cache = {
+        let mut var_values = OriginalQueryValues::default();
+        Some((infcx.canonicalize_query_keep_static(param_env, &mut var_values), var_values))
+    };
+
     let CreateResult {
         universal_region_relations,
         region_bound_pairs,
         normalized_inputs_and_output,
     } = free_region_relations::create(
         infcx,
         param_env,
+        &canon_cache,
         implicit_region_bound,
         universal_regions,
         &mut constraints,
@@ -183,6 +189,7 @@ pub(crate) fn type_check<'mir, 'tcx>(
         infcx,
         body,
         param_env,
+        canon_cache,
         &region_bound_pairs,
         implicit_region_bound,
         &mut borrowck_context,
@@ -840,6 +847,7 @@ impl<'a, 'b, 'tcx> TypeVerifier<'a, 'b, 'tcx> {
 struct TypeChecker<'a, 'tcx> {
     infcx: &'a BorrowckInferCtxt<'a, 'tcx>,
     param_env: ty::ParamEnv<'tcx>,
+    canon_cache: CanonCache<'tcx>,
     last_span: Span,
     body: &'a Body<'tcx>,
     /// User type annotations are shared between the main MIR and the MIR of
@@ -1002,6 +1010,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
         infcx: &'a BorrowckInferCtxt<'a, 'tcx>,
         body: &'a Body<'tcx>,
         param_env: ty::ParamEnv<'tcx>,
+        canon_cache: CanonCache<'tcx>,
         region_bound_pairs: &'a RegionBoundPairs<'tcx>,
         implicit_region_bound: ty::Region<'tcx>,
         borrowck_context: &'a mut BorrowCheckContext<'a, 'tcx>,
@@ -1012,6 +1021,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
             body,
             user_type_annotations: &body.user_type_annotations,
             param_env,
+            canon_cache,
             region_bound_pairs,
             implicit_region_bound,
             borrowck_context,
@@ -2764,6 +2774,7 @@ impl<'tcx> TypeOp<'tcx> for InstantiateOpaqueType<'tcx> {
         mut self,
         infcx: &InferCtxt<'tcx>,
         span: Span,
+        _: &CanonCache<'tcx>,
     ) -> Result<TypeOpOutput<'tcx, Self>, ErrorGuaranteed> {
         let (mut output, region_constraints) = scrape_region_constraints(
             infcx,

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

@@ -13,6 +13,7 @@ use rustc_middle::ty::flags::FlagComputation;
 use rustc_middle::ty::fold::{TypeFoldable, TypeFolder, TypeSuperFoldable};
 use rustc_middle::ty::GenericArg;
 use rustc_middle::ty::{self, BoundVar, InferConst, List, Ty, TyCtxt, TypeFlags, TypeVisitableExt};
+use std::borrow::Cow;
 use std::sync::atomic::Ordering;
 
 use rustc_data_structures::fx::FxHashMap;
@@ -148,6 +149,27 @@ impl<'tcx> InferCtxt<'tcx> {
             query_state,
         )
     }
+
+    pub fn canonicalize_query_keep_static_continue<U, V>(
+        &self,
+        base: Canonical<'tcx, U>,
+        value: V,
+        query_state: &mut Cow<'_, OriginalQueryValues<'tcx>>,
+    ) -> Canonical<'tcx, (U, V)>
+    where
+        V: TypeFoldable<TyCtxt<'tcx>>,
+    {
+        self.tcx.sess.perf_stats.queries_canonicalized.fetch_add(1, Ordering::Relaxed);
+
+        Canonicalizer::canonicalize_continue(
+            base,
+            value,
+            self,
+            self.tcx,
+            &CanonicalizeFreeRegionsOtherThanStatic,
+            query_state,
+        )
+    }
 }
 
 /// Controls how we canonicalize "free regions" that are not inference
@@ -616,6 +638,66 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
         Canonical { max_universe, variables: canonical_variables, value: out_value }
     }
 
+    fn canonicalize_continue<U, V>(
+        base: Canonical<'tcx, U>,
+        value: V,
+        infcx: &InferCtxt<'tcx>,
+        tcx: TyCtxt<'tcx>,
+        canonicalize_region_mode: &dyn CanonicalizeMode,
+        query_state: &mut Cow<'_, OriginalQueryValues<'tcx>>,
+    ) -> Canonical<'tcx, (U, V)>
+    where
+        V: TypeFoldable<TyCtxt<'tcx>>,
+    {
+        let needs_canonical_flags = if canonicalize_region_mode.any() {
+            TypeFlags::HAS_INFER | TypeFlags::HAS_PLACEHOLDER | TypeFlags::HAS_FREE_REGIONS
+        } else {
+            TypeFlags::HAS_INFER | TypeFlags::HAS_PLACEHOLDER
+        };
+
+        // Fast path: nothing that needs to be canonicalized.
+        if !value.has_type_flags(needs_canonical_flags) {
+            return base.unchecked_map(|b| (b, value));
+        }
+
+        let mut canonicalizer = Canonicalizer {
+            infcx,
+            tcx,
+            canonicalize_mode: canonicalize_region_mode,
+            needs_canonical_flags,
+            variables: SmallVec::from_slice(base.variables),
+            query_state: query_state.to_mut(),
+            indices: FxHashMap::default(),
+            binder_index: ty::INNERMOST,
+        };
+        if canonicalizer.query_state.var_values.spilled() {
+            canonicalizer.indices = canonicalizer
+                .query_state
+                .var_values
+                .iter()
+                .enumerate()
+                .map(|(i, &kind)| (kind, BoundVar::new(i)))
+                .collect();
+        }
+        let out_value = value.fold_with(&mut canonicalizer);
+
+        // Once we have canonicalized `out_value`, it should not
+        // contain anything that ties it to this inference context
+        // anymore.
+        debug_assert!(!out_value.has_infer() && !out_value.has_placeholders());
+
+        let canonical_variables =
+            tcx.mk_canonical_var_infos(&canonicalizer.universe_canonicalized_variables());
+
+        let max_universe = canonical_variables
+            .iter()
+            .map(|cvar| cvar.universe())
+            .max()
+            .unwrap_or(ty::UniverseIndex::ROOT);
+
+        Canonical { max_universe, variables: canonical_variables, value: (base.value, out_value) }
+    }
+
     /// Creates a canonical variable replacing `kind` from the input,
     /// or returns an existing variable if `kind` has already been
     /// seen. `kind` is expected to be an unbound variable (or

compiler/rustc_trait_selection/src/traits/query/type_op/custom.rs

@@ -1,6 +1,6 @@
 use crate::infer::canonical::query_response;
 use crate::infer::InferCtxt;
-use crate::traits::query::type_op::TypeOpOutput;
+use crate::traits::query::type_op::{CanonCache, TypeOpOutput};
 use crate::traits::ObligationCtxt;
 use rustc_errors::ErrorGuaranteed;
 use rustc_infer::infer::region_constraints::RegionConstraintData;
@@ -41,6 +41,7 @@ where
         self,
         infcx: &InferCtxt<'tcx>,
         span: Span,
+        _: &CanonCache<'tcx>,
     ) -> Result<TypeOpOutput<'tcx, Self>, ErrorGuaranteed> {
         if cfg!(debug_assertions) {
             info!("fully_perform({:?})", self);

compiler/rustc_trait_selection/src/traits/query/type_op/mod.rs

@@ -8,8 +8,9 @@ use rustc_infer::infer::canonical::Certainty;
 use rustc_infer::traits::PredicateObligations;
 use rustc_middle::traits::query::NoSolution;
 use rustc_middle::ty::fold::TypeFoldable;
-use rustc_middle::ty::{ParamEnvAnd, TyCtxt};
+use rustc_middle::ty::{ParamEnv, ParamEnvAnd, TyCtxt};
 use rustc_span::Span;
+use std::borrow::Cow;
 use std::fmt;
 
 pub mod ascribe_user_type;
@@ -25,6 +26,8 @@ pub use rustc_middle::traits::query::type_op::*;
 
 use self::custom::scrape_region_constraints;
 
+pub type CanonCache<'tcx> = Option<(Canonical<'tcx, ParamEnv<'tcx>>, OriginalQueryValues<'tcx>)>;
+
 /// "Type ops" are used in NLL to perform some particular action and
 /// extract out the resulting region constraints (or an error if it
 /// cannot be completed).
@@ -39,6 +42,7 @@ pub trait TypeOp<'tcx>: Sized + fmt::Debug {
         self,
         infcx: &InferCtxt<'tcx>,
         span: Span,
+        cache: &CanonCache<'tcx>,
     ) -> Result<TypeOpOutput<'tcx, Self>, ErrorGuaranteed>;
 }
 
@@ -98,6 +102,7 @@ pub trait QueryTypeOp<'tcx>: fmt::Debug + Copy + TypeFoldable<TyCtxt<'tcx>> + 't
         query_key: ParamEnvAnd<'tcx, Self>,
         infcx: &InferCtxt<'tcx>,
         output_query_region_constraints: &mut QueryRegionConstraints<'tcx>,
+        cache: &CanonCache<'tcx>,
     ) -> Result<
         (
             Self::QueryResponse,
@@ -115,10 +120,18 @@ pub trait QueryTypeOp<'tcx>: fmt::Debug + Copy + TypeFoldable<TyCtxt<'tcx>> + 't
         // `canonicalize_query_keep_static` here because of things
         // like the subtype query, which go awry around
         // `'static` otherwise.
-        let mut canonical_var_values = OriginalQueryValues::default();
+        let mut canonical_var_values;
         let old_param_env = query_key.param_env;
-        let canonical_self =
-            infcx.canonicalize_query_keep_static(query_key, &mut canonical_var_values);
+        let canonical_self = if let &Some((base, ref cache_var_values)) = cache {
+            canonical_var_values = Cow::Borrowed(cache_var_values);
+            let ParamEnvAnd { param_env: _, value } = query_key;
+            infcx
+                .canonicalize_query_keep_static_continue(base, value, &mut canonical_var_values)
+                .unchecked_map(|(param_env, value)| param_env.and(value))
+        } else {
+            canonical_var_values = Cow::Owned(OriginalQueryValues::default());
+            infcx.canonicalize_query_keep_static(query_key, canonical_var_values.to_mut())
+        };
         let canonical_result = Self::perform_query(infcx.tcx, canonical_self)?;
 
         let InferOk { value, obligations } = infcx
@@ -145,6 +158,7 @@ where
         self,
         infcx: &InferCtxt<'tcx>,
         span: Span,
+        cache: &CanonCache<'tcx>,
     ) -> Result<TypeOpOutput<'tcx, Self>, ErrorGuaranteed> {
         if infcx.next_trait_solver() {
             return Ok(scrape_region_constraints(
@@ -158,7 +172,7 @@ where
 
         let mut region_constraints = QueryRegionConstraints::default();
         let (output, error_info, mut obligations) =
-            Q::fully_perform_into(self, infcx, &mut region_constraints)
+            Q::fully_perform_into(self, infcx, &mut region_constraints, cache)
                 .map_err(|_| {
                     infcx.tcx.sess.delay_span_bug(span, format!("error performing {self:?}"))
                 })
@@ -184,6 +198,7 @@ where
                     obligation.param_env.and(ProvePredicate::new(obligation.predicate)),
                     infcx,
                     &mut region_constraints,
+                    cache,
                 ) {
                     Ok(((), _, new, certainty)) => {
                         obligations.extend(new);