Skip to content

Commit 5a95672

Browse files
lcnrlcnr
committed
very cool, 10/10
1 parent 1b5aa96 commit 5a95672

File tree

2 files changed

+141
-109
lines changed

2 files changed

+141
-109
lines changed

‎compiler/rustc_next_trait_solver/src/canonicalizer.rs‎

Lines changed: 135 additions & 109 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use std::cmp::Ordering;
22

3+
use rustc_type_ir::data_structures::HashMap;
34
use rustc_type_ir::fold::{TypeFoldable, TypeFolder, TypeSuperFoldable};
45
use rustc_type_ir::inherent::*;
56
use rustc_type_ir::visit::TypeVisitableExt;
@@ -44,8 +45,12 @@ pub struct Canonicalizer<'a, D: SolverDelegate<Interner = I>, I: Interner> {
4445
canonicalize_mode: CanonicalizeMode,
4546

4647
variables: &'a mut Vec<I::GenericArg>,
48+
variable_lookup_table: HashMap<I::GenericArg, usize>,
49+
4750
primitive_var_infos: Vec<CanonicalVarInfo<I>>,
4851
binder_index: ty::DebruijnIndex,
52+
53+
cache: HashMap<(ty::DebruijnIndex, I::Ty), I::Ty>,
4954
}
5055

5156
impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
@@ -60,12 +65,14 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
6065
canonicalize_mode,
6166

6267
variables,
68+
variable_lookup_table: Default::default(),
6369
primitive_var_infos: Vec::new(),
6470
binder_index: ty::INNERMOST,
71+
72+
cache: Default::default(),
6573
};
6674

6775
let value = value.fold_with(&mut canonicalizer);
68-
// FIXME: Restore these assertions. Should we uplift type flags?
6976
assert!(!value.has_infer(), "unexpected infer in {value:?}");
7077
assert!(!value.has_placeholders(), "unexpected placeholders in {value:?}");
7178

@@ -75,6 +82,37 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
7582
Canonical { defining_opaque_types, max_universe, variables, value }
7683
}
7784

85+
fn get_or_insert_bound_var(
86+
&mut self,
87+
arg: impl Into<I::GenericArg>,
88+
canonical_var_info: CanonicalVarInfo<I>,
89+
) -> ty::BoundVar {
90+
// FIXME: 16 is made up and arbitrary. We should look at some
91+
// perf data here.
92+
let arg = arg.into();
93+
let idx = if self.variables.len() > 16 {
94+
if self.variable_lookup_table.is_empty() {
95+
self.variable_lookup_table.extend(self.variables.iter().copied().zip(0..));
96+
}
97+
98+
*self.variable_lookup_table.entry(arg).or_insert_with(|| {
99+
let var = self.variables.len();
100+
self.variables.push(arg);
101+
self.primitive_var_infos.push(canonical_var_info);
102+
var
103+
})
104+
} else {
105+
self.variables.iter().position(|&v| v == arg).unwrap_or_else(|| {
106+
let var = self.variables.len();
107+
self.variables.push(arg);
108+
self.primitive_var_infos.push(canonical_var_info);
109+
var
110+
})
111+
};
112+
113+
ty::BoundVar::from(idx)
114+
}
115+
78116
fn finalize(self) -> (ty::UniverseIndex, I::CanonicalVars) {
79117
let mut var_infos = self.primitive_var_infos;
80118
// See the rustc-dev-guide section about how we deal with universes
@@ -124,8 +162,8 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
124162
// - var_infos: [E0, U1, E2, U1, E1, E6, U6], curr_compressed_uv: 2, next_orig_uv: 6
125163
// - var_infos: [E0, U1, E1, U1, E1, E3, U3], curr_compressed_uv: 2, next_orig_uv: -
126164
//
127-
// This algorithm runs in `O(n²)` where `n` is the number of different universe
128-
// indices in the input. This should be fine as `n` is expected to be small.
165+
// This algorithm runs in `O(mn)` where `n` is the number of different universes and
166+
// `m` the number of variables. This should be fine as both are expected to be small.
129167
let mut curr_compressed_uv = ty::UniverseIndex::ROOT;
130168
let mut existential_in_new_uv = None;
131169
let mut next_orig_uv = Some(ty::UniverseIndex::ROOT);
@@ -185,14 +223,16 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
185223
for var in var_infos.iter_mut() {
186224
// We simply put all regions from the input into the highest
187225
// compressed universe, so we only deal with them at the end.
188-
if !var.is_region() && is_existential == var.is_existential() {
189-
update_uv(var, orig_uv, is_existential)
226+
if !var.is_region() {
227+
if is_existential == var.is_existential() {
228+
update_uv(var, orig_uv, is_existential)
229+
}
190230
}
191231
}
192232
}
193233
}
194234

