add universe assert in instantiation

Author: adwinwhite

compiler/rustc_infer/src/infer/context.rs

@@ -255,21 +255,13 @@ impl<'tcx> rustc_type_ir::InferCtxtLike for InferCtxt<'tcx> {
     }
 
     fn instantiate_ty_var_raw(&self, vid: ty::TyVid, ty: Ty<'tcx>) {
-        let ty = ty.fold_with(&mut LowerUniverseFolder {
-            infcx: self,
-            for_universe: self.try_resolve_ty_var(vid).unwrap_err(),
-            cache: Default::default(),
-        });
+        let ty = lower_universe(self, self.try_resolve_ty_var(vid).unwrap_err(), ty);
 
         self.inner.borrow_mut().type_variables().instantiate(vid, ty);
     }
 
     fn instantiate_const_var_raw(&self, vid: ty::ConstVid, ct: ty::Const<'tcx>) {
-        let ct = ct.fold_with(&mut LowerUniverseFolder {
-            infcx: self,
-            for_universe: self.try_resolve_const_var(vid).unwrap_err(),
-            cache: Default::default(),
-        });
+        let ct = lower_universe(self, self.try_resolve_const_var(vid).unwrap_err(), ct);
 
         self.inner
             .borrow_mut()
@@ -438,6 +430,33 @@ impl<'tcx> rustc_type_ir::InferCtxtLike for InferCtxt<'tcx> {
     }
 }
 
+fn lower_universe<'tcx, T: TypeFoldable<TyCtxt<'tcx>> + Copy>(
+    infcx: &InferCtxt<'tcx>,
+    for_universe: ty::UniverseIndex,
+    value: T,
+) -> T {
+    let value = value.fold_with(&mut LowerUniverseFolder {
+        infcx,
+        for_universe,
+        cache: Default::default(),
+    });
+
+    // This assertion is needed because we don't lower the universes of placeholders
+    // in the folder.
+    #[cfg(debug_assertions)]
+    {
+        let value_universe = ty::max_universe(infcx, value);
+        assert!(
+            for_universe.can_name(value_universe),
+            "variable in universe {:?} can't name value in universe {:?}",
+            for_universe,
+            value_universe,
+        );
+    }
+
+    value
+}
+
 /// Canonicalizing inputs puts all inference variables and placeholders
 /// into the root universe.
 ///
@@ -487,15 +506,6 @@ impl<'a, 'tcx> ty::TypeFolder<TyCtxt<'tcx>> for LowerUniverseFolder<'a, 'tcx> {
                     }
                 }
             }
-            ty::Placeholder(p) => {
-                debug_assert!(
-                    self.for_universe.can_name(p.universe),
-                    "variable in universe {:?} can't name type in universe {:?}",
-                    self.for_universe,
-                    p.universe
-                );
-                t
-            }
             _ => t.super_fold_with(self),
         };
 
@@ -532,15 +542,6 @@ impl<'a, 'tcx> ty::TypeFolder<TyCtxt<'tcx>> for LowerUniverseFolder<'a, 'tcx> {
                     ty::Const::new_var(self.cx(), new_var_id)
                 }
             }
-            ty::ConstKind::Placeholder(p) => {
-                debug_assert!(
-                    self.for_universe.can_name(p.universe),
-                    "variable in universe {:?} can't name const in universe {:?}",
-                    self.for_universe,
-                    p.universe
-                );
-                c
-            }
             _ => c.super_fold_with(self),
         }
     }

compiler/rustc_infer/src/infer/relate/generalize.rs

