anonymize all bound vars, not just regions

Author: lcnr

compiler/rustc_infer/src/infer/canonical/substitute.rs

@@ -71,7 +71,7 @@ where
     if var_values.var_values.is_empty() {
         value
     } else {
-        let delegate = FnMutDelegate {
+        let mut delegate = FnMutDelegate {
             regions: |br: ty::BoundRegion| match var_values.var_values[br.var].unpack() {
                 GenericArgKind::Lifetime(l) => l,
                 r => bug!("{:?} is a region but value is {:?}", br, r),
@@ -86,6 +86,6 @@ where
             },
         };
 
-        tcx.replace_escaping_bound_vars_uncached(value, delegate)
+        tcx.replace_escaping_bound_vars_uncached(value, &mut delegate)
     }
 }

compiler/rustc_infer/src/traits/util.rs

@@ -11,7 +11,7 @@ pub fn anonymize_predicate<'tcx>(
     tcx: TyCtxt<'tcx>,
     pred: ty::Predicate<'tcx>,
 ) -> ty::Predicate<'tcx> {
-    let new = tcx.anonymize_late_bound_regions(pred.kind());
+    let new = tcx.anonymize_bound_vars(pred.kind());
     tcx.reuse_or_mk_predicate(pred, new)
 }
 
@@ -334,7 +334,7 @@ pub fn transitive_bounds_that_define_assoc_type<'tcx>(
 
     std::iter::from_fn(move || {
         while let Some(trait_ref) = stack.pop() {
-            let anon_trait_ref = tcx.anonymize_late_bound_regions(trait_ref);
+            let anon_trait_ref = tcx.anonymize_bound_vars(trait_ref);
             if visited.insert(anon_trait_ref) {
                 let super_predicates = tcx.super_predicates_that_define_assoc_type((
                     trait_ref.def_id(),

compiler/rustc_middle/src/ty/erase_regions.rs

@@ -49,7 +49,7 @@ impl<'tcx> TypeFolder<'tcx> for RegionEraserVisitor<'tcx> {
     where
         T: TypeFoldable<'tcx>,
     {
-        let u = self.tcx.anonymize_late_bound_regions(t);
+        let u = self.tcx.anonymize_bound_vars(t);
         u.super_fold_with(self)
     }
 

compiler/rustc_middle/src/ty/fold.rs

@@ -44,7 +44,8 @@
 //! - u.fold_with(folder)
 //! ```
 use crate::mir;
-use crate::ty::{self, Binder, Ty, TyCtxt, TypeVisitable};
+use crate::ty::{self, Binder, BoundTy, Ty, TyCtxt, TypeVisitable};
+use rustc_data_structures::fx::FxIndexMap;
 use rustc_hir::def_id::DefId;
 
 use std::collections::BTreeMap;
@@ -533,12 +534,12 @@ impl<'tcx> TyCtxt<'tcx> {
     pub fn replace_escaping_bound_vars_uncached<T: TypeFoldable<'tcx>>(
         self,
         value: T,
-        mut delegate: impl BoundVarReplacerDelegate<'tcx>,
+        delegate: &mut impl BoundVarReplacerDelegate<'tcx>,
     ) -> T {
         if !value.has_escaping_bound_vars() {
             value
         } else {
-            let mut replacer = BoundVarReplacer::new(self, &mut delegate);
+            let mut replacer = BoundVarReplacer::new(self, delegate);
             value.fold_with(&mut replacer)
         }
     }
@@ -549,9 +550,9 @@ impl<'tcx> TyCtxt<'tcx> {
     pub fn replace_bound_vars_uncached<T: TypeFoldable<'tcx>>(
         self,
         value: Binder<'tcx, T>,
-        delegate: impl BoundVarReplacerDelegate<'tcx>,
+        mut delegate: impl BoundVarReplacerDelegate<'tcx>,
     ) -> T {
-        self.replace_escaping_bound_vars_uncached(value.skip_binder(), delegate)
+        self.replace_escaping_bound_vars_uncached(value.skip_binder(), &mut delegate)
     }
 
     /// Replaces any late-bound regions bound in `value` with
@@ -579,7 +580,7 @@ impl<'tcx> TyCtxt<'tcx> {
         let shift_bv = |bv: ty::BoundVar| ty::BoundVar::from_usize(bv.as_usize() + bound_vars);
         self.replace_escaping_bound_vars_uncached(
             value,
-            FnMutDelegate {
+            &mut FnMutDelegate {
                 regions: |r: ty::BoundRegion| {
                     self.mk_region(ty::ReLateBound(
                         ty::INNERMOST,
@@ -640,6 +641,50 @@ impl<'tcx> TyCtxt<'tcx> {
         );
         Binder::bind_with_vars(inner, bound_vars)
     }
+
+    /// Anonymize all bound variables in `value`, this is mostly used to improve caching.
+    pub fn anonymize_bound_vars<T>(self, value: Binder<'tcx, T>) -> Binder<'tcx, T>
+    where
+        T: TypeFoldable<'tcx>,
+    {
+        struct Anonymize<'tcx> {
+            tcx: TyCtxt<'tcx>,
+            map: FxIndexMap<ty::BoundVar, ty::BoundVariableKind>,
+        }
+        impl<'tcx> BoundVarReplacerDelegate<'tcx> for Anonymize<'tcx> {
+            fn replace_region(&mut self, br: ty::BoundRegion) -> ty::Region<'tcx> {
+                let entry = self.map.entry(br.var);
+                let index = entry.index();
+                let var = ty::BoundVar::from_usize(index);
+                let kind = entry
+                    .or_insert_with(|| ty::BoundVariableKind::Region(ty::BrAnon(index as u32)))
+                    .expect_region();
+                let br = ty::BoundRegion { var, kind };
+                self.tcx.mk_region(ty::ReLateBound(ty::INNERMOST, br))
+            }
+            fn replace_ty(&mut self, bt: ty::BoundTy) -> Ty<'tcx> {
+                let entry = self.map.entry(bt.var);
+                let index = entry.index();
+                let var = ty::BoundVar::from_usize(index);
+                let kind = entry
+                    .or_insert_with(|| ty::BoundVariableKind::Ty(ty::BoundTyKind::Anon))
+                    .expect_ty();
+                self.tcx.mk_ty(ty::Bound(ty::INNERMOST, BoundTy { var, kind }))
+            }
+            fn replace_const(&mut self, bv: ty::BoundVar, ty: Ty<'tcx>) -> ty::Const<'tcx> {
+                let entry = self.map.entry(bv);
+                let index = entry.index();
+                let var = ty::BoundVar::from_usize(index);
+                let () = entry.or_insert_with(|| ty::BoundVariableKind::Const).expect_const();
+                self.tcx.mk_const(ty::ConstS { ty, kind: ty::ConstKind::Bound(ty::INNERMOST, var) })
+            }
+        }
+
+        let mut delegate = Anonymize { tcx: self, map: Default::default() };
+        let inner = self.replace_escaping_bound_vars_uncached(value.skip_binder(), &mut delegate);
+        let bound_vars = self.mk_bound_variable_kinds(delegate.map.into_values());
+        Binder::bind_with_vars(inner, bound_vars)
+    }
 }
 
 ///////////////////////////////////////////////////////////////////////////

compiler/rustc_middle/src/ty/sty.rs

@@ -976,6 +976,29 @@ pub enum BoundVariableKind {
     Const,
 }
 
+impl BoundVariableKind {
+    pub fn expect_region(self) -> BoundRegionKind {
+        match self {
+            BoundVariableKind::Region(lt) => lt,
+            _ => bug!("expected a region, but found another kind"),
+        }
+    }
+
+    pub fn expect_ty(self) -> BoundTyKind {
+        match self {
+            BoundVariableKind::Ty(ty) => ty,
+            _ => bug!("expected a type, but found another kind"),
+        }
+    }
+
+    pub fn expect_const(self) {
+        match self {
+            BoundVariableKind::Const => (),
+            _ => bug!("expected a const, but found another kind"),
+        }
+    }
+}
+
 /// Binder is a binder for higher-ranked lifetimes or types. It is part of the
 /// compiler's representation for things like `for<'a> Fn(&'a isize)`
 /// (which would be represented by the type `PolyTraitRef ==

compiler/rustc_typeck/src/check/dropck.rs

@@ -318,8 +318,8 @@ impl<'tcx> TypeRelation<'tcx> for SimpleEqRelation<'tcx> {
 
         // Anonymizing the LBRs is necessary to solve (Issue #59497).
         // After we do so, it should be totally fine to skip the binders.
-        let anon_a = self.tcx.anonymize_late_bound_regions(a);
-        let anon_b = self.tcx.anonymize_late_bound_regions(b);
+        let anon_a = self.tcx.anonymize_bound_vars(a);
+        let anon_b = self.tcx.anonymize_bound_vars(b);
         self.relate(anon_a.skip_binder(), anon_b.skip_binder())?;
 
         Ok(a)

src/test/ui/generic-associated-types/anonymize-bound-vars.rs

@@ -0,0 +1,14 @@
+// check-pass
+//
+// regression test for #98702
+#![feature(generic_associated_types)]
+
+trait Foo {
+    type Assoc<T>;
+}
+
+impl Foo for () {
+    type Assoc<T> = [T; 2*2];
+}
+
+fn main() {}