195-
// We uniquify regions and always put them into their own universe
235+
// We put all regions into a separate universe.
196236
let mut first_region = true;
197237
for var in var_infos.iter_mut() {
198238
if var.is_region() {
@@ -208,93 +248,8 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
208248
let var_infos = self.delegate.cx().mk_canonical_var_infos(&var_infos);
209249
(curr_compressed_uv, var_infos)
210250
}
211-
}
212-
213-
impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicalizer<'_, D, I> {
214-
fn cx(&self) -> I {
215-
self.delegate.cx()
216-
}
217-
218-
fn fold_binder<T>(&mut self, t: ty::Binder<I, T>) -> ty::Binder<I, T>
219-
where
220-
T: TypeFoldable<I>,
221-
{
222-
self.binder_index.shift_in(1);
223-
let t = t.super_fold_with(self);
224-
self.binder_index.shift_out(1);
225-
t
226-
}
227-
228-
fn fold_region(&mut self, r: I::Region) -> I::Region {
229-
let kind = match r.kind() {
230-
ty::ReBound(..) => return r,
231-
232-
// We may encounter `ReStatic` in item signatures or the hidden type
233-
// of an opaque. `ReErased` should only be encountered in the hidden
234-
// type of an opaque for regions that are ignored for the purposes of
235-
// captures.
236-
//
237-
// FIXME: We should investigate the perf implications of not uniquifying
238-
// `ReErased`. We may be able to short-circuit registering region
239-
// obligations if we encounter a `ReErased` on one side, for example.
240-
ty::ReStatic | ty::ReErased | ty::ReError(_) => match self.canonicalize_mode {
241-
CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
242-
CanonicalizeMode::Response { .. } => return r,
243-
},
244-
245-
ty::ReEarlyParam(_) | ty::ReLateParam(_) => match self.canonicalize_mode {
246-
CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
247-
CanonicalizeMode::Response { .. } => {
248-
panic!("unexpected region in response: {r:?}")
249-
}
250-
},
251-
252-
ty::RePlaceholder(placeholder) => match self.canonicalize_mode {
253-
// We canonicalize placeholder regions as existentials in query inputs.
254-
CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
255-
CanonicalizeMode::Response { max_input_universe } => {
256-
// If we have a placeholder region inside of a query, it must be from
257-
// a new universe.
258-
if max_input_universe.can_name(placeholder.universe()) {
259-
panic!("new placeholder in universe {max_input_universe:?}: {r:?}");
260-
}
261-
CanonicalVarKind::PlaceholderRegion(placeholder)
262-
}
263-
},
264-
265-
ty::ReVar(vid) => {
266-
assert_eq!(
267-
self.delegate.opportunistic_resolve_lt_var(vid),
268-
r,
269-
"region vid should have been resolved fully before canonicalization"
270-
);
271-
match self.canonicalize_mode {
272-
CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
273-
CanonicalizeMode::Response { .. } => {
274-
CanonicalVarKind::Region(self.delegate.universe_of_lt(vid).unwrap())
275-
}
276-
}
277-
}
278-
};
279-
280-
let existing_bound_var = match self.canonicalize_mode {
281-
CanonicalizeMode::Input => None,
282-
CanonicalizeMode::Response { .. } => {
283-
self.variables.iter().position(|&v| v == r.into()).map(ty::BoundVar::from)
284-
}
285-
};
286-
287-
let var = existing_bound_var.unwrap_or_else(|| {
288-
let var = ty::BoundVar::from(self.variables.len());
289-
self.variables.push(r.into());
290-
self.primitive_var_infos.push(CanonicalVarInfo { kind });
291-
var
292-
});
293251

294-
Region::new_anon_bound(self.cx(), self.binder_index, var)
295-
}
296-
297-
fn fold_ty(&mut self, t: I::Ty) -> I::Ty {
252+
fn cached_fold_ty(&mut self, t: I::Ty) -> I::Ty {
298253
let kind = match t.kind() {
299254
ty::Infer(i) => match i {
300255
ty::TyVar(vid) => {
@@ -368,20 +323,98 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
368323
| ty::Tuple(_)
369324
| ty::Alias(_, _)
370325
| ty::Bound(_, _)
371-
| ty::Error(_) => return t.super_fold_with(self),
326+
| ty::Error(_) => {
327+
return t.super_fold_with(self);
328+
}
372329
};
373330

374-
let var = ty::BoundVar::from(
375-
self.variables.iter().position(|&v| v == t.into()).unwrap_or_else(|| {
376-
let var = self.variables.len();
377-
self.variables.push(t.into());
378-
self.primitive_var_infos.push(CanonicalVarInfo { kind });
379-
var
380-
}),
381-
);
331+
let var = self.get_or_insert_bound_var(t, CanonicalVarInfo { kind });
382332

383333
Ty::new_anon_bound(self.cx(), self.binder_index, var)
384334
}
335+
}
336+
337+
impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicalizer<'_, D, I> {
338+
fn cx(&self) -> I {
339+
self.delegate.cx()
340+
}
341+
342+
fn fold_binder<T>(&mut self, t: ty::Binder<I, T>) -> ty::Binder<I, T>
343+
where
344+
T: TypeFoldable<I>,
345+
{
346+
self.binder_index.shift_in(1);
347+
let t = t.super_fold_with(self);
348+
self.binder_index.shift_out(1);
349+
t
350+
}
351+
352+
fn fold_region(&mut self, r: I::Region) -> I::Region {
353+
let kind = match r.kind() {
354+
ty::ReBound(..) => return r,
355+
356+
// We may encounter `ReStatic` in item signatures or the hidden type
357+
// of an opaque. `ReErased` should only be encountered in the hidden
358+
// type of an opaque for regions that are ignored for the purposes of
359+
// captures.
360+
//
361+
// FIXME: We should investigate the perf implications of not uniquifying
362+
// `ReErased`. We may be able to short-circuit registering region
363+
// obligations if we encounter a `ReErased` on one side, for example.
364+
ty::ReStatic | ty::ReErased | ty::ReError(_) => match self.canonicalize_mode {
365+
CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
366+
CanonicalizeMode::Response { .. } => return r,
367+
},
368+
369+
ty::ReEarlyParam(_) | ty::ReLateParam(_) => match self.canonicalize_mode {
370+
CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
371+
CanonicalizeMode::Response { .. } => {
372+
panic!("unexpected region in response: {r:?}")
373+
}
374+
},
375+
376+
ty::RePlaceholder(placeholder) => match self.canonicalize_mode {
377+
// We canonicalize placeholder regions as existentials in query inputs.
378+
CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
379+
CanonicalizeMode::Response { max_input_universe } => {
380+
// If we have a placeholder region inside of a query, it must be from
381+
// a new universe.
382+
if max_input_universe.can_name(placeholder.universe()) {
383+
panic!("new placeholder in universe {max_input_universe:?}: {r:?}");
384+
}
385+
CanonicalVarKind::PlaceholderRegion(placeholder)
386+
}
387+
},
388+
389+
ty::ReVar(vid) => {
390+
assert_eq!(
391+
self.delegate.opportunistic_resolve_lt_var(vid),
392+
r,
393+
"region vid should have been resolved fully before canonicalization"
394+
);
395+
match self.canonicalize_mode {
396+
CanonicalizeMode::Input => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
397+
CanonicalizeMode::Response { .. } => {
398+
CanonicalVarKind::Region(self.delegate.universe_of_lt(vid).unwrap())
399+
}
400+
}
401+
}
402+
};
403+
404+
let var = self.get_or_insert_bound_var(r, CanonicalVarInfo { kind });
405+
406+
Region::new_anon_bound(self.cx(), self.binder_index, var)
407+
}
408+
409+
fn fold_ty(&mut self, t: I::Ty) -> I::Ty {
410+
if let Some(&ty) = self.cache.get(&(self.binder_index, t)) {
411+
ty
412+
} else {
413+
let res = self.cached_fold_ty(t);
414+
assert!(self.cache.insert((self.binder_index, t), res).is_none());
415+
res
416+
}
417+
}
385418

386419
fn fold_const(&mut self, c: I::Const) -> I::Const {
387420
let kind = match c.kind() {
@@ -419,14 +452,7 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
419452
| ty::ConstKind::Expr(_) => return c.super_fold_with(self),
420453
};
421454

422-
let var = ty::BoundVar::from(
423-
self.variables.iter().position(|&v| v == c.into()).unwrap_or_else(|| {
424-
let var = self.variables.len();
425-
self.variables.push(c.into());
426-
self.primitive_var_infos.push(CanonicalVarInfo { kind });
427-
var
428-
}),
429-
);
455+
let var = self.get_or_insert_bound_var(c, CanonicalVarInfo { kind });
430456

431457
Const::new_anon_bound(self.cx(), self.binder_index, var)
432458
}

‎compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -362,6 +362,12 @@ where
362362
goal: Goal<I, I::Predicate>,
363363
) -> Result<(NestedNormalizationGoals<I>, bool, Certainty), NoSolution> {
364364
let (orig_values, canonical_goal) = self.canonicalize_goal(goal);
365+
366+
let num_orig_values = orig_values.iter().filter(|v| v.as_region().is_none()).count();
367+
if num_orig_values >= 128 {
368+
println!("bail due to too many orig_values: {num_orig_values}");
369+
return Ok((Default::default(), false, Certainty::overflow(false)));
370+
}
365371
let mut goal_evaluation =
366372
self.inspect.new_goal_evaluation(goal, &orig_values, goal_evaluation_kind);
367373
let canonical_response = EvalCtxt::evaluate_canonical_goal(

0 commit comments

Comments
 (0)