Skip to content

Commit 46487b0

Browse files
lcnrlcnr
committed
track relevant sub_relations in canonical queries
1 parent a1ccd9b commit 46487b0

File tree

10 files changed

+124
-77
lines changed

10 files changed

+124
-77
lines changed

‎compiler/rustc_infer/src/infer/canonical/canonicalizer.rs

+16-2
Original file line numberDiff line numberDiff line change
@@ -298,6 +298,7 @@ struct Canonicalizer<'cx, 'tcx> {
298298
// Note that indices is only used once `var_values` is big enough to be
299299
// heap-allocated.
300300
indices: FxHashMap<GenericArg<'tcx>, BoundVar>,
301+
sub_root_lookup_table: FxHashMap<ty::TyVid, usize>,
301302
canonicalize_mode: &'cx dyn CanonicalizeMode,
302303
needs_canonical_flags: TypeFlags,
303304

@@ -366,8 +367,11 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
366367
// FIXME: perf problem described in #55921.
367368
ui = ty::UniverseIndex::ROOT;
368369
}
370+
let sub_root = self.get_or_insert_sub_root(vid);
369371
self.canonicalize_ty_var(
370-
CanonicalVarInfo { kind: CanonicalVarKind::Ty(ui) },
372+
CanonicalVarInfo {
373+
kind: CanonicalVarKind::Ty { universe: ui, sub_root },
374+
},
371375
t,
372376
)
373377
}
@@ -567,6 +571,7 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
567571
variables: SmallVec::from_slice(base.variables),
568572
query_state,
569573
indices: FxHashMap::default(),
574+
sub_root_lookup_table: Default::default(),
570575
binder_index: ty::INNERMOST,
571576
};
572577
if canonicalizer.query_state.var_values.spilled() {
@@ -661,6 +666,13 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
661666
}
662667
}
663668

