Normalize associated types with bound vars

2 files changed, 115 insertions, 70 deletions

diff --git a/compiler/rustc_trait_selection/src/traits/project.rs b/compiler/rustc_trait_selection/src/traits/project.rs
index 19e3c1afb9f..df1dca9831a 100644
--- a/compiler/rustc_trait_selection/src/traits/project.rs
+++ b/compiler/rustc_trait_selection/src/traits/project.rs
@@ -362,25 +362,25 @@ impl<'a, 'b, 'tcx> TypeFolder<'tcx> for AssocTypeNormalizer<'a, 'b, 'tcx> {
         if !needs_normalization(&ty, self.param_env.reveal()) {
             return ty;
         }
-        // We don't want to normalize associated types that occur inside of region
-        // binders, because they may contain bound regions, and we can't cope with that.
-        //
-        // Example:
-        //
-        //     for<'a> fn(<T as Foo<&'a>>::A)
-        //
-        // Instead of normalizing `<T as Foo<&'a>>::A` here, we'll
-        // normalize it when we instantiate those bound regions (which
-        // should occur eventually).
-
-        let ty = ty.super_fold_with(self);
+
+        // N.b. while we want to call `super_fold_with(self)` on `ty` before
+        // normalization, we wait until we know whether we need to normalize the
+        // current type. If we do, then we only fold the ty *after* replacing bound
+        // vars with placeholders. This means that nested types don't need to replace
+        // bound vars at the current binder level or above. A key assumption here is
+        // that folding the type can't introduce new bound vars.
+
         match *ty.kind() {
-            ty::Opaque(def_id, substs) if !substs.has_escaping_bound_vars() => {
+            ty::Opaque(def_id, substs) => {
                 // Only normalize `impl Trait` after type-checking, usually in codegen.
                 match self.param_env.reveal() {
-                    Reveal::UserFacing => ty,
+                    Reveal::UserFacing => ty.super_fold_with(self),
 
                     Reveal::All => {
+                        // N.b. there is an assumption here all this code can handle
+                        // escaping bound vars.
+
+                        let substs = substs.super_fold_with(self);
                         let recursion_limit = self.tcx().recursion_limit();
                         if !recursion_limit.value_within_limit(self.depth) {
                             let obligation = Obligation::with_depth(
@@ -403,18 +403,13 @@ impl<'a, 'b, 'tcx> TypeFolder<'tcx> for AssocTypeNormalizer<'a, 'b, 'tcx> {
             }
 
             ty::Projection(data) if !data.has_escaping_bound_vars() => {
-                // This is kind of hacky -- we need to be able to
-                // handle normalization within binders because
-                // otherwise we wind up a need to normalize when doing
-                // trait matching (since you can have a trait
-                // obligation like `for<'a> T::B: Fn(&'a i32)`), but
-                // we can't normalize with bound regions in scope. So
-                // far now we just ignore binders but only normalize
-                // if all bound regions are gone (and then we still
-                // have to renormalize whenever we instantiate a
-                // binder). It would be better to normalize in a
-                // binding-aware fashion.
+                // This branch is *mostly* just an optimization: when we don't
+                // have escaping bound vars, we don't need to replace them with
+                // placeholders (see branch below). *Also*, we know that we can
+                // register an obligation to *later* project, since we know
+                // there won't be bound vars there.
 
+                let data = data.super_fold_with(self);
                 let normalized_ty = normalize_projection_type(
                     self.selcx,
                     self.param_env,
@@ -433,22 +428,19 @@ impl<'a, 'b, 'tcx> TypeFolder<'tcx> for AssocTypeNormalizer<'a, 'b, 'tcx> {
                 normalized_ty
             }
 
-            ty::Projection(data) if !data.trait_ref(self.tcx()).has_escaping_bound_vars() => {
-                // Okay, so you thought the previous branch was hacky. Well, to
-                // extend upon this, when the *trait ref* doesn't have escaping
-                // bound vars, but the associated item *does* (can only occur
-                // with GATs), then we might still be able to project the type.
-                // For this, we temporarily replace the bound vars with
-                // placeholders. Note though, that in the case that we still
-                // can't project for whatever reason (e.g. self type isn't
-                // known enough), we *can't* register an obligation and return
-                // an inference variable (since then that obligation would have
-                // bound vars and that's a can of worms). Instead, we just
-                // give up and fall back to pretending like we never tried!
+            ty::Projection(data) => {
+                // If there are escaping bound vars, we temporarily replace the
+                // bound vars with placeholders. Note though, that in the cas
+                // that we still can't project for whatever reason (e.g. self
+                // type isn't known enough), we *can't* register an obligation
+                // and return an inference variable (since then that obligation
+                // would have bound vars and that's a can of worms). Instead,
+                // we just give up and fall back to pretending like we never tried!
 
                 let infcx = self.selcx.infcx();
                 let (data, mapped_regions, mapped_types, mapped_consts) =
                     BoundVarReplacer::replace_bound_vars(infcx, &mut self.universes, data);
+                let data = data.super_fold_with(self);
                 let normalized_ty = opt_normalize_projection_type(
                     self.selcx,
                     self.param_env,
@@ -459,7 +451,7 @@ impl<'a, 'b, 'tcx> TypeFolder<'tcx> for AssocTypeNormalizer<'a, 'b, 'tcx> {
                 )
                 .ok()
                 .flatten()
-                .unwrap_or_else(|| ty);
+                .unwrap_or_else(|| ty.super_fold_with(self));
 
                 let normalized_ty = PlaceholderReplacer::replace_placeholders(
                     infcx,
@@ -479,7 +471,7 @@ impl<'a, 'b, 'tcx> TypeFolder<'tcx> for AssocTypeNormalizer<'a, 'b, 'tcx> {
                 normalized_ty
             }
 
-            _ => ty,
+            _ => ty.super_fold_with(self),
         }
     }
 
@@ -908,6 +900,7 @@ fn opt_normalize_projection_type<'a, 'b, 'tcx>(
             // an impl, where-clause etc) and hence we must
             // re-normalize it
 
+            let projected_ty = selcx.infcx().resolve_vars_if_possible(projected_ty);
             debug!(?projected_ty, ?depth, ?projected_obligations);
 
             let result = if projected_ty.has_projections() {
diff --git a/compiler/rustc_trait_selection/src/traits/query/normalize.rs b/compiler/rustc_trait_selection/src/traits/query/normalize.rs
index 3f6efa03b3a..8f8de24981b 100644
--- a/compiler/rustc_trait_selection/src/traits/query/normalize.rs
+++ b/compiler/rustc_trait_selection/src/traits/query/normalize.rs
@@ -14,7 +14,9 @@ use rustc_infer::traits::Normalized;
 use rustc_middle::mir;
 use rustc_middle::ty::fold::{TypeFoldable, TypeFolder};
 use rustc_middle::ty::subst::Subst;
-use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_middle::ty::{self, Ty, TyCtxt, TypeVisitor};
+
+use std::ops::ControlFlow;
 
 use super::NoSolution;
 
@@ -65,6 +67,14 @@ impl<'cx, 'tcx> AtExt<'tcx> for At<'cx, 'tcx> {
             universes: vec![],
         };
 
+        if value.has_escaping_bound_vars() {
+            let mut max_visitor =
+                MaxEscapingBoundVarVisitor { outer_index: ty::INNERMOST, escaping: 0 };
+            value.visit_with(&mut max_visitor);
+            if max_visitor.escaping > 0 {
+                normalizer.universes.extend((0..max_visitor.escaping).map(|_| None));
+            }
+        }
         let result = value.fold_with(&mut normalizer);
         info!(
             "normalize::<{}>: result={:?} with {} obligations",
@@ -85,6 +95,58 @@ impl<'cx, 'tcx> AtExt<'tcx> for At<'cx, 'tcx> {
     }
 }
 
+/// Visitor to find the maximum escaping bound var
+struct MaxEscapingBoundVarVisitor {
+    // The index which would count as escaping
+    outer_index: ty::DebruijnIndex,
+    escaping: usize,
+}
+
+impl<'tcx> TypeVisitor<'tcx> for MaxEscapingBoundVarVisitor {
+    fn visit_binder<T: TypeFoldable<'tcx>>(
+        &mut self,
+        t: &ty::Binder<'tcx, T>,
+    ) -> ControlFlow<Self::BreakTy> {
+        self.outer_index.shift_in(1);
+        let result = t.super_visit_with(self);
+        self.outer_index.shift_out(1);
+        result
+    }
+
+    #[inline]
+    fn visit_ty(&mut self, t: Ty<'tcx>) -> ControlFlow<Self::BreakTy> {
+        if t.outer_exclusive_binder() > self.outer_index {
+            self.escaping = self
+                .escaping
+                .max(t.outer_exclusive_binder().as_usize() - self.outer_index.as_usize());
+        }
+        ControlFlow::CONTINUE
+    }
+
+    #[inline]
+    fn visit_region(&mut self, r: ty::Region<'tcx>) -> ControlFlow<Self::BreakTy> {
+        match *r {
+            ty::ReLateBound(debruijn, _) if debruijn > self.outer_index => {
+                self.escaping =
+                    self.escaping.max(debruijn.as_usize() - self.outer_index.as_usize());
+            }
+            _ => {}
+        }
+        ControlFlow::CONTINUE
+    }
+
+    fn visit_const(&mut self, ct: &'tcx ty::Const<'tcx>) -> ControlFlow<Self::BreakTy> {
+        match ct.val {
+            ty::ConstKind::Bound(debruijn, _) if debruijn >= self.outer_index => {
+                self.escaping =
+                    self.escaping.max(debruijn.as_usize() - self.outer_index.as_usize());
+                ControlFlow::CONTINUE
+            }
+            _ => ct.super_visit_with(self),
+        }
+    }
+}
+
 struct QueryNormalizer<'cx, 'tcx> {
     infcx: &'cx InferCtxt<'cx, 'tcx>,
     cause: &'cx ObligationCause<'tcx>,
@@ -121,14 +183,25 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for QueryNormalizer<'cx, 'tcx> {
             return ty;
         }
 
-        let ty = ty.super_fold_with(self);
+        // N.b. while we want to call `super_fold_with(self)` on `ty` before
+        // normalization, we wait until we know whether we need to normalize the
+        // current type. If we do, then we only fold the ty *after* replacing bound
+        // vars with placeholders. This means that nested types don't need to replace
+        // bound vars at the current binder level or above. A key assumption here is
+        // that folding the type can't introduce new bound vars.
+
+        // Wrap this in a closure so we don't accidentally return from the outer function
         let res = (|| match *ty.kind() {
-            ty::Opaque(def_id, substs) if !substs.has_escaping_bound_vars() => {
+            ty::Opaque(def_id, substs) => {
                 // Only normalize `impl Trait` after type-checking, usually in codegen.
                 match self.param_env.reveal() {
-                    Reveal::UserFacing => ty,
+                    Reveal::UserFacing => ty.super_fold_with(self),
 
                     Reveal::All => {
+                        // N.b. there is an assumption here all this code can handle
+                        // escaping bound vars.
+
+                        let substs = substs.super_fold_with(self);
                         let recursion_limit = self.tcx().recursion_limit();
                         if !recursion_limit.value_within_limit(self.anon_depth) {
                             let obligation = Obligation::with_depth(
@@ -161,19 +234,11 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for QueryNormalizer<'cx, 'tcx> {
             }
 
             ty::Projection(data) if !data.has_escaping_bound_vars() => {
-                // This is kind of hacky -- we need to be able to
-                // handle normalization within binders because
-                // otherwise we wind up a need to normalize when doing
-                // trait matching (since you can have a trait
-                // obligation like `for<'a> T::B: Fn(&'a i32)`), but
-                // we can't normalize with bound regions in scope. So
-                // far now we just ignore binders but only normalize
-                // if all bound regions are gone (and then we still
-                // have to renormalize whenever we instantiate a
-                // binder). It would be better to normalize in a
-                // binding-aware fashion.
+                // This branch is just an optimization: when we don't have escaping bound vars,
+                // we don't need to replace them with placeholders (see branch below).
 
                 let tcx = self.infcx.tcx;
+                let data = data.super_fold_with(self);
 
                 let mut orig_values = OriginalQueryValues::default();
                 // HACK(matthewjasper) `'static` is special-cased in selection,
@@ -217,22 +282,9 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for QueryNormalizer<'cx, 'tcx> {
                     }
                 }
             }
-            ty::Projection(data) if !data.trait_ref(self.infcx.tcx).has_escaping_bound_vars() => {
-                // See note in `rustc_trait_selection::traits::project`
-
-                // One other point mentioning: In `traits::project`, if a
-                // projection can't be normalized, we return an inference variable
-                // and register an obligation to later resolve that. Here, the query
-                // will just return ambiguity. In both cases, the effect is the same: we only want
-                // to return `ty` because there are bound vars that we aren't yet handling in a more
-                // complete way.
 
-                // `BoundVarReplacer` can't handle escaping bound vars. Ideally, we want this before even calling
-                // `QueryNormalizer`, but some const-generics tests pass escaping bound vars.
-                // Also, use `ty` so we get that sweet `outer_exclusive_binder` optimization
-                assert!(!ty.has_vars_bound_at_or_above(ty::DebruijnIndex::from_usize(
-                    self.universes.len()
-                )));
+            ty::Projection(data) => {
+                // See note in `rustc_trait_selection::traits::project`
 
                 let tcx = self.infcx.tcx;
                 let infcx = self.infcx;
@@ -292,7 +344,7 @@ impl<'cx, 'tcx> TypeFolder<'tcx> for QueryNormalizer<'cx, 'tcx> {
                 )
             }
 
-            _ => ty,
+            _ => ty.super_fold_with(self),
         })();
         self.cache.insert(ty, res);
         res