1use rustc_hash::FxHashMap;
9use rustc_index::Idx;
10use rustc_type_ir::InferTy::{self, FloatVar, IntVar, TyVar};
11use rustc_type_ir::inherent::{Const as _, IntoKind as _, Region as _, SliceLike, Ty as _};
12use rustc_type_ir::{
13 BoundVar, BoundVarIndexKind, DebruijnIndex, Flags, InferConst, RegionKind, TyVid, TypeFlags,
14 TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitableExt, UniverseIndex,
15};
16use smallvec::SmallVec;
17use tracing::debug;
18
19use crate::next_solver::infer::InferCtxt;
20use crate::next_solver::{
21 Binder, Canonical, CanonicalVarKind, CanonicalVars, Const, ConstKind, DbInterner, GenericArg,
22 Placeholder, Region, Ty, TyKind,
23};
24
25#[derive(Clone, Debug)]
30pub struct OriginalQueryValues<'db> {
31 pub universe_map: SmallVec<[UniverseIndex; 4]>,
36
37 pub var_values: SmallVec<[GenericArg<'db>; 8]>,
40}
41
42impl<'db> Default for OriginalQueryValues<'db> {
43 fn default() -> Self {
44 let mut universe_map = SmallVec::default();
45 universe_map.push(UniverseIndex::ROOT);
46
47 Self { universe_map, var_values: SmallVec::default() }
48 }
49}
50
51impl<'db> InferCtxt<'db> {
52 pub fn canonicalize_query<V>(
68 &self,
69 value: V,
70 query_state: &mut OriginalQueryValues<'db>,
71 ) -> Canonical<'db, V>
72 where
73 V: TypeFoldable<DbInterner<'db>>,
74 {
75 Canonicalizer::canonicalize(
76 value,
77 self,
78 self.interner,
79 &CanonicalizeAllFreeRegions,
80 query_state,
81 )
82 }
83
84 pub fn canonicalize_response<V>(&self, value: V) -> Canonical<'db, V>
110 where
111 V: TypeFoldable<DbInterner<'db>>,
112 {
113 let mut query_state = OriginalQueryValues::default();
114 Canonicalizer::canonicalize(
115 value,
116 self,
117 self.interner,
118 &CanonicalizeQueryResponse,
119 &mut query_state,
120 )
121 }
122
123 pub fn canonicalize_user_type_annotation<V>(&self, value: V) -> Canonical<'db, V>
124 where
125 V: TypeFoldable<DbInterner<'db>>,
126 {
127 let mut query_state = OriginalQueryValues::default();
128 Canonicalizer::canonicalize(
129 value,
130 self,
131 self.interner,
132 &CanonicalizeUserTypeAnnotation,
133 &mut query_state,
134 )
135 }
136}
137
138trait CanonicalizeMode {
146 fn canonicalize_free_region<'db>(
147 &self,
148 canonicalizer: &mut Canonicalizer<'_, 'db>,
149 r: Region<'db>,
150 ) -> Region<'db>;
151
152 fn any(&self) -> bool;
153
154 fn preserve_universes(&self) -> bool;
156}
157
158struct CanonicalizeQueryResponse;
159
160impl CanonicalizeMode for CanonicalizeQueryResponse {
161 fn canonicalize_free_region<'db>(
162 &self,
163 canonicalizer: &mut Canonicalizer<'_, 'db>,
164 mut r: Region<'db>,
165 ) -> Region<'db> {
166 let infcx = canonicalizer.infcx;
167
168 if let RegionKind::ReVar(vid) = r.kind() {
169 r = infcx
170 .inner
171 .borrow_mut()
172 .unwrap_region_constraints()
173 .opportunistic_resolve_var(canonicalizer.tcx, vid);
174 debug!(
175 "canonical: region var found with vid {vid:?}, \
176 opportunistically resolved to {r:?}",
177 );
178 };
179
180 match r.kind() {
181 RegionKind::ReLateParam(_)
182 | RegionKind::ReErased
183 | RegionKind::ReStatic
184 | RegionKind::ReEarlyParam(..)
185 | RegionKind::ReError(..) => r,
186
187 RegionKind::RePlaceholder(placeholder) => canonicalizer
188 .canonical_var_for_region(CanonicalVarKind::PlaceholderRegion(placeholder), r),
189
190 RegionKind::ReVar(vid) => {
191 let universe = infcx
192 .inner
193 .borrow_mut()
194 .unwrap_region_constraints()
195 .probe_value(vid)
196 .unwrap_err();
197 canonicalizer.canonical_var_for_region(CanonicalVarKind::Region(universe), r)
198 }
199
200 _ => {
201 panic!("unexpected region in query response: `{r:?}`");
210 }
211 }
212 }
213
214 fn any(&self) -> bool {
215 false
216 }
217
218 fn preserve_universes(&self) -> bool {
219 true
220 }
221}
222
223struct CanonicalizeUserTypeAnnotation;
224
225impl CanonicalizeMode for CanonicalizeUserTypeAnnotation {
226 fn canonicalize_free_region<'db>(
227 &self,
228 canonicalizer: &mut Canonicalizer<'_, 'db>,
229 r: Region<'db>,
230 ) -> Region<'db> {
231 match r.kind() {
232 RegionKind::ReEarlyParam(_)
233 | RegionKind::ReLateParam(_)
234 | RegionKind::ReErased
235 | RegionKind::ReStatic
236 | RegionKind::ReError(_) => r,
237 RegionKind::ReVar(_) => canonicalizer.canonical_var_for_region_in_root_universe(r),
238 RegionKind::RePlaceholder(..) | RegionKind::ReBound(..) => {
239 panic!("unexpected region in query response: `{r:?}`")
241 }
242 }
243 }
244
245 fn any(&self) -> bool {
246 false
247 }
248
249 fn preserve_universes(&self) -> bool {
250 false
251 }
252}
253
254struct CanonicalizeAllFreeRegions;
255
256impl CanonicalizeMode for CanonicalizeAllFreeRegions {
257 fn canonicalize_free_region<'db>(
258 &self,
259 canonicalizer: &mut Canonicalizer<'_, 'db>,
260 r: Region<'db>,
261 ) -> Region<'db> {
262 canonicalizer.canonical_var_for_region_in_root_universe(r)
263 }
264
265 fn any(&self) -> bool {
266 true
267 }
268
269 fn preserve_universes(&self) -> bool {
270 false
271 }
272}
273
274struct Canonicalizer<'cx, 'db> {
275 infcx: &'cx InferCtxt<'db>,
277 tcx: DbInterner<'db>,
278 variables: SmallVec<[CanonicalVarKind<'db>; 8]>,
279 query_state: &'cx mut OriginalQueryValues<'db>,
280 indices: FxHashMap<GenericArg<'db>, BoundVar>,
283 sub_root_lookup_table: FxHashMap<TyVid, usize>,
290 canonicalize_mode: &'cx dyn CanonicalizeMode,
291 needs_canonical_flags: TypeFlags,
292
293 binder_index: DebruijnIndex,
294}
295
296impl<'cx, 'db> TypeFolder<DbInterner<'db>> for Canonicalizer<'cx, 'db> {
297 fn cx(&self) -> DbInterner<'db> {
298 self.tcx
299 }
300
301 fn fold_binder<T>(&mut self, t: Binder<'db, T>) -> Binder<'db, T>
302 where
303 T: TypeFoldable<DbInterner<'db>>,
304 {
305 self.binder_index.shift_in(1);
306 let t = t.super_fold_with(self);
307 self.binder_index.shift_out(1);
308 t
309 }
310
311 fn fold_region(&mut self, r: Region<'db>) -> Region<'db> {
312 match r.kind() {
313 RegionKind::ReBound(BoundVarIndexKind::Bound(..), ..) => r,
314 RegionKind::ReBound(BoundVarIndexKind::Canonical, ..) => {
315 panic!("canonicalized bound var found during canonicalization");
316 }
317
318 RegionKind::ReStatic
319 | RegionKind::ReEarlyParam(..)
320 | RegionKind::ReError(_)
321 | RegionKind::ReLateParam(_)
322 | RegionKind::RePlaceholder(..)
323 | RegionKind::ReVar(_)
324 | RegionKind::ReErased => self.canonicalize_mode.canonicalize_free_region(self, r),
325 }
326 }
327
328 fn fold_ty(&mut self, mut t: Ty<'db>) -> Ty<'db> {
329 match t.kind() {
330 TyKind::Infer(TyVar(mut vid)) => {
331 let root_vid = self.infcx.root_var(vid);
335 if root_vid != vid {
336 t = Ty::new_var(self.tcx, root_vid);
337 vid = root_vid;
338 }
339
340 debug!("canonical: type var found with vid {:?}", vid);
341 match self.infcx.probe_ty_var(vid) {
342 Ok(t) => {
344 debug!("(resolved to {:?})", t);
345 self.fold_ty(t)
346 }
347
348 Err(mut ui) => {
351 if !self.canonicalize_mode.preserve_universes() {
352 ui = UniverseIndex::ROOT;
354 }
355
356 let sub_root = self.get_or_insert_sub_root(vid);
357 self.canonicalize_ty_var(CanonicalVarKind::Ty { ui, sub_root }, t)
358 }
359 }
360 }
361
362 TyKind::Infer(IntVar(vid)) => {
363 let nt = self.infcx.opportunistic_resolve_int_var(vid);
364 if nt != t {
365 self.fold_ty(nt)
366 } else {
367 self.canonicalize_ty_var(CanonicalVarKind::Int, t)
368 }
369 }
370 TyKind::Infer(FloatVar(vid)) => {
371 let nt = self.infcx.opportunistic_resolve_float_var(vid);
372 if nt != t {
373 self.fold_ty(nt)
374 } else {
375 self.canonicalize_ty_var(CanonicalVarKind::Float, t)
376 }
377 }
378
379 TyKind::Infer(
380 InferTy::FreshTy(_) | InferTy::FreshIntTy(_) | InferTy::FreshFloatTy(_),
381 ) => {
382 panic!("encountered a fresh type during canonicalization")
383 }
384
385 TyKind::Placeholder(mut placeholder) => {
386 if !self.canonicalize_mode.preserve_universes() {
387 placeholder.universe = UniverseIndex::ROOT;
388 }
389 self.canonicalize_ty_var(CanonicalVarKind::PlaceholderTy(placeholder), t)
390 }
391
392 TyKind::Bound(BoundVarIndexKind::Bound(..), _) => t,
393 TyKind::Bound(BoundVarIndexKind::Canonical, ..) => {
394 panic!("canonicalized bound var found during canonicalization");
395 }
396
397 TyKind::Closure(..)
398 | TyKind::CoroutineClosure(..)
399 | TyKind::Coroutine(..)
400 | TyKind::CoroutineWitness(..)
401 | TyKind::Bool
402 | TyKind::Char
403 | TyKind::Int(..)
404 | TyKind::Uint(..)
405 | TyKind::Float(..)
406 | TyKind::Adt(..)
407 | TyKind::Str
408 | TyKind::Error(_)
409 | TyKind::Array(..)
410 | TyKind::Slice(..)
411 | TyKind::RawPtr(..)
412 | TyKind::Ref(..)
413 | TyKind::FnDef(..)
414 | TyKind::FnPtr(..)
415 | TyKind::Dynamic(..)
416 | TyKind::UnsafeBinder(_)
417 | TyKind::Never
418 | TyKind::Tuple(..)
419 | TyKind::Alias(..)
420 | TyKind::Foreign(..)
421 | TyKind::Pat(..)
422 | TyKind::Param(..) => {
423 if t.flags().intersects(self.needs_canonical_flags) {
424 t.super_fold_with(self)
425 } else {
426 t
427 }
428 }
429 }
430 }
431
432 fn fold_const(&mut self, mut ct: Const<'db>) -> Const<'db> {
433 match ct.kind() {
434 ConstKind::Infer(InferConst::Var(mut vid)) => {
435 let root_vid = self.infcx.root_const_var(vid);
439 if root_vid != vid {
440 ct = Const::new_var(self.tcx, root_vid);
441 vid = root_vid;
442 }
443
444 debug!("canonical: const var found with vid {:?}", vid);
445 match self.infcx.probe_const_var(vid) {
446 Ok(c) => {
447 debug!("(resolved to {:?})", c);
448 return self.fold_const(c);
449 }
450
451 Err(mut ui) => {
454 if !self.canonicalize_mode.preserve_universes() {
455 ui = UniverseIndex::ROOT;
457 }
458 return self.canonicalize_const_var(CanonicalVarKind::Const(ui), ct);
459 }
460 }
461 }
462 ConstKind::Infer(InferConst::Fresh(_)) => {
463 panic!("encountered a fresh const during canonicalization")
464 }
465 ConstKind::Bound(BoundVarIndexKind::Bound(..), _) => {
466 return ct;
467 }
468 ConstKind::Bound(BoundVarIndexKind::Canonical, ..) => {
469 panic!("canonicalized bound var found during canonicalization");
470 }
471 ConstKind::Placeholder(placeholder) => {
472 return self
473 .canonicalize_const_var(CanonicalVarKind::PlaceholderConst(placeholder), ct);
474 }
475 _ => {}
476 }
477
478 if ct.flags().intersects(self.needs_canonical_flags) {
479 ct.super_fold_with(self)
480 } else {
481 ct
482 }
483 }
484}
485
486impl<'cx, 'db> Canonicalizer<'cx, 'db> {
487 fn canonicalize<V>(
490 value: V,
491 infcx: &InferCtxt<'db>,
492 tcx: DbInterner<'db>,
493 canonicalize_region_mode: &dyn CanonicalizeMode,
494 query_state: &mut OriginalQueryValues<'db>,
495 ) -> Canonical<'db, V>
496 where
497 V: TypeFoldable<DbInterner<'db>>,
498 {
499 let base = Canonical {
500 max_universe: UniverseIndex::ROOT,
501 variables: CanonicalVars::new_from_iter(tcx, []),
502 value: (),
503 };
504 Canonicalizer::canonicalize_with_base(
505 base,
506 value,
507 infcx,
508 tcx,
509 canonicalize_region_mode,
510 query_state,
511 )
512 .unchecked_map(|((), val)| val)
513 }
514
515 fn canonicalize_with_base<U, V>(
516 base: Canonical<'db, U>,
517 value: V,
518 infcx: &InferCtxt<'db>,
519 tcx: DbInterner<'db>,
520 canonicalize_region_mode: &dyn CanonicalizeMode,
521 query_state: &mut OriginalQueryValues<'db>,
522 ) -> Canonical<'db, (U, V)>
523 where
524 V: TypeFoldable<DbInterner<'db>>,
525 {
526 let needs_canonical_flags = if canonicalize_region_mode.any() {
527 TypeFlags::HAS_INFER | TypeFlags::HAS_PLACEHOLDER | TypeFlags::HAS_FREE_REGIONS
528 } else {
529 TypeFlags::HAS_INFER | TypeFlags::HAS_PLACEHOLDER
530 };
531
532 if !value.has_type_flags(needs_canonical_flags) {
534 return base.unchecked_map(|b| (b, value));
535 }
536
537 let mut canonicalizer = Canonicalizer {
538 infcx,
539 tcx,
540 canonicalize_mode: canonicalize_region_mode,
541 needs_canonical_flags,
542 variables: SmallVec::from_slice(base.variables.as_slice()),
543 query_state,
544 indices: FxHashMap::default(),
545 sub_root_lookup_table: Default::default(),
546 binder_index: DebruijnIndex::ZERO,
547 };
548 if canonicalizer.query_state.var_values.spilled() {
549 canonicalizer.indices = canonicalizer
550 .query_state
551 .var_values
552 .iter()
553 .enumerate()
554 .map(|(i, &kind)| (kind, BoundVar::from(i)))
555 .collect();
556 }
557 let out_value = value.fold_with(&mut canonicalizer);
558
559 debug_assert!(!out_value.has_infer() && !out_value.has_placeholders());
563
564 let canonical_variables =
565 CanonicalVars::new_from_iter(tcx, canonicalizer.universe_canonicalized_variables());
566
567 let max_universe = canonical_variables
568 .iter()
569 .map(|cvar| cvar.universe())
570 .max()
571 .unwrap_or(UniverseIndex::ROOT);
572
573 Canonical { max_universe, variables: canonical_variables, value: (base.value, out_value) }
574 }
575
576 fn canonical_var(&mut self, info: CanonicalVarKind<'db>, kind: GenericArg<'db>) -> BoundVar {
581 let Canonicalizer { variables, query_state, indices, .. } = self;
582
583 let var_values = &mut query_state.var_values;
584
585 let universe = info.universe();
586 if universe != UniverseIndex::ROOT {
587 assert!(self.canonicalize_mode.preserve_universes());
588
589 match query_state.universe_map.binary_search(&universe) {
593 Err(idx) => query_state.universe_map.insert(idx, universe),
594 Ok(_) => {}
595 }
596 }
597
598 if !var_values.spilled() {
604 if let Some(idx) = var_values.iter().position(|&k| k == kind) {
607 BoundVar::new(idx)
609 } else {
610 variables.push(info);
613 var_values.push(kind);
614 assert_eq!(variables.len(), var_values.len());
615
616 if var_values.spilled() {
619 assert!(indices.is_empty());
620 *indices = var_values
621 .iter()
622 .enumerate()
623 .map(|(i, &kind)| (kind, BoundVar::new(i)))
624 .collect();
625 }
626 BoundVar::new(var_values.len() - 1)
628 }
629 } else {
630 *indices.entry(kind).or_insert_with(|| {
632 variables.push(info);
633 var_values.push(kind);
634 assert_eq!(variables.len(), var_values.len());
635 BoundVar::new(variables.len() - 1)
636 })
637 }
638 }
639
640 fn get_or_insert_sub_root(&mut self, vid: TyVid) -> BoundVar {
641 let root_vid = self.infcx.sub_unification_table_root_var(vid);
642 let idx =
643 *self.sub_root_lookup_table.entry(root_vid).or_insert_with(|| self.variables.len());
644 BoundVar::from(idx)
645 }
646
647 fn universe_canonicalized_variables(self) -> SmallVec<[CanonicalVarKind<'db>; 8]> {
651 if self.query_state.universe_map.len() == 1 {
652 return self.variables;
653 }
654
655 let reverse_universe_map: FxHashMap<UniverseIndex, UniverseIndex> = self
656 .query_state
657 .universe_map
658 .iter()
659 .enumerate()
660 .map(|(idx, universe)| (*universe, UniverseIndex::from_usize(idx)))
661 .collect();
662
663 self.variables
664 .iter()
665 .map(|v| match *v {
666 CanonicalVarKind::Int | CanonicalVarKind::Float => *v,
667 CanonicalVarKind::Ty { ui, sub_root } => {
668 CanonicalVarKind::Ty { ui: reverse_universe_map[&ui], sub_root }
669 }
670 CanonicalVarKind::Region(u) => CanonicalVarKind::Region(reverse_universe_map[&u]),
671 CanonicalVarKind::Const(u) => CanonicalVarKind::Const(reverse_universe_map[&u]),
672 CanonicalVarKind::PlaceholderTy(placeholder) => {
673 CanonicalVarKind::PlaceholderTy(Placeholder {
674 universe: reverse_universe_map[&placeholder.universe],
675 ..placeholder
676 })
677 }
678 CanonicalVarKind::PlaceholderRegion(placeholder) => {
679 CanonicalVarKind::PlaceholderRegion(Placeholder {
680 universe: reverse_universe_map[&placeholder.universe],
681 ..placeholder
682 })
683 }
684 CanonicalVarKind::PlaceholderConst(placeholder) => {
685 CanonicalVarKind::PlaceholderConst(Placeholder {
686 universe: reverse_universe_map[&placeholder.universe],
687 ..placeholder
688 })
689 }
690 })
691 .collect()
692 }
693
694 fn canonical_var_for_region_in_root_universe(&mut self, r: Region<'db>) -> Region<'db> {
708 self.canonical_var_for_region(CanonicalVarKind::Region(UniverseIndex::ROOT), r)
709 }
710
711 fn canonical_var_for_region(
714 &mut self,
715 info: CanonicalVarKind<'db>,
716 r: Region<'db>,
717 ) -> Region<'db> {
718 let var = self.canonical_var(info, r.into());
719 Region::new_canonical_bound(self.cx(), var)
720 }
721
722 fn canonicalize_ty_var(&mut self, info: CanonicalVarKind<'db>, ty_var: Ty<'db>) -> Ty<'db> {
727 debug_assert_eq!(ty_var, self.infcx.shallow_resolve(ty_var));
728 let var = self.canonical_var(info, ty_var.into());
729 Ty::new_canonical_bound(self.cx(), var)
730 }
731
732 fn canonicalize_const_var(
737 &mut self,
738 info: CanonicalVarKind<'db>,
739 const_var: Const<'db>,
740 ) -> Const<'db> {
741 debug_assert_eq!(const_var, self.infcx.shallow_resolve_const(const_var));
742 let var = self.canonical_var(info, const_var.into());
743 Const::new_canonical_bound(self.cx(), var)
744 }
745}