canonicalizer: rm region uniquification, add caching

Author: lcnr

compiler/rustc_next_trait_solver/src/canonicalizer.rs

@@ -1,5 +1,6 @@
 use std::cmp::Ordering;
 
+use rustc_type_ir::data_structures::HashMap;
 use rustc_type_ir::fold::{TypeFoldable,TypeFolder,TypeSuperFoldable};
 use rustc_type_ir::inherent::*;
 use rustc_type_ir::visit::TypeVisitableExt;
@@ -44,8 +45,15 @@ pub struct Canonicalizer<'a, D: SolverDelegate<Interner = I>, I: Interner> {
 canonicalize_mode:CanonicalizeMode,
 
 variables:&'amutVec<I::GenericArg>,
+variable_lookup_table:HashMap<I::GenericArg,usize>,
+
 primitive_var_infos:Vec<CanonicalVarInfo<I>>,
 binder_index: ty::DebruijnIndex,
+
+/// We only use the debruijn index during lookup as all other fields
+/// should not be impacted by whether a type is folded once or multiple
+/// times.
+cache:HashMap<(ty::DebruijnIndex,I::Ty),I::Ty>,
 }
 
 impl<'a,D:SolverDelegate<Interner = I>,I:Interner>Canonicalizer<'a,D,I>{
@@ -60,12 +68,14 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
  canonicalize_mode,
 
  variables,
+variable_lookup_table:Default::default(),
 primitive_var_infos:Vec::new(),
 binder_index: ty::INNERMOST,
+
+cache:Default::default(),
 };
 
 let value = value.fold_with(&mut canonicalizer);
-// FIXME: Restore these assertions. Should we uplift type flags?
 assert!(!value.has_infer(),"unexpected infer in {value:?}");
 assert!(!value.has_placeholders(),"unexpected placeholders in {value:?}");
 
@@ -75,6 +85,37 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
 Canonical{ defining_opaque_types, max_universe, variables, value }
 }
 
+fnget_or_insert_bound_var(
+&mutself,
+arg:implInto<I::GenericArg>,
+canonical_var_info:CanonicalVarInfo<I>,
+) -> ty::BoundVar{
+// FIXME: 16 is made up and arbitrary. We should look at some
+// perf data here.
+let arg = arg.into();
+let idx = ifself.variables.len() > 16{
+ifself.variable_lookup_table.is_empty(){
+self.variable_lookup_table.extend(self.variables.iter().copied().zip(0..));
+}
+
+*self.variable_lookup_table.entry(arg).or_insert_with(|| {
+let var = self.variables.len();
+self.variables.push(arg);
+self.primitive_var_infos.push(canonical_var_info);
+ var
+})
+}else{
+self.variables.iter().position(|&v| v == arg).unwrap_or_else(|| {
+let var = self.variables.len();
+self.variables.push(arg);
+self.primitive_var_infos.push(canonical_var_info);
+ var
+})
+};
+
+ ty::BoundVar::from(idx)
+}
+
 fnfinalize(self) -> (ty::UniverseIndex,I::CanonicalVars){
 letmut var_infos = self.primitive_var_infos;
 // See the rustc-dev-guide section about how we deal with universes
@@ -124,8 +165,8 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
 // - var_infos: [E0, U1, E2, U1, E1, E6, U6], curr_compressed_uv: 2, next_orig_uv: 6
 // - var_infos: [E0, U1, E1, U1, E1, E3, U3], curr_compressed_uv: 2, next_orig_uv: -
 //
-// This algorithm runs in `O(n²)` where `n` is the number of different universe
-// indices in the input. This should be fine as `n` is expected to be small.
+// This algorithm runs in `O(mn)` where `n` is the number of different universes and
+// `m` the number of variables. This should be fine as both are expected to be small.
 letmut curr_compressed_uv = ty::UniverseIndex::ROOT;
 letmut existential_in_new_uv = None;
 letmut next_orig_uv = Some(ty::UniverseIndex::ROOT);
@@ -185,14 +226,16 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
 for var in var_infos.iter_mut(){
 // We simply put all regions from the input into the highest
 // compressed universe, so we only deal with them at the end.
-if !var.is_region() && is_existential == var.is_existential(){
-update_uv(var, orig_uv, is_existential)
+if !var.is_region(){
+if is_existential == var.is_existential(){
+update_uv(var, orig_uv, is_existential)
+}
 }
 }
 }
 }
 