669+
fn get_or_insert_sub_root(&mut self, vid: ty::TyVid) -> ty::BoundVar {
670+
let root_vid = self.infcx.unwrap().sub_root_var(vid);
671+
let idx =
672+
*self.sub_root_lookup_table.entry(root_vid).or_insert_with(|| self.variables.len());
673+
ty::BoundVar::from(idx)
674+
}
675+
664676
/// Replaces the universe indexes used in `var_values` with their index in
665677
/// `query_state.universe_map`. This minimizes the maximum universe used in
666678
/// the canonicalized value.
@@ -684,7 +696,9 @@ impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {
684696
CanonicalVarKind::Int | CanonicalVarKind::Float => {
685697
return *v;
686698
}
687-
CanonicalVarKind::Ty(u) => CanonicalVarKind::Ty(reverse_universe_map[&u]),
699+
CanonicalVarKind::Ty { universe, sub_root } => {
700+
CanonicalVarKind::Ty { universe: reverse_universe_map[&universe], sub_root }
701+
}
688702
CanonicalVarKind::Region(u) => {
689703
CanonicalVarKind::Region(reverse_universe_map[&u])
690704
}

‎compiler/rustc_infer/src/infer/canonical/mod.rs

+16-7
Original file line numberDiff line numberDiff line change
@@ -84,13 +84,12 @@ impl<'tcx> InferCtxt<'tcx> {
8484
variables: &List<CanonicalVarInfo<'tcx>>,
8585
universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
8686
) -> CanonicalVarValues<'tcx> {
87-
CanonicalVarValues {
88-
var_values: self.tcx.mk_args_from_iter(
89-
variables
90-
.iter()
91-
.map(|info| self.instantiate_canonical_var(span, info, &universe_map)),
92-
),
87+
let mut var_values = Vec::new();
88+
for info in variables.iter() {
89+
let value = self.instantiate_canonical_var(span, info, &var_values, &universe_map);
90+
var_values.push(value);
9391
}
92+
CanonicalVarValues { var_values: self.tcx.mk_args(&var_values) }
9493
}
9594

9695
/// Given the "info" about a canonical variable, creates a fresh
@@ -105,10 +104,20 @@ impl<'tcx> InferCtxt<'tcx> {
105104
&self,
106105
span: Span,
107106
cv_info: CanonicalVarInfo<'tcx>,
107+
previous_var_values: &[GenericArg<'tcx>],
108108
universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
109109
) -> GenericArg<'tcx> {
110110
match cv_info.kind {
111-
CanonicalVarKind::Ty(ui) => self.next_ty_var_in_universe(span, universe_map(ui)).into(),
111+
CanonicalVarKind::Ty { universe, sub_root } => {
112+
let vid = self.next_ty_var_id_in_universe(span, universe_map(universe));
113+
if let Some(prev) = previous_var_values.get(sub_root.as_usize()) {
114+
let &ty::Infer(ty::TyVar(sub_root)) = prev.expect_ty().kind() else {
115+
unreachable!("expected `sub_root` to be an inference variable");
116+
};
117+
self.inner.borrow_mut().type_variables().sub(vid, sub_root);
118+
}
119+
Ty::new_var(self.tcx, vid).into()
120+
}
112121
CanonicalVarKind::Int => self.next_int_var().into(),
113122
CanonicalVarKind::Float => self.next_float_var().into(),
114123

‎compiler/rustc_infer/src/infer/canonical/query_response.rs

+25-26
Original file line numberDiff line numberDiff line change
@@ -455,32 +455,31 @@ impl<'tcx> InferCtxt<'tcx> {
455455
// Create result arguments: if we found a value for a
456456
// given variable in the loop above, use that. Otherwise, use
457457
// a fresh inference variable.
458-
let result_args = CanonicalVarValues {
459-
var_values: self.tcx.mk_args_from_iter(
460-
query_response.variables.iter().enumerate().map(|(index, info)| {
461-
if info.universe() != ty::UniverseIndex::ROOT {
462-
// A variable from inside a binder of the query. While ideally these shouldn't
463-
// exist at all, we have to deal with them for now.
464-
self.instantiate_canonical_var(cause.span, info, |u| {
465-
universe_map[u.as_usize()]
466-
})
467-
} else if info.is_existential() {
468-
match opt_values[BoundVar::new(index)] {
469-
Some(k) => k,
470-
None => self.instantiate_canonical_var(cause.span, info, |u| {
471-
universe_map[u.as_usize()]
472-
}),
473-
}
474-
} else {
475-
// For placeholders which were already part of the input, we simply map this
476-
// universal bound variable back the placeholder of the input.
477-
opt_values[BoundVar::new(index)].expect(
478-
"expected placeholder to be unified with itself during response",
479-
)
480-
}
481-
}),
482-
),
483-
};
458+
let mut var_values = Vec::new();
459+
for (index, info) in query_response.variables.iter().enumerate() {
460+
let value = if info.universe() != ty::UniverseIndex::ROOT {
461+
// A variable from inside a binder of the query. While ideally these shouldn't
462+
// exist at all, we have to deal with them for now.
463+
self.instantiate_canonical_var(cause.span, info, &var_values, |u| {
464+
universe_map[u.as_usize()]
465+
})
466+
} else if info.is_existential() {
467+
match opt_values[BoundVar::new(index)] {
468+
Some(k) => k,
469+
None => self.instantiate_canonical_var(cause.span, info, &var_values, |u| {
470+
universe_map[u.as_usize()]
471+
}),
472+
}
473+
} else {
474+
// For placeholders which were already part of the input, we simply map this
475+
// universal bound variable back the placeholder of the input.
476+
opt_values[BoundVar::new(index)]
477+
.expect("expected placeholder to be unified with itself during response")
478+
};
479+
var_values.push(value)
480+
}
481+
482+
let result_args = CanonicalVarValues { var_values: self.tcx.mk_args(&var_values) };
484483

485484
let mut obligations = PredicateObligations::new();
486485

‎compiler/rustc_infer/src/infer/context.rs

+3
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,9 @@ impl<'tcx> rustc_type_ir::InferCtxtLike for InferCtxt<'tcx> {
5555
fn root_ty_var(&self, var: ty::TyVid) -> ty::TyVid {
5656
self.root_var(var)
5757
}
58+
fn sub_root_ty_var(&self, var: ty::TyVid) -> ty::TyVid {
59+
self.sub_root_var(var)
60+
}
5861

5962
fn root_const_var(&self, var: ty::ConstVid) -> ty::ConstVid {
6063
self.root_const_var(var)

‎compiler/rustc_next_trait_solver/src/canonicalizer.rs

+17-5
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ pub struct Canonicalizer<'a, D: SolverDelegate<Interner = I>, I: Interner> {
5252
variables: &'a mut Vec<I::GenericArg>,
5353
primitive_var_infos: Vec<CanonicalVarInfo<I>>,
5454
variable_lookup_table: HashMap<I::GenericArg, usize>,
55+
sub_root_lookup_table: HashMap<ty::TyVid, usize>,
5556
binder_index: ty::DebruijnIndex,
5657

5758
/// We only use the debruijn index during lookup. We don't need to
@@ -73,6 +74,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
7374

7475
variables,
7576
variable_lookup_table: Default::default(),
77+
sub_root_lookup_table: Default::default(),
7678
primitive_var_infos: Vec::new(),
7779
binder_index: ty::INNERMOST,
7880

@@ -106,6 +108,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
106108

107109
variables,
108110
variable_lookup_table: Default::default(),
111+
sub_root_lookup_table: Default::default(),
109112
primitive_var_infos: Vec::new(),
110113
binder_index: ty::INNERMOST,
111114

@@ -123,6 +126,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
123126
// We're able to reuse the `variable_lookup_table` as whether or not
124127
// it already contains an entry for `'static` does not matter.
125128
variable_lookup_table: env_canonicalizer.variable_lookup_table,
129+
sub_root_lookup_table: Default::default(),
126130
primitive_var_infos: env_canonicalizer.primitive_var_infos,
127131
binder_index: ty::INNERMOST,
128132

@@ -177,6 +181,13 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
177181
ty::BoundVar::from(idx)
178182
}
179183

184+
fn get_or_insert_sub_root(&mut self, vid: ty::TyVid) -> ty::BoundVar {
185+
let root_vid = self.delegate.sub_root_ty_var(vid);
186+
let idx =
187+
*self.sub_root_lookup_table.entry(root_vid).or_insert_with(|| self.variables.len());
188+
ty::BoundVar::from(idx)
189+
}
190+
180191
fn finalize(self) -> (ty::UniverseIndex, I::CanonicalVars) {
181192
let mut var_infos = self.primitive_var_infos;
182193
// See the rustc-dev-guide section about how we deal with universes
@@ -323,11 +334,12 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
323334
"ty vid should have been resolved fully before canonicalization"
324335
);
325336

326-
CanonicalVarKind::Ty(
327-
self.delegate
328-
.universe_of_ty(vid)
329-
.unwrap_or_else(|| panic!("ty var should have been resolved: {t:?}")),
330-
)
337+
let universe = self
338+
.delegate
339+
.universe_of_ty(vid)
340+
.unwrap_or_else(|| panic!("ty var should have been resolved: {t:?}"));
341+
let sub_root = self.get_or_insert_sub_root(vid);
342+
CanonicalVarKind::Ty { universe, sub_root }
331343
}
332344
ty::IntVar(vid) => {
333345
assert_eq!(

‎compiler/rustc_next_trait_solver/src/delegate.rs

+1
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,7 @@ pub trait SolverDelegate: Deref<Target = Self::Infcx> + Sized {
5959
&self,
6060
cv_info: ty::CanonicalVarInfo<Self::Interner>,
6161
span: <Self::Interner as Interner>::Span,
62+
var_values: &[<Self::Interner as Interner>::GenericArg],
6263
universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
6364
) -> <Self::Interner as Interner>::GenericArg;
6465

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

+30-29
Original file line numberDiff line numberDiff line change
@@ -344,37 +344,38 @@ where
344344
}
345345
}
346346

347-
let var_values = delegate.cx().mk_args_from_iter(
348-
response.variables.iter().enumerate().map(|(index, info)| {
349-
if info.universe() != ty::UniverseIndex::ROOT {
350-
// A variable from inside a binder of the query. While ideally these shouldn't
351-
// exist at all (see the FIXME at the start of this method), we have to deal with
352-
// them for now.
353-
delegate.instantiate_canonical_var_with_infer(info, span, |idx| {
354-
prev_universe + idx.index()
355-
})
356-
} else if info.is_existential() {
357-
// As an optimization we sometimes avoid creating a new inference variable here.
358-
//
359-
// All new inference variables we create start out in the current universe of the caller.
360-
// This is conceptually wrong as these inference variables would be able to name
361-
// more placeholders then they should be able to. However the inference variables have
362-
// to "come from somewhere", so by equating them with the original values of the caller
363-
// later on, we pull them down into their correct universe again.
364-
if let Some(v) = opt_values[ty::BoundVar::from_usize(index)] {
365-
v
366-
} else {
367-
delegate.instantiate_canonical_var_with_infer(info, span, |_| prev_universe)
368-
}
347+
let mut var_values = Vec::new();
348+
for (index, info) in response.variables.iter().enumerate() {
349+
let value = if info.universe() != ty::UniverseIndex::ROOT {
350+
// A variable from inside a binder of the query. While ideally these shouldn't
351+
// exist at all (see the FIXME at the start of this method), we have to deal with
352+
// them for now.
353+
delegate.instantiate_canonical_var_with_infer(info, span, &var_values, |idx| {
354+
prev_universe + idx.index()
355+
})
356+
} else if info.is_existential() {
357+
// As an optimization we sometimes avoid creating a new inference variable here.
358+
//
359+
// All new inference variables we create start out in the current universe of the caller.
360+
// This is conceptually wrong as these inference variables would be able to name
361+
// more placeholders then they should be able to. However the inference variables have
362+
// to "come from somewhere", so by equating them with the original values of the caller
363+
// later on, we pull them down into their correct universe again.
364+
if let Some(v) = opt_values[ty::BoundVar::from_usize(index)] {
365+
v
369366
} else {
370-
// For placeholders which were already part of the input, we simply map this
371-
// universal bound variable back the placeholder of the input.
372-
original_values[info.expect_placeholder_index()]
367+
delegate.instantiate_canonical_var_with_infer(info, span, &var_values, |_| {
368+
prev_universe
369+
})
373370
}
374-
}),
375-
);
376-
377-
CanonicalVarValues { var_values }
371+
} else {
372+
// For placeholders which were already part of the input, we simply map this
373+
// universal bound variable back the placeholder of the input.
374+
original_values[info.expect_placeholder_index()]
375+
};
376+
var_values.push(value)
377+
}
378+
CanonicalVarValues { var_values: delegate.cx().mk_args(&var_values) }
378379
}
379380

380381
/// Unify the `original_values` with the `var_values` returned by the canonical query..

‎compiler/rustc_trait_selection/src/solve/delegate.rs

+2-1
Original file line numberDiff line numberDiff line change
@@ -148,9 +148,10 @@ impl<'tcx> rustc_next_trait_solver::delegate::SolverDelegate for SolverDelegate<
148148
&self,
149149
cv_info: CanonicalVarInfo<'tcx>,
150150
span: Span,
151+
var_values: &[ty::GenericArg<'tcx>],
151152
universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
152153
) -> ty::GenericArg<'tcx> {
153-
self.0.instantiate_canonical_var(span, cv_info, universe_map)
154+
self.0.instantiate_canonical_var(span, cv_info, var_values, universe_map)
154155
}
155156

156157
fn register_hidden_type_in_storage(

‎compiler/rustc_type_ir/src/canonical.rs

+13-7
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ impl<I: Interner> CanonicalVarInfo<I> {
110110

111111
pub fn is_existential(&self) -> bool {
112112
match self.kind {
113-
CanonicalVarKind::Ty(_) | CanonicalVarKind::Int | CanonicalVarKind::Float => true,
113+
CanonicalVarKind::Ty { .. } | CanonicalVarKind::Int | CanonicalVarKind::Float => true,
114114
CanonicalVarKind::PlaceholderTy(_) => false,
115115
CanonicalVarKind::Region(_) => true,
116116
CanonicalVarKind::PlaceholderRegion(..) => false,
@@ -122,7 +122,7 @@ impl<I: Interner> CanonicalVarInfo<I> {
122122
pub fn is_region(&self) -> bool {
123123
match self.kind {
124124
CanonicalVarKind::Region(_) | CanonicalVarKind::PlaceholderRegion(_) => true,
125-
CanonicalVarKind::Ty(_)
125+
CanonicalVarKind::Ty { .. }
126126
| CanonicalVarKind::Int
127127
| CanonicalVarKind::Float
128128
| CanonicalVarKind::PlaceholderTy(_)
@@ -133,7 +133,7 @@ impl<I: Interner> CanonicalVarInfo<I> {
133133

134134
pub fn expect_placeholder_index(self) -> usize {
135135
match self.kind {
136-
CanonicalVarKind::Ty(_)
136+
CanonicalVarKind::Ty { .. }
137137
| CanonicalVarKind::Int
138138
| CanonicalVarKind::Float
139139
| CanonicalVarKind::Region(_)
@@ -158,7 +158,11 @@ impl<I: Interner> CanonicalVarInfo<I> {
158158
)]
159159
pub enum CanonicalVarKind<I: Interner> {
160160
/// A general type variable `?T` that can be unified with arbitrary types.
161-
Ty(UniverseIndex),
161+
///
162+
/// We also store the index of the first type variable which is sub-unified
163+
/// with this one. If there is no inference variable related to this one,
164+
/// its `sub_root` just points to itself.
165+
Ty { universe: UniverseIndex, sub_root: ty::BoundVar },
162166

163167
/// Integral type variable `?I` (that can only be unified with integral types).
164168
Int,
@@ -187,7 +191,7 @@ pub enum CanonicalVarKind<I: Interner> {
187191
impl<I: Interner> CanonicalVarKind<I> {
188192
pub fn universe(self) -> UniverseIndex {
189193
match self {
190-
CanonicalVarKind::Ty(ui) => ui,
194+
CanonicalVarKind::Ty { universe, sub_root: _ } => universe,
191195
CanonicalVarKind::Region(ui) => ui,
192196
CanonicalVarKind::Const(ui) => ui,
193197
CanonicalVarKind::PlaceholderTy(placeholder) => placeholder.universe(),
@@ -203,7 +207,9 @@ impl<I: Interner> CanonicalVarKind<I> {
203207
/// the updated universe is not the root.
204208
pub fn with_updated_universe(self, ui: UniverseIndex) -> CanonicalVarKind<I> {
205209
match self {
206-
CanonicalVarKind::Ty(_) => CanonicalVarKind::Ty(ui),
210+
CanonicalVarKind::Ty { universe: _, sub_root } => {
211+
CanonicalVarKind::Ty { universe: ui, sub_root }
212+
}
207213
CanonicalVarKind::Region(_) => CanonicalVarKind::Region(ui),
208214
CanonicalVarKind::Const(_) => CanonicalVarKind::Const(ui),
209215

@@ -297,7 +303,7 @@ impl<I: Interner> CanonicalVarValues<I> {
297303
var_values: cx.mk_args_from_iter(infos.iter().enumerate().map(
298304
|(i, info)| -> I::GenericArg {
299305
match info.kind {
300-
CanonicalVarKind::Ty(_)
306+
CanonicalVarKind::Ty { .. }
301307
| CanonicalVarKind::Int
302308
| CanonicalVarKind::Float
303309
| CanonicalVarKind::PlaceholderTy(_) => {

‎compiler/rustc_type_ir/src/infer_ctxt.rs

+1
Original file line numberDiff line numberDiff line change
@@ -149,6 +149,7 @@ pub trait InferCtxtLike: Sized {
149149
fn universe_of_ct(&self, ct: ty::ConstVid) -> Option<ty::UniverseIndex>;
150150

151151
fn root_ty_var(&self, var: ty::TyVid) -> ty::TyVid;
152+
fn sub_root_ty_var(&self, var: ty::TyVid) -> ty::TyVid;
152153
fn root_const_var(&self, var: ty::ConstVid) -> ty::ConstVid;
153154

154155
fn opportunistic_resolve_ty_var(&self, vid: ty::TyVid) -> <Self::Interner as Interner>::Ty;

0 commit comments

Comments
 (0)