From 332098c53c295a89e595078b0dafb0c2e464e3e9 Mon Sep 17 00:00:00 2001 From: lcnr Date: Fri, 16 Aug 2024 15:27:38 +0200 Subject: [PATCH] add canonicalizer lookup table --- .../src/canonicalizer.rs | 77 ++++++++++++------- 1 file changed, 49 insertions(+), 28 deletions(-) diff --git a/compiler/rustc_next_trait_solver/src/canonicalizer.rs b/compiler/rustc_next_trait_solver/src/canonicalizer.rs index 394518daa4250..fcf36ff5d04ef 100644 --- a/compiler/rustc_next_trait_solver/src/canonicalizer.rs +++ b/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,6 +45,8 @@ pub struct Canonicalizer<'a, D: SolverDelegate, I: Interner> { canonicalize_mode: CanonicalizeMode, variables: &'a mut Vec, + variable_lookup_table: HashMap, + primitive_var_infos: Vec>, binder_index: ty::DebruijnIndex, } @@ -60,6 +63,7 @@ impl<'a, D: SolverDelegate, I: Interner> Canonicalizer<'a, D, I> { canonicalize_mode, variables, + variable_lookup_table: Default::default(), primitive_var_infos: Vec::new(), binder_index: ty::INNERMOST, }; @@ -75,6 +79,37 @@ impl<'a, D: SolverDelegate, I: Interner> Canonicalizer<'a, D, I> { Canonical { defining_opaque_types, max_universe, variables, value } } + fn get_or_insert_bound_var( + &mut self, + arg: impl Into, + canonical_var_info: CanonicalVarInfo, + ) -> ty::BoundVar { + // FIXME: 16 is made up and arbitrary. We should look at some + // perf data here. + let arg = arg.into(); + let idx = if self.variables.len() > 16 { + if self.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) + } + fn finalize(self) -> (ty::UniverseIndex, I::CanonicalVars) { let mut var_infos = self.primitive_var_infos; // See the rustc-dev-guide section about how we deal with universes @@ -124,8 +159,8 @@ impl<'a, D: SolverDelegate, 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. let mut curr_compressed_uv = ty::UniverseIndex::ROOT; let mut existential_in_new_uv = None; let mut next_orig_uv = Some(ty::UniverseIndex::ROOT); @@ -279,20 +314,20 @@ impl, I: Interner> TypeFolder for Canonicaliz } }; - let existing_bound_var = match self.canonicalize_mode { - CanonicalizeMode::Input => None, + let var = match self.canonicalize_mode { + CanonicalizeMode::Input => { + // It's fine to not add `r` to the lookup table, as we will + // never lookup regions when canonicalizing inputs. + let var = ty::BoundVar::from(self.variables.len()); + self.variables.push(r.into()); + self.primitive_var_infos.push(CanonicalVarInfo { kind }); + var + } CanonicalizeMode::Response { .. } => { - self.variables.iter().position(|&v| v == r.into()).map(ty::BoundVar::from) + self.get_or_insert_bound_var(r, CanonicalVarInfo { kind }) } }; - 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) } @@ -373,14 +408,7 @@ impl, I: Interner> TypeFolder for Canonicaliz | 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) } @@ -421,14 +449,7 @@ impl, I: Interner> TypeFolder 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) }