From 4bcca3294a2074fa07631c662067d5ccf31b8cfa Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Thu, 13 Jul 2023 18:24:37 +0000
Subject: [PATCH 1/2] Allow escaping bound vars during
 normalize_erasing_regions in new solver

---
 .../rustc_trait_selection/src/solve/mod.rs    |  2 +-
 .../src/solve/normalize.rs                    | 12 +++++-
 .../src/traits/query/normalize.rs             | 41 ++++++++++---------
 ...g-bound-vars-in-writeback-normalization.rs | 18 ++++++++
 4 files changed, 51 insertions(+), 22 deletions(-)
 create mode 100644 tests/ui/traits/new-solver/escaping-bound-vars-in-writeback-normalization.rs

diff --git a/compiler/rustc_trait_selection/src/solve/mod.rs b/compiler/rustc_trait_selection/src/solve/mod.rs
index 9dfb793d0a490..50b2fe7eb332c 100644
--- a/compiler/rustc_trait_selection/src/solve/mod.rs
+++ b/compiler/rustc_trait_selection/src/solve/mod.rs
@@ -37,7 +37,7 @@ pub use eval_ctxt::{
     EvalCtxt, GenerateProofTree, InferCtxtEvalExt, InferCtxtSelectExt, UseGlobalCache,
 };
 pub use fulfill::FulfillmentCtxt;
-pub(crate) use normalize::deeply_normalize;
+pub(crate) use normalize::{deeply_normalize, deeply_normalize_with_skipped_universes};
 
 #[derive(Debug, Clone, Copy)]
 enum SolverMode {
diff --git a/compiler/rustc_trait_selection/src/solve/normalize.rs b/compiler/rustc_trait_selection/src/solve/normalize.rs
index c388850d831e2..524bac19547da 100644
--- a/compiler/rustc_trait_selection/src/solve/normalize.rs
+++ b/compiler/rustc_trait_selection/src/solve/normalize.rs
@@ -19,9 +19,19 @@ use super::FulfillmentCtxt;
 pub(crate) fn deeply_normalize<'tcx, T: TypeFoldable<TyCtxt<'tcx>>>(
     at: At<'_, 'tcx>,
     value: T,
+) -> Result<T, Vec<FulfillmentError<'tcx>>> {
+    deeply_normalize_with_skipped_universes(at, value, vec![])
+}
+
+/// Deeply normalize all aliases in `value`. This does not handle inference and expects
+/// its input to be already fully resolved.
+pub(crate) fn deeply_normalize_with_skipped_universes<'tcx, T: TypeFoldable<TyCtxt<'tcx>>>(
+    at: At<'_, 'tcx>,
+    value: T,
+    universes: Vec<Option<UniverseIndex>>,
 ) -> Result<T, Vec<FulfillmentError<'tcx>>> {
     let fulfill_cx = FulfillmentCtxt::new(at.infcx);
-    let mut folder = NormalizationFolder { at, fulfill_cx, depth: 0, universes: Vec::new() };
+    let mut folder = NormalizationFolder { at, fulfill_cx, depth: 0, universes };
 
     value.try_fold_with(&mut folder)
 }
diff --git a/compiler/rustc_trait_selection/src/traits/query/normalize.rs b/compiler/rustc_trait_selection/src/traits/query/normalize.rs
index 7fe79fd865ce3..5edfa71591559 100644
--- a/compiler/rustc_trait_selection/src/traits/query/normalize.rs
+++ b/compiler/rustc_trait_selection/src/traits/query/normalize.rs
@@ -61,8 +61,27 @@ impl<'cx, 'tcx> QueryNormalizeExt<'tcx> for At<'cx, 'tcx> {
             self.cause,
         );
 