-// We uniquify regions and always put them into their own universe
+// We put all regions into a separate universe.
 letmut first_region = true;
 for var in var_infos.iter_mut(){
 if var.is_region(){
@@ -208,93 +251,8 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
 let var_infos = self.delegate.cx().mk_canonical_var_infos(&var_infos);
 (curr_compressed_uv, var_infos)
 }
-}
-
-impl<D:SolverDelegate<Interner = I>,I:Interner>TypeFolder<I>forCanonicalizer<'_,D,I>{
-fncx(&self) -> I{
-self.delegate.cx()
-}
-
-fnfold_binder<T>(&mutself,t: ty::Binder<I,T>) -> ty::Binder<I,T>
-where
-T:TypeFoldable<I>,
-{
-self.binder_index.shift_in(1);
-let t = t.super_fold_with(self);
-self.binder_index.shift_out(1);
- t
-}
-
-fnfold_region(&mutself,r:I::Region) -> I::Region{
-let kind = match r.kind(){
- ty::ReBound(..) => return r,
-
-// We may encounter `ReStatic` in item signatures or the hidden type
-// of an opaque. `ReErased` should only be encountered in the hidden
-// type of an opaque for regions that are ignored for the purposes of
-// captures.
-//
-// FIXME: We should investigate the perf implications of not uniquifying
-// `ReErased`. We may be able to short-circuit registering region
-// obligations if we encounter a `ReErased` on one side, for example.
- ty::ReStatic | ty::ReErased | ty::ReError(_) => matchself.canonicalize_mode{
-CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
-CanonicalizeMode::Response{ .. } => return r,
-},
-
- ty::ReEarlyParam(_) | ty::ReLateParam(_) => matchself.canonicalize_mode{
-CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
-CanonicalizeMode::Response{ .. } => {
-panic!("unexpected region in response: {r:?}")
-}
-},
-
- ty::RePlaceholder(placeholder) => matchself.canonicalize_mode{
-// We canonicalize placeholder regions as existentials in query inputs.
-CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
-CanonicalizeMode::Response{ max_input_universe } => {
-// If we have a placeholder region inside of a query, it must be from
-// a new universe.
-if max_input_universe.can_name(placeholder.universe()){
-panic!("new placeholder in universe {max_input_universe:?}: {r:?}");
-}
-CanonicalVarKind::PlaceholderRegion(placeholder)
-}
-},
-
- ty::ReVar(vid) => {
-assert_eq!(
-self.delegate.opportunistic_resolve_lt_var(vid),
- r,
-"region vid should have been resolved fully before canonicalization"
-);
-matchself.canonicalize_mode{
-CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
-CanonicalizeMode::Response{ .. } => {
-CanonicalVarKind::Region(self.delegate.universe_of_lt(vid).unwrap())
-}
-}
-}
-};
-
-let existing_bound_var = matchself.canonicalize_mode{
-CanonicalizeMode::Input => None,
-CanonicalizeMode::Response{ .. } => {
-self.variables.iter().position(|&v| v == r.into()).map(ty::BoundVar::from)
-}
-};
-
-let var = existing_bound_var.unwrap_or_else(|| {
-let var = ty::BoundVar::from(self.variables.len());
-self.variables.push(r.into());
-self.primitive_var_infos.push(CanonicalVarInfo{ kind });
- var
-});
 
-Region::new_anon_bound(self.cx(),self.binder_index, var)
-}
-
-fnfold_ty(&mutself,t:I::Ty) -> I::Ty{
+fncached_fold_ty(&mutself,t:I::Ty) -> I::Ty{
 let kind = match t.kind(){
  ty::Infer(i) => match i {
  ty::TyVar(vid) => {
@@ -368,20 +326,98 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
  | ty::Tuple(_)
  | ty::Alias(_, _)
  | ty::Bound(_, _)
- | ty::Error(_) => return t.super_fold_with(self),
+ | ty::Error(_) => {
+return t.super_fold_with(self);
+}
 };
 
-let var = ty::BoundVar::from(
-self.variables.iter().position(|&v| v == t.into()).unwrap_or_else(|| {
-let var = self.variables.len();
-self.variables.push(t.into());
-self.primitive_var_infos.push(CanonicalVarInfo{ kind });
- var
-}),
-);
+let var = self.get_or_insert_bound_var(t,CanonicalVarInfo{ kind });
 
 Ty::new_anon_bound(self.cx(),self.binder_index, var)
 }
+}
+
+impl<D:SolverDelegate<Interner = I>,I:Interner>TypeFolder<I>forCanonicalizer<'_,D,I>{
+fncx(&self) -> I{
+self.delegate.cx()
+}
+
+fnfold_binder<T>(&mutself,t: ty::Binder<I,T>) -> ty::Binder<I,T>
+where
+T:TypeFoldable<I>,
+{
+self.binder_index.shift_in(1);
+let t = t.super_fold_with(self);
+self.binder_index.shift_out(1);
+ t
+}
+
+fnfold_region(&mutself,r:I::Region) -> I::Region{
+let kind = match r.kind(){
+ ty::ReBound(..) => return r,
+
+// We may encounter `ReStatic` in item signatures or the hidden type
+// of an opaque. `ReErased` should only be encountered in the hidden
+// type of an opaque for regions that are ignored for the purposes of
+// captures.
+//
+// FIXME: We should investigate the perf implications of not uniquifying
+// `ReErased`. We may be able to short-circuit registering region
+// obligations if we encounter a `ReErased` on one side, for example.
+ ty::ReStatic | ty::ReErased | ty::ReError(_) => matchself.canonicalize_mode{
+CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
+CanonicalizeMode::Response{ .. } => return r,
+},
+
+ ty::ReEarlyParam(_) | ty::ReLateParam(_) => matchself.canonicalize_mode{
+CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
+CanonicalizeMode::Response{ .. } => {
+panic!("unexpected region in response: {r:?}")
+}
+},
+
+ ty::RePlaceholder(placeholder) => matchself.canonicalize_mode{
+// We canonicalize placeholder regions as existentials in query inputs.
+CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
+CanonicalizeMode::Response{ max_input_universe } => {
+// If we have a placeholder region inside of a query, it must be from
+// a new universe.
+if max_input_universe.can_name(placeholder.universe()){
+panic!("new placeholder in universe {max_input_universe:?}: {r:?}");
+}
+CanonicalVarKind::PlaceholderRegion(placeholder)
+}
+},
+
+ ty::ReVar(vid) => {
+assert_eq!(
+self.delegate.opportunistic_resolve_lt_var(vid),
+ r,
+"region vid should have been resolved fully before canonicalization"
+);
+matchself.canonicalize_mode{
+CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
+CanonicalizeMode::Response{ .. } => {
+CanonicalVarKind::Region(self.delegate.universe_of_lt(vid).unwrap())
+}
+}
+}
+};
+
+let var = self.get_or_insert_bound_var(r,CanonicalVarInfo{ kind });
+
+Region::new_anon_bound(self.cx(),self.binder_index, var)
+}
+
+fnfold_ty(&mutself,t:I::Ty) -> I::Ty{
+ifletSome(&ty) = self.cache.get(&(self.binder_index, t)){
+ ty
+}else{
+let res = self.cached_fold_ty(t);
+assert!(self.cache.insert((self.binder_index, t), res).is_none());
+ res
+}
+}
 
 fnfold_const(&mutself,c:I::Const) -> I::Const{
 let kind = match c.kind(){
@@ -419,14 +455,7 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
  | ty::ConstKind::Expr(_) => return c.super_fold_with(self),
 };
 
-let var = ty::BoundVar::from(
-self.variables.iter().position(|&v| v == c.into()).unwrap_or_else(|| {
-let var = self.variables.len();
-self.variables.push(c.into());
-self.primitive_var_infos.push(CanonicalVarInfo{ kind });
- var
-}),
-);
+let var = self.get_or_insert_bound_var(c,CanonicalVarInfo{ kind });
 
 Const::new_anon_bound(self.cx(),self.binder_index, var)
 }