Skip to content

Commit 667d26c

Browse files
lcnrlcnr
committed
track relevant sub_relations in canonical queries
1 parent 945aef7 commit 667d26c

File tree

10 files changed

+156
-76
lines changed

10 files changed

+156
-76
lines changed

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

Lines changed: 16 additions & 2 deletions
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

Lines changed: 16 additions & 7 deletions
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

Lines changed: 43 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,9 @@ use std::iter;
1313
use rustc_index::{Idx, IndexVec};
1414
use rustc_middle::arena::ArenaAllocatable;
1515
use rustc_middle::mir::ConstraintCategory;
16-
use rustc_middle::ty::{self, BoundVar, GenericArg, GenericArgKind, Ty, TyCtxt, TypeFoldable};
16+
use rustc_middle::ty::{
17+
self, BoundVar, CanonicalVarKind, GenericArg, GenericArgKind, Ty, TyCtxt, TypeFoldable,
18+
};
1719
use rustc_middle::{bug, span_bug};
1820
use tracing::{debug, instrument};
1921

@@ -443,32 +445,48 @@ impl<'tcx> InferCtxt<'tcx> {
443445
// Create result arguments: if we found a value for a
444446
// given variable in the loop above, use that. Otherwise, use
445447
// a fresh inference variable.
446-
let result_args = CanonicalVarValues {
447-
var_values: self.tcx.mk_args_from_iter(
448-
query_response.variables.iter().enumerate().map(|(index, info)| {
449-
if info.universe() != ty::UniverseIndex::ROOT {
450-
// A variable from inside a binder of the query. While ideally these shouldn't
451-
// exist at all, we have to deal with them for now.
452-
self.instantiate_canonical_var(cause.span, info, |u| {
453-
universe_map[u.as_usize()]
454-
})
455-
} else if info.is_existential() {
456-
match opt_values[BoundVar::new(index)] {
457-
Some(k) => k,
458-
None => self.instantiate_canonical_var(cause.span, info, |u| {
459-
universe_map[u.as_usize()]
460-
}),
448+
let mut var_values = Vec::new();
449+
for (index, info) in query_response.variables.iter().enumerate() {
450+
let value = if info.universe() != ty::UniverseIndex::ROOT {
451+
// A variable from inside a binder of the query. While ideally these shouldn't
452+
// exist at all, we have to deal with them for now.
453+
self.instantiate_canonical_var(cause.span, info, &var_values, |u| {
454+
universe_map[u.as_usize()]
455+
})
456+
} else if info.is_existential() {
457+
// As an optimization we sometimes avoid creating a new inference variable here.
458+
// We need to still make sure to register any subtype relations returned by the
459+
// query.
460+
match opt_values[BoundVar::new(index)] {
461+
Some(v) => {
462+
if let CanonicalVarKind::Ty { universe: _, sub_root } = info.kind {
463+
if let Some(prev) = var_values.get(sub_root.as_usize()) {
464+
let &ty::Infer(ty::TyVar(vid)) = v.expect_ty().kind() else {
465+
unreachable!("expected `sub_root` to be an inference variable");
466+
};
467+
let &ty::Infer(ty::TyVar(sub_root)) = prev.expect_ty().kind()
468+
else {
469+
unreachable!("expected `sub_root` to be an inference variable");
470+
};
471+
self.inner.borrow_mut().type_variables().sub(vid, sub_root);
472+
}
461473
}
462-
} else {
463-
// For placeholders which were already part of the input, we simply map this
464-
// universal bound variable back the placeholder of the input.
465-
opt_values[BoundVar::new(index)].expect(
466-
"expected placeholder to be unified with itself during response",
467-
)
474+
v
468475
}
469-
}),
470-
),
471-
};
476+
None => self.instantiate_canonical_var(cause.span, info, &var_values, |u| {
477+
universe_map[u.as_usize()]
478+
}),
479+
}
480+
} else {
481+
// For placeholders which were already part of the input, we simply map this
482+
// universal bound variable back the placeholder of the input.
483+
opt_values[BoundVar::new(index)]
484+
.expect("expected placeholder to be unified with itself during response")
485+
};
486+
var_values.push(value)
487+
}
488+
489+
let result_args = CanonicalVarValues { var_values: self.tcx.mk_args(&var_values) };
472490

473491
let mut obligations = PredicateObligations::new();
474492

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,9 @@ impl<'tcx> rustc_type_ir::InferCtxtLike for InferCtxt<'tcx> {
5858
fn root_ty_var(&self, var: ty::TyVid) -> ty::TyVid {
5959
self.root_var(var)
6060
}
61+
fn sub_root_ty_var(&self, var: ty::TyVid) -> ty::TyVid {
62+
self.sub_root_var(var)
63+
}
6164

6265
fn root_const_var(&self, var: ty::ConstVid) -> ty::ConstVid {
6366
self.root_const_var(var)

‎compiler/rustc_next_trait_solver/src/canonicalizer.rs

Lines changed: 17 additions & 5 deletions
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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ pub trait SolverDelegate: Deref<Target = Self::Infcx> + Sized {
5555
&self,
5656
cv_info: ty::CanonicalVarInfo<Self::Interner>,
5757
span: <Self::Interner as Interner>::Span,
58+
var_values: &[<Self::Interner as Interner>::GenericArg],
5859
universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
5960
) -> <Self::Interner as Interner>::GenericArg;
6061
fn add_item_bounds_for_hidden_type(

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

Lines changed: 44 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@ use rustc_index::IndexVec;
1515
use rustc_type_ir::inherent::*;
1616
use rustc_type_ir::relate::solver_relating::RelateExt;
1717
use rustc_type_ir::{
18-
self as ty, Canonical, CanonicalVarValues, InferCtxtLike, Interner, TypeFoldable,
18+
self as ty, Canonical, CanonicalVarKind, CanonicalVarValues, InferCtxtLike, Interner,
19+
TypeFoldable,
1920
};
2021
use tracing::{debug, instrument, trace};
2122

@@ -356,37 +357,51 @@ where
356357
}
357358
}
358359

359-
let var_values = delegate.cx().mk_args_from_iter(
360-
response.variables.iter().enumerate().map(|(index, info)| {
361-
if info.universe() != ty::UniverseIndex::ROOT {
362-
// A variable from inside a binder of the query. While ideally these shouldn't
363-
// exist at all (see the FIXME at the start of this method), we have to deal with
364-
// them for now.
365-
delegate.instantiate_canonical_var_with_infer(info, span, |idx| {
366-
prev_universe + idx.index()
367-
})
368-
} else if info.is_existential() {
369-
// As an optimization we sometimes avoid creating a new inference variable here.
370-
//
371-
// All new inference variables we create start out in the current universe of the caller.
372-
// This is conceptually wrong as these inference variables would be able to name
373-
// more placeholders then they should be able to. However the inference variables have
374-
// to "come from somewhere", so by equating them with the original values of the caller
375-
// later on, we pull them down into their correct universe again.
376-
if let Some(v) = opt_values[ty::BoundVar::from_usize(index)] {
377-
v
378-
} else {
379-
delegate.instantiate_canonical_var_with_infer(info, span, |_| prev_universe)
360+
let mut var_values = Vec::new();
361+
for (index, info) in response.variables.iter().enumerate() {
362+
let value = if info.universe() != ty::UniverseIndex::ROOT {
363+
// A variable from inside a binder of the query. While ideally these shouldn't
364+
// exist at all (see the FIXME at the start of this method), we have to deal with
365+
// them for now.
366+
delegate.instantiate_canonical_var_with_infer(info, span, &var_values, |idx| {
367+
prev_universe + idx.index()
368+
})
369+
} else if info.is_existential() {
370+
// As an optimization we sometimes avoid creating a new inference variable here.
371+
// We need to still make sure to register any subtype relations returned by the
372+
// query.
373+
if let Some(v) = opt_values[ty::BoundVar::from_usize(index)] {
374+
if let CanonicalVarKind::Ty { universe: _, sub_root } = info.kind {
375+
if let Some(prev) = var_values.get(sub_root.as_usize()) {
376+
let ty::Infer(ty::TyVar(vid)) = v.expect_ty().kind() else {
377+
unreachable!("expected `sub_root` to be an inference variable");
378+
};
379+
let ty::Infer(ty::TyVar(sub_root)) = prev.expect_ty().kind() else {
380+
unreachable!("expected `sub_root` to be an inference variable");
381+
};
382+
delegate.sub_ty_vids_raw(vid, sub_root);
383+
}
380384
}
385+
v
381386
} else {
382-
// For placeholders which were already part of the input, we simply map this
383-
// universal bound variable back the placeholder of the input.
384-
original_values[info.expect_placeholder_index()]
387+
// All new inference variables we create start out in the current universe
388+
// of the caller. This is conceptually wrong as these inference variables
389+
// would be able to name more placeholders then they should be able to.
390+
// However the inference variables have to "come from somewhere", so by
391+
// equating them with the original values of the caller later on, we pull
392+
// them down into their correct universe again.
393+
delegate.instantiate_canonical_var_with_infer(info, span, &var_values, |_| {
394+
prev_universe
395+
})
385396
}
386-
}),
387-
);
388-
389-
CanonicalVarValues { var_values }
397+
} else {
398+
// For placeholders which were already part of the input, we simply map this
399+
// universal bound variable back the placeholder of the input.
400+
original_values[info.expect_placeholder_index()]
401+
};
402+
var_values.push(value)
403+
}
404+
CanonicalVarValues { var_values: delegate.cx().mk_args(&var_values) }
390405
}
391406

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

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -144,9 +144,10 @@ impl<'tcx> rustc_next_trait_solver::delegate::SolverDelegate for SolverDelegate<
144144
&self,
145145
cv_info: CanonicalVarInfo<'tcx>,
146146
span: Span,
147+
var_values: &[ty::GenericArg<'tcx>],
147148
universe_map: impl Fn(ty::UniverseIndex) -> ty::UniverseIndex,
148149
) -> ty::GenericArg<'tcx> {
149-
self.0.instantiate_canonical_var(span, cv_info, universe_map)
150+
self.0.instantiate_canonical_var(span, cv_info, var_values, universe_map)
150151
}
151152

152153
fn add_item_bounds_for_hidden_type(

0 commit comments

Comments
 (0)