1
Fork 0

Auto merge of #117749 - aliemjay:perf-canon-cache, r=lcnr

cache param env canonicalization

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.

Best to review commits individually. Some commits have a short description.

Initial implementation had a soundness bug (https://github.com/rust-lang/rust/pull/117749#issuecomment-1840453387) due to cache invalidation:
- When canonicalizing `Ty<'?0>` we first try to resolve region variables in the current InferCtxt which may have a constraint `?0 == 'static`. This means that we register `Ty<'?0> => Canonical<Ty<'static>>` in the cache, which is obviously incorrect in another inference context.
- This is fixed by not doing region resolution when canonicalizing the query *input* (vs. response), which is the only place where ParamEnv is used, and then in a later commit we *statically* guard against any form of inference variable resolution of the cached canonical ParamEnv's.

r? `@ghost`
This commit is contained in:
bors 2023-12-14 04:04:10 +00:00
commit d23e1a6894
8 changed files with 199 additions and 125 deletions

View file

@ -35,36 +35,13 @@ impl<'tcx> InferCtxt<'tcx> {
/// [c]: https://rust-lang.github.io/chalk/book/canonical_queries/canonicalization.html#canonicalizing-the-query
pub fn canonicalize_query<V>(
&self,
value: V,
value: ty::ParamEnvAnd<'tcx, V>,
query_state: &mut OriginalQueryValues<'tcx>,
) -> Canonical<'tcx, V>
) -> Canonical<'tcx, ty::ParamEnvAnd<'tcx, V>>
where
V: TypeFoldable<TyCtxt<'tcx>>,
{
Canonicalizer::canonicalize(value, self, self.tcx, &CanonicalizeAllFreeRegions, query_state)
}
/// Like [Self::canonicalize_query], but preserves distinct universes. For
/// example, canonicalizing `&'?0: Trait<'?1>`, where `'?0` is in `U1` and
/// `'?1` is in `U3` would be canonicalized to have `?0` in `U1` and `'?1`
/// in `U2`.
///
/// This is used for Chalk integration.
pub fn canonicalize_query_preserving_universes<V>(
&self,
value: V,
query_state: &mut OriginalQueryValues<'tcx>,
) -> Canonical<'tcx, V>
where
V: TypeFoldable<TyCtxt<'tcx>>,
{
Canonicalizer::canonicalize(
value,
self,
self.tcx,
&CanonicalizeAllFreeRegionsPreservingUniverses,
query_state,
)
self.canonicalize_query_with_mode(value, query_state, &CanonicalizeAllFreeRegions)
}
/// Canonicalizes a query *response* `V`. When we canonicalize a
@ -99,7 +76,7 @@ impl<'tcx> InferCtxt<'tcx> {
let mut query_state = OriginalQueryValues::default();
Canonicalizer::canonicalize(
value,
self,
Some(self),
self.tcx,
&CanonicalizeQueryResponse,
&mut query_state,
@ -113,7 +90,7 @@ impl<'tcx> InferCtxt<'tcx> {
let mut query_state = OriginalQueryValues::default();
Canonicalizer::canonicalize(
value,
self,
Some(self),
self.tcx,
&CanonicalizeUserTypeAnnotation,
&mut query_state,
@ -126,19 +103,53 @@ impl<'tcx> InferCtxt<'tcx> {
/// handling of `'static` regions (e.g. trait evaluation).
pub fn canonicalize_query_keep_static<V>(
&self,
value: V,
value: ty::ParamEnvAnd<'tcx, V>,
query_state: &mut OriginalQueryValues<'tcx>,
) -> Canonical<'tcx, V>
) -> Canonical<'tcx, ty::ParamEnvAnd<'tcx, V>>
where
V: TypeFoldable<TyCtxt<'tcx>>,
{
Canonicalizer::canonicalize(
self.canonicalize_query_with_mode(
value,
self,
self.tcx,
query_state,
&CanonicalizeFreeRegionsOtherThanStatic,
)
}
fn canonicalize_query_with_mode<V>(
&self,
value: ty::ParamEnvAnd<'tcx, V>,
query_state: &mut OriginalQueryValues<'tcx>,
canonicalize_region_mode: &dyn CanonicalizeMode,
) -> Canonical<'tcx, ty::ParamEnvAnd<'tcx, V>>
where
V: TypeFoldable<TyCtxt<'tcx>>,
{
let (param_env, value) = value.into_parts();
let base = self.tcx.canonical_param_env_cache.get_or_insert(
self.tcx,
param_env,
query_state,
|tcx, param_env, query_state| {
Canonicalizer::canonicalize(
param_env,
None,
tcx,
&CanonicalizeFreeRegionsOtherThanStatic,
query_state,
)
},
);
Canonicalizer::canonicalize_with_base(
base,
value,
Some(self),
self.tcx,
canonicalize_region_mode,
query_state,
)
.unchecked_map(|(param_env, value)| param_env.and(value))
}
}
@ -168,8 +179,22 @@ impl CanonicalizeMode for CanonicalizeQueryResponse {
fn canonicalize_free_region<'tcx>(
&self,
canonicalizer: &mut Canonicalizer<'_, 'tcx>,
r: ty::Region<'tcx>,
mut r: ty::Region<'tcx>,
) -> ty::Region<'tcx> {
let infcx = canonicalizer.infcx.unwrap();
if let ty::ReVar(vid) = *r {
r = infcx
.inner
.borrow_mut()
.unwrap_region_constraints()
.opportunistic_resolve_var(canonicalizer.tcx, vid);
debug!(
"canonical: region var found with vid {vid:?}, \
opportunistically resolved to {r:?}",
);
};
match *r {
ty::ReLateParam(_) | ty::ReErased | ty::ReStatic | ty::ReEarlyParam(..) => r,
@ -179,7 +204,8 @@ impl CanonicalizeMode for CanonicalizeQueryResponse {
),
ty::ReVar(vid) => {
let universe = canonicalizer.region_var_universe(vid);
let universe =
infcx.inner.borrow_mut().unwrap_region_constraints().var_universe(vid);
canonicalizer.canonical_var_for_region(
CanonicalVarInfo { kind: CanonicalVarKind::Region(universe) },
r,
@ -264,30 +290,6 @@ impl CanonicalizeMode for CanonicalizeAllFreeRegions {
}
}
struct CanonicalizeAllFreeRegionsPreservingUniverses;
impl CanonicalizeMode for CanonicalizeAllFreeRegionsPreservingUniverses {
fn canonicalize_free_region<'tcx>(
&self,
canonicalizer: &mut Canonicalizer<'_, 'tcx>,
r: ty::Region<'tcx>,
) -> ty::Region<'tcx> {
let universe = canonicalizer.infcx.universe_of_region(r);
canonicalizer.canonical_var_for_region(
CanonicalVarInfo { kind: CanonicalVarKind::Region(universe) },
r,
)
}
fn any(&self) -> bool {
true
}
fn preserve_universes(&self) -> bool {
true
}
}
struct CanonicalizeFreeRegionsOtherThanStatic;
impl CanonicalizeMode for CanonicalizeFreeRegionsOtherThanStatic {
@ -309,7 +311,8 @@ impl CanonicalizeMode for CanonicalizeFreeRegionsOtherThanStatic {
}
struct Canonicalizer<'cx, 'tcx> {
infcx: &'cx InferCtxt<'tcx>,
/// Set to `None` to disable the resolution of inference variables.
infcx: Option<&'cx InferCtxt<'tcx>>,
tcx: TyCtxt<'tcx>,
variables: SmallVec<[CanonicalVarInfo<'tcx>; 8]>,
query_state: &'cx mut OriginalQueryValues<'tcx>,
@ -347,25 +350,12 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
}
}
ty::ReVar(vid) => {
let resolved = self
.infcx
.inner
.borrow_mut()
.unwrap_region_constraints()
.opportunistic_resolve_var(self.tcx, vid);
debug!(
"canonical: region var found with vid {vid:?}, \
opportunistically resolved to {resolved:?}",
);
self.canonicalize_mode.canonicalize_free_region(self, resolved)
}
ty::ReStatic
| ty::ReEarlyParam(..)
| ty::ReError(_)
| ty::ReLateParam(_)
| ty::RePlaceholder(..)
| ty::ReVar(_)
| ty::ReErased => self.canonicalize_mode.canonicalize_free_region(self, r),
}
}
@ -376,14 +366,14 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
// We need to canonicalize the *root* of our ty var.
// This is so that our canonical response correctly reflects
// any equated inference vars correctly!
let root_vid = self.infcx.root_var(vid);
let root_vid = self.infcx.unwrap().root_var(vid);
if root_vid != vid {
t = Ty::new_var(self.infcx.tcx, root_vid);
t = Ty::new_var(self.tcx, root_vid);
vid = root_vid;
}
debug!("canonical: type var found with vid {:?}", vid);
match self.infcx.probe_ty_var(vid) {
match self.infcx.unwrap().probe_ty_var(vid) {
// `t` could be a float / int variable; canonicalize that instead.
Ok(t) => {
debug!("(resolved to {:?})", t);
@ -408,7 +398,7 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
}
ty::Infer(ty::IntVar(vid)) => {
let nt = self.infcx.opportunistic_resolve_int_var(vid);
let nt = self.infcx.unwrap().opportunistic_resolve_int_var(vid);
if nt != t {
return self.fold_ty(nt);
} else {
@ -419,7 +409,7 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
}
}
ty::Infer(ty::FloatVar(vid)) => {
let nt = self.infcx.opportunistic_resolve_float_var(vid);
let nt = self.infcx.unwrap().opportunistic_resolve_float_var(vid);
if nt != t {
return self.fold_ty(nt);
} else {
@ -490,14 +480,14 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
// We need to canonicalize the *root* of our const var.
// This is so that our canonical response correctly reflects
// any equated inference vars correctly!
let root_vid = self.infcx.root_const_var(vid);
let root_vid = self.infcx.unwrap().root_const_var(vid);
if root_vid != vid {
ct = ty::Const::new_var(self.infcx.tcx, root_vid, ct.ty());
ct = ty::Const::new_var(self.tcx, root_vid, ct.ty());
vid = root_vid;
}
debug!("canonical: const var found with vid {:?}", vid);
match self.infcx.probe_const_var(vid) {
match self.infcx.unwrap().probe_const_var(vid) {
Ok(c) => {
debug!("(resolved to {:?})", c);
return self.fold_const(c);
@ -518,8 +508,8 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
}
}
ty::ConstKind::Infer(InferConst::EffectVar(vid)) => {
match self.infcx.probe_effect_var(vid) {
Some(value) => return self.fold_const(value.as_const(self.infcx.tcx)),
match self.infcx.unwrap().probe_effect_var(vid) {
Some(value) => return self.fold_const(value.as_const(self.tcx)),
None => {
return self.canonicalize_const_var(
CanonicalVarInfo { kind: CanonicalVarKind::Effect },
@ -562,11 +552,38 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
/// `canonicalize_query` and `canonicalize_response`.
fn canonicalize<V>(
value: V,
infcx: &InferCtxt<'tcx>,
infcx: Option<&InferCtxt<'tcx>>,
tcx: TyCtxt<'tcx>,
canonicalize_region_mode: &dyn CanonicalizeMode,
query_state: &mut OriginalQueryValues<'tcx>,
) -> Canonical<'tcx, V>
where
V: TypeFoldable<TyCtxt<'tcx>>,
{
let base = Canonical {
max_universe: ty::UniverseIndex::ROOT,
variables: List::empty(),
value: (),
};
Canonicalizer::canonicalize_with_base(
base,
value,
infcx,
tcx,
canonicalize_region_mode,
query_state,
)
.unchecked_map(|((), val)| val)
}
fn canonicalize_with_base<U, V>(
base: Canonical<'tcx, U>,
value: V,
infcx: Option<&InferCtxt<'tcx>>,
tcx: TyCtxt<'tcx>,
canonicalize_region_mode: &dyn CanonicalizeMode,
query_state: &mut OriginalQueryValues<'tcx>,
) -> Canonical<'tcx, (U, V)>
where
V: TypeFoldable<TyCtxt<'tcx>>,
{
@ -578,12 +595,7 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
// Fast path: nothing that needs to be canonicalized.
if !value.has_type_flags(needs_canonical_flags) {
let canon_value = Canonical {
max_universe: ty::UniverseIndex::ROOT,
variables: List::empty(),
value,
};
return canon_value;
return base.unchecked_map(|b| (b, value));
}
let mut canonicalizer = Canonicalizer {
@ -591,11 +603,20 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
tcx,
canonicalize_mode: canonicalize_region_mode,
needs_canonical_flags,
variables: SmallVec::new(),
variables: SmallVec::from_slice(base.variables),
query_state,
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
@ -612,7 +633,7 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
.max()
.unwrap_or(ty::UniverseIndex::ROOT);
Canonical { max_universe, variables: canonical_variables, value: out_value }
Canonical { max_universe, variables: canonical_variables, value: (base.value, out_value) }
}
/// Creates a canonical variable replacing `kind` from the input,
@ -761,11 +782,6 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
)
}
/// Returns the universe in which `vid` is defined.
fn region_var_universe(&self, vid: ty::RegionVid) -> ty::UniverseIndex {
self.infcx.inner.borrow_mut().unwrap_region_constraints().var_universe(vid)
}
/// Creates a canonical variable (with the given `info`)
/// representing the region `r`; return a region referencing it.
fn canonical_var_for_region(
@ -783,14 +799,9 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
/// *that*. Otherwise, create a new canonical variable for
/// `ty_var`.
fn canonicalize_ty_var(&mut self, info: CanonicalVarInfo<'tcx>, ty_var: Ty<'tcx>) -> Ty<'tcx> {
let infcx = self.infcx;
let bound_to = infcx.shallow_resolve(ty_var);
if bound_to != ty_var {
self.fold_ty(bound_to)
} else {
let var = self.canonical_var(info, ty_var.into());
Ty::new_bound(self.tcx, self.binder_index, var.into())
}
debug_assert!(!self.infcx.is_some_and(|infcx| ty_var != infcx.shallow_resolve(ty_var)));
let var = self.canonical_var(info, ty_var.into());
Ty::new_bound(self.tcx, self.binder_index, var.into())
}
/// Given a type variable `const_var` of the given kind, first check
@ -802,13 +813,10 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
info: CanonicalVarInfo<'tcx>,
const_var: ty::Const<'tcx>,
) -> ty::Const<'tcx> {
let infcx = self.infcx;
let bound_to = infcx.shallow_resolve(const_var);
if bound_to != const_var {
self.fold_const(bound_to)
} else {
let var = self.canonical_var(info, const_var.into());
ty::Const::new_bound(self.tcx, self.binder_index, var, self.fold_ty(const_var.ty()))
}
debug_assert!(
!self.infcx.is_some_and(|infcx| const_var != infcx.shallow_resolve(const_var))
);
let var = self.canonical_var(info, const_var.into());
ty::Const::new_bound(self.tcx, self.binder_index, var, self.fold_ty(const_var.ty()))
}
}

View file

@ -172,7 +172,7 @@ impl<'tcx> InferCtxt<'tcx> {
// two const param's types are able to be equal has to go through a canonical query with the actual logic
// in `rustc_trait_selection`.
let canonical = self.canonicalize_query(
(relation.param_env(), a.ty(), b.ty()),
relation.param_env().and((a.ty(), b.ty())),
&mut OriginalQueryValues::default(),
);
self.tcx.check_tys_might_be_eq(canonical).map_err(|_| {