+        // This is actually a consequence by the way `normalize_erasing_regions` works currently.
+        // Because it needs to call the `normalize_generic_arg_after_erasing_regions`, it folds
+        // through tys and consts in a `TypeFoldable`. Importantly, it skips binders, leaving us
+        // with trying to normalize with escaping bound vars.
+        //
+        // Here, we just add the universes that we *would* have created had we passed through the binders.
+        //
+        // We *could* replace escaping bound vars eagerly here, but it doesn't seem really necessary.
+        // The rest of the code is already set up to be lazy about replacing bound vars,
+        // and only when we actually have to normalize.
+        let universes = if value.has_escaping_bound_vars() {
+            let mut max_visitor =
+                MaxEscapingBoundVarVisitor { outer_index: ty::INNERMOST, escaping: 0 };
+            value.visit_with(&mut max_visitor);
+            vec![None; max_visitor.escaping]
+        } else {
+            vec![]
+        };
+
         if self.infcx.next_trait_solver() {
-            match crate::solve::deeply_normalize(self, value) {
+            match crate::solve::deeply_normalize_with_skipped_universes(self, value, universes) {
                 Ok(value) => return Ok(Normalized { value, obligations: vec![] }),
                 Err(_errors) => {
                     return Err(NoSolution);
@@ -81,27 +100,9 @@ impl<'cx, 'tcx> QueryNormalizeExt<'tcx> for At<'cx, 'tcx> {
             obligations: vec![],
             cache: SsoHashMap::new(),
             anon_depth: 0,
-            universes: vec![],
+            universes,
         };
 
-        // This is actually a consequence by the way `normalize_erasing_regions` works currently.
-        // Because it needs to call the `normalize_generic_arg_after_erasing_regions`, it folds
-        // through tys and consts in a `TypeFoldable`. Importantly, it skips binders, leaving us
-        // with trying to normalize with escaping bound vars.
-        //
-        // Here, we just add the universes that we *would* have created had we passed through the binders.
-        //
-        // We *could* replace escaping bound vars eagerly here, but it doesn't seem really necessary.
-        // The rest of the code is already set up to be lazy about replacing bound vars,
-        // and only when we actually have to normalize.
-        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.try_fold_with(&mut normalizer);
         info!(
             "normalize::<{}>: result={:?} with {} obligations",
diff --git a/tests/ui/traits/new-solver/escaping-bound-vars-in-writeback-normalization.rs b/tests/ui/traits/new-solver/escaping-bound-vars-in-writeback-normalization.rs
new file mode 100644
index 0000000000000..29784c32a1b69
--- /dev/null
+++ b/tests/ui/traits/new-solver/escaping-bound-vars-in-writeback-normalization.rs
@@ -0,0 +1,18 @@
+// compile-flags: -Ztrait-solver=next
+// check-pass
+
+trait Trivial {
+    type Assoc;
+}
+
+impl<T: ?Sized> Trivial for T {
+    type Assoc = ();
+}
+
+fn main() {
+    // During writeback, we call `normalize_erasing_regions`, which will walk past
+    // the `for<'a>` binder and try to normalize `<&'a () as Trivial>::Assoc` directly.
+    // We need to handle this case in the new deep normalizer similarly to how it
+    // is handled in the old solver.
+    let x: Option<for<'a> fn(<&'a () as Trivial>::Assoc)> = None;
+}

From 1ef85d82e0904367a603f97c100bcbb23f8ce2ff Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Thu, 13 Jul 2023 18:27:40 +0000
Subject: [PATCH 2/2] assertion, comment

---
 compiler/rustc_trait_selection/src/solve/normalize.rs | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/compiler/rustc_trait_selection/src/solve/normalize.rs b/compiler/rustc_trait_selection/src/solve/normalize.rs
index 524bac19547da..67e6f5246c840 100644
--- a/compiler/rustc_trait_selection/src/solve/normalize.rs
+++ b/compiler/rustc_trait_selection/src/solve/normalize.rs
@@ -20,11 +20,15 @@ pub(crate) fn deeply_normalize<'tcx, T: TypeFoldable<TyCtxt<'tcx>>>(
     at: At<'_, 'tcx>,
     value: T,
 ) -> Result<T, Vec<FulfillmentError<'tcx>>> {
+    assert!(!value.has_escaping_bound_vars());
     deeply_normalize_with_skipped_universes(at, value, vec![])
 }
 
 /// Deeply normalize all aliases in `value`. This does not handle inference and expects
 /// its input to be already fully resolved.
+///
+/// Additionally takes a list of universes which represents the binders which have been
+/// entered before passing `value` to the function.
 pub(crate) fn deeply_normalize_with_skipped_universes<'tcx, T: TypeFoldable<TyCtxt<'tcx>>>(
     at: At<'_, 'tcx>,
     value: T,