@@ -297,13 +297,24 @@ impl<'tcx> InferCtxt<'tcx> {
                 if let Some(r) = r.ty_vid() {
                     self.inner.borrow_mut().type_variables().equate(l, r)
                 } else {
+                    // Ideally, we put this assert into `type_variables().instantiate()`.
+                    // But we can't pass the infcx into it as the infcx is already
+                    // mutably borrowed.
+                    debug_assert!(
+                        self.try_resolve_ty_var(l).unwrap_err().can_name(ty::max_universe(self, r))
+                    );
                     self.inner.borrow_mut().type_variables().instantiate(l, r)
                 }
             }
             (TermVid::Const(l), ty::TermKind::Const(r)) => {
                 if let Some(r) = r.ct_vid() {
                     self.inner.borrow_mut().const_unification_table().union(l, r)
                 } else {
+                    debug_assert!(
+                        self.try_resolve_const_var(l)
+                            .unwrap_err()
+                            .can_name(ty::max_universe(self, r))
+                    );
                     self.inner
                         .borrow_mut()
                         .const_unification_table()

compiler/rustc_type_ir/src/universe.rs

@@ -1,5 +1,6 @@
 use tracing::{debug, instrument};
 
+use crate::data_structures::HashSet;
 use crate::inherent::*;
 use crate::visit::TypeVisitableExt;
 use crate::{
@@ -49,37 +50,48 @@ fn max_universe_inner<
     infcx: &Infcx,
     t: T,
 ) -> UniverseIndex {
-    if !MaxUniverse::<Infcx, VISIT_PLACEHOLDER, VISIT_INFER>::needs_visit(&t) {
+    if !MaxUniverse::<Infcx, I, VISIT_PLACEHOLDER, VISIT_INFER>::needs_visit(&t) {
         return UniverseIndex::ROOT;
     }
 
-    let mut visitor = MaxUniverse::<_, VISIT_PLACEHOLDER, VISIT_INFER>::new(infcx);
-    // FIXME: make this a debug_assert and let callers resolve vars if there's
-    // perf win here.
+    let mut visitor = MaxUniverse::<_, _, VISIT_PLACEHOLDER, VISIT_INFER>::new(infcx);
+    // FIXME: make this a debug_assert and let callers resolve vars. Then the input only needs to
+    // be `TypeVisitable`.
     let t = infcx.resolve_vars_if_possible(t);
     t.visit_with(&mut visitor);
     visitor.max_universe()
 }
 
-struct MaxUniverse<'a, Infcx: InferCtxtLike, const VISIT_PLACEHOLDER: bool, const VISIT_INFER: bool>
-{
+struct MaxUniverse<
+    'a,
+    Infcx: InferCtxtLike<Interner = I>,
+    I: Interner,
+    const VISIT_PLACEHOLDER: bool,
+    const VISIT_INFER: bool,
+> {
     max_universe: UniverseIndex,
     infcx: &'a Infcx,
+    cache: HashSet<I::Ty>,
 }
 
-impl<'a, Infcx: InferCtxtLike, const VISIT_PLACEHOLDER: bool, const VISIT_INFER: bool>
-    MaxUniverse<'a, Infcx, VISIT_PLACEHOLDER, VISIT_INFER>
+impl<
+    'a,
+    Infcx: InferCtxtLike<Interner = I>,
+    I: Interner,
+    const VISIT_PLACEHOLDER: bool,
+    const VISIT_INFER: bool,
+> MaxUniverse<'a, Infcx, I, VISIT_PLACEHOLDER, VISIT_INFER>
 {
     fn new(infcx: &'a Infcx) -> Self {
-        MaxUniverse { infcx, max_universe: UniverseIndex::ROOT }
+        MaxUniverse { infcx, max_universe: UniverseIndex::ROOT, cache: Default::default() }
     }
 
     fn max_universe(self) -> UniverseIndex {
         self.max_universe
     }
 
     #[instrument(ret, level = "debug")]
-    fn needs_visit<T: TypeVisitable<I>, I: Interner>(t: &T) -> bool {
+    fn needs_visit<T: TypeVisitable<I>>(t: &T) -> bool {
         (VISIT_PLACEHOLDER && t.has_placeholders()) || (VISIT_INFER && t.has_infer())
     }
 }
@@ -90,7 +102,7 @@ impl<
     I: Interner,
     const VISIT_PLACEHOLDER: bool,
     const VISIT_INFER: bool,
-> TypeVisitor<I> for MaxUniverse<'a, Infcx, VISIT_PLACEHOLDER, VISIT_INFER>
+> TypeVisitor<I> for MaxUniverse<'a, Infcx, I, VISIT_PLACEHOLDER, VISIT_INFER>
 {
     type Result = ();
 
@@ -99,6 +111,10 @@ impl<
             return;
         }
 
+        if self.cache.contains(&t) {
+            return;
+        }
+
         match t.kind() {
             TyKind::Placeholder(p) if VISIT_PLACEHOLDER => {
                 self.max_universe = self.max_universe.max(p.universe)
@@ -110,6 +126,8 @@ impl<
             }
             _ => t.super_visit_with(self),
         }
+
+        assert!(self.cache.insert(t), "we shouldn't visit {t:?} twice");
     }
 
     fn visit_const(&mut self, c: I::Const) {