From b93a98e0fca57c85bdf1ccbb96efb85d4b8d7d42 Mon Sep 17 00:00:00 2001 From: Ali MJ Al-Nasrawy Date: Thu, 9 Nov 2023 11:11:38 +0000 Subject: [PATCH] canonicalize param env only once --- .../src/type_check/canonical.rs | 2 +- .../src/type_check/free_region_relations.rs | 9 +- .../src/type_check/liveness/trace.rs | 10 +-- compiler/rustc_borrowck/src/type_check/mod.rs | 15 +++- .../src/infer/canonical/canonicalizer.rs | 82 +++++++++++++++++++ .../src/traits/query/type_op/custom.rs | 3 +- .../src/traits/query/type_op/mod.rs | 25 ++++-- 7 files changed, 129 insertions(+), 17 deletions(-) diff --git a/compiler/rustc_borrowck/src/type_check/canonical.rs b/compiler/rustc_borrowck/src/type_check/canonical.rs index fc600af1b76f0..12c9eeb06bb7d 100644 --- a/compiler/rustc_borrowck/src/type_check/canonical.rs +++ b/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); diff --git a/compiler/rustc_borrowck/src/type_check/free_region_relations.rs b/compiler/rustc_borrowck/src/type_check/free_region_relations.rs index c84256f8ff321..a2a809957afaf 100644 --- a/compiler/rustc_borrowck/src/type_check/free_region_relations.rs +++ b/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,6 +50,7 @@ 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>, constraints: &mut MirTypeckRegionConstraints<'tcx>, @@ -57,6 +58,7 @@ pub(crate) fn create<'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>, 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); diff --git a/compiler/rustc_borrowck/src/type_check/liveness/trace.rs b/compiler/rustc_borrowck/src/type_check/liveness/trace.rs index e616449ccd412..429ff3744ebeb 100644 --- a/compiler/rustc_borrowck/src/type_check/liveness/trace.rs +++ b/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 } } diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs index 9f30d9d8ba12b..d19749eed6de5 100644 --- a/compiler/rustc_borrowck/src/type_check/mod.rs +++ b/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,6 +152,11 @@ 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, @@ -159,6 +164,7 @@ pub(crate) fn type_check<'mir, 'tcx>( } = 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, ®ion_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, ErrorGuaranteed> { let (mut output, region_constraints) = scrape_region_constraints( infcx, diff --git a/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs b/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs index 0e2f9ba70fecb..78fa745214893 100644 --- a/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs +++ b/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( + &self, + base: Canonical<'tcx, U>, + value: V, + query_state: &mut Cow<'_, OriginalQueryValues<'tcx>>, + ) -> Canonical<'tcx, (U, V)> + where + V: TypeFoldable>, + { + 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( + 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>, + { + 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 diff --git a/compiler/rustc_trait_selection/src/traits/query/type_op/custom.rs b/compiler/rustc_trait_selection/src/traits/query/type_op/custom.rs index c81bc57900358..6e37f55194a1f 100644 --- a/compiler/rustc_trait_selection/src/traits/query/type_op/custom.rs +++ b/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, ErrorGuaranteed> { if cfg!(debug_assertions) { info!("fully_perform({:?})", self); diff --git a/compiler/rustc_trait_selection/src/traits/query/type_op/mod.rs b/compiler/rustc_trait_selection/src/traits/query/type_op/mod.rs index a5ccf62608edb..6479a1a5e6862 100644 --- a/compiler/rustc_trait_selection/src/traits/query/type_op/mod.rs +++ b/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, ErrorGuaranteed>; } @@ -98,6 +102,7 @@ pub trait QueryTypeOp<'tcx>: fmt::Debug + Copy + TypeFoldable> + '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> + '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, 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);