Auto merge of #144405 - lcnr:hir-typeck-uniquify, r=BoxyUwU

uniquify root goals during HIR typeck

We need to rely on region identity to deal with hangs such as https://github.com/rust-lang/trait-system-refactor-initiative/issues/210 and to keep the current behavior of `fn try_merge_responses`.

This is a problem as borrowck starts by replacing each *occurrence* of a region with a unique inference variable. This frequently splits a single region during HIR typeck into multiple distinct regions. As we assume goals to always succeed during borrowck, relying on two occurances of a region being identical during HIR typeck causes ICE. See the now fixed examples in https://github.com/rust-lang/trait-system-refactor-initiative/issues/27 and rust-lang/rust#139409.

We've previously tried to avoid this issue by always *uniquifying* regions when canonicalizing goals. This prevents caching subtrees during canonicalization which resulted in hangs for very large types. People rely on such types in practice, which caused us to revert our attempt to reinstate `#[type_length_limit]` in https://github.com/rust-lang/rust/pull/127670. The complete list of changes here:
- rust-lang/rust#107981
- rust-lang/rust#110180
- rust-lang/rust#114117
- rust-lang/rust#130821

After more consideration, all occurrences of such large types need to happen outside of typeck/borrowck. We know this as we already walk over all types in the MIR body when replacing their regions with nll vars.

This PR therefore enables us to rely on region identity inside of the trait solver by exclusively **uniquifying root goals during HIR typeck**. These are the only goals we assume to hold during borrowck. This is insufficient as type inference variables may "hide" regions we later uniquify. Because of this, we now stash proven goals which depend on inference variables in HIR typeck and reprove them after writeback. This closes https://github.com/rust-lang/trait-system-refactor-initiative/issues/127.

This was originally part of rust-lang/rust#144258 but I've moved it into a separate PR. While I believe we need to rely on region identity to fix the performance issues in some way, I don't know whether rust-lang/rust#144258 is the best approach to actually do so. Regardless of how we deal with the hangs however, this change is necessary and desirable regardless.

r? `@compiler-errors` or `@BoxyUwU`

4 files changed, 100 insertions, 34 deletions

diff --git a/compiler/rustc_next_trait_solver/src/canonicalizer.rs b/compiler/rustc_next_trait_solver/src/canonicalizer.rs
index a418aa82100..1bc35e599c7 100644
--- a/compiler/rustc_next_trait_solver/src/canonicalizer.rs
+++ b/compiler/rustc_next_trait_solver/src/canonicalizer.rs
@@ -19,6 +19,20 @@ const NEEDS_CANONICAL: TypeFlags = TypeFlags::from_bits(
 )
 .unwrap();
 
+#[derive(Debug, Clone, Copy)]
+enum CanonicalizeInputKind {
+    /// When canonicalizing the `param_env`, we keep `'static` as merging
+    /// trait candidates relies on it when deciding whether a where-bound
+    /// is trivial.
+    ParamEnv,
+    /// When canonicalizing predicates, we don't keep `'static`. If we're
+    /// currently outside of the trait solver and canonicalize the root goal
+    /// during HIR typeck, we replace each occurance of a region with a
+    /// unique region variable. See the comment on `InferCtxt::in_hir_typeck`
+    /// for more details.
+    Predicate { is_hir_typeck_root_goal: bool },
+}
+
 /// Whether we're canonicalizing a query input or the query response.
 ///
 /// When canonicalizing an input we're in the context of the caller
@@ -26,10 +40,7 @@ const NEEDS_CANONICAL: TypeFlags = TypeFlags::from_bits(
 /// query.
 #[derive(Debug, Clone, Copy)]
 enum CanonicalizeMode {
-    /// When canonicalizing the `param_env`, we keep `'static` as merging
-    /// trait candidates relies on it when deciding whether a where-bound
-    /// is trivial.
-    Input { keep_static: bool },
+    Input(CanonicalizeInputKind),
     /// FIXME: We currently return region constraints referring to
     /// placeholders and inference variables from a binder instantiated
     /// inside of the query.
@@ -122,7 +133,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
                     let mut variables = Vec::new();
                     let mut env_canonicalizer = Canonicalizer {
                         delegate,
-                        canonicalize_mode: CanonicalizeMode::Input { keep_static: true },
+                        canonicalize_mode: CanonicalizeMode::Input(CanonicalizeInputKind::ParamEnv),
 
                         variables: &mut variables,
                         variable_lookup_table: Default::default(),
@@ -154,7 +165,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
         } else {
             let mut env_canonicalizer = Canonicalizer {
                 delegate,
-                canonicalize_mode: CanonicalizeMode::Input { keep_static: true },
+                canonicalize_mode: CanonicalizeMode::Input(CanonicalizeInputKind::ParamEnv),
 
                 variables,
                 variable_lookup_table: Default::default(),
@@ -180,6 +191,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
     pub fn canonicalize_input<P: TypeFoldable<I>>(
         delegate: &'a D,
         variables: &'a mut Vec<I::GenericArg>,
+        is_hir_typeck_root_goal: bool,
         input: QueryInput<I, P>,
     ) -> ty::Canonical<I, QueryInput<I, P>> {
         // First canonicalize the `param_env` while keeping `'static`
@@ -189,7 +201,9 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
         // while *mostly* reusing the canonicalizer from above.
         let mut rest_canonicalizer = Canonicalizer {
             delegate,
-            canonicalize_mode: CanonicalizeMode::Input { keep_static: false },
+            canonicalize_mode: CanonicalizeMode::Input(CanonicalizeInputKind::Predicate {
+                is_hir_typeck_root_goal,
+            }),
 
             variables,
             variable_lookup_table,
@@ -296,7 +310,7 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
         }
     }
 
-    fn cached_fold_ty(&mut self, t: I::Ty) -> I::Ty {
+    fn inner_fold_ty(&mut self, t: I::Ty) -> I::Ty {
         let kind = match t.kind() {
             ty::Infer(i) => match i {
                 ty::TyVar(vid) => {
@@ -413,10 +427,10 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
             // We don't canonicalize `ReStatic` in the `param_env` as we use it
             // when checking whether a `ParamEnv` candidate is global.
             ty::ReStatic => match self.canonicalize_mode {
-                CanonicalizeMode::Input { keep_static: false } => {
+                CanonicalizeMode::Input(CanonicalizeInputKind::Predicate { .. }) => {
                     CanonicalVarKind::Region(ty::UniverseIndex::ROOT)
                 }
-                CanonicalizeMode::Input { keep_static: true }
+                CanonicalizeMode::Input(CanonicalizeInputKind::ParamEnv)
                 | CanonicalizeMode::Response { .. } => return r,
             },
 
@@ -428,12 +442,12 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
             // `ReErased`. We may be able to short-circuit registering region
             // obligations if we encounter a `ReErased` on one side, for example.
             ty::ReErased | ty::ReError(_) => match self.canonicalize_mode {
-                CanonicalizeMode::Input { .. } => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
+                CanonicalizeMode::Input(_) => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
                 CanonicalizeMode::Response { .. } => return r,
             },
 
             ty::ReEarlyParam(_) | ty::ReLateParam(_) => match self.canonicalize_mode {
-                CanonicalizeMode::Input { .. } => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
+                CanonicalizeMode::Input(_) => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
                 CanonicalizeMode::Response { .. } => {
                     panic!("unexpected region in response: {r:?}")
                 }
@@ -441,7 +455,7 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
 
             ty::RePlaceholder(placeholder) => match self.canonicalize_mode {
                 // We canonicalize placeholder regions as existentials in query inputs.
-                CanonicalizeMode::Input { .. } => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
+                CanonicalizeMode::Input(_) => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
                 CanonicalizeMode::Response { max_input_universe } => {
                     // If we have a placeholder region inside of a query, it must be from
                     // a new universe.
@@ -459,9 +473,7 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
                     "region vid should have been resolved fully before canonicalization"
                 );
                 match self.canonicalize_mode {
-                    CanonicalizeMode::Input { keep_static: _ } => {
-                        CanonicalVarKind::Region(ty::UniverseIndex::ROOT)
-                    }
+                    CanonicalizeMode::Input(_) => CanonicalVarKind::Region(ty::UniverseIndex::ROOT),
                     CanonicalizeMode::Response { .. } => {
                         CanonicalVarKind::Region(self.delegate.universe_of_lt(vid).unwrap())
                     }
@@ -469,16 +481,34 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
             }
         };
 
-        let var = self.get_or_insert_bound_var(r, kind);
+        let var = if let CanonicalizeMode::Input(CanonicalizeInputKind::Predicate {
+            is_hir_typeck_root_goal: true,
+        }) = self.canonicalize_mode
+        {
+            let var = ty::BoundVar::from(self.variables.len());
+            self.variables.push(r.into());
+            self.var_kinds.push(kind);
+            var
+        } else {
+            self.get_or_insert_bound_var(r, kind)
+        };
 
         Region::new_anon_bound(self.cx(), self.binder_index, var)
     }
 
     fn fold_ty(&mut self, t: I::Ty) -> I::Ty {
-        if let Some(&ty) = self.cache.get(&(self.binder_index, t)) {
+        if let CanonicalizeMode::Input(CanonicalizeInputKind::Predicate {
+            is_hir_typeck_root_goal: true,
+        }) = self.canonicalize_mode
+        {
+            // If we're canonicalizing a root goal during HIR typeck, we
+            // must not use the `cache` as we want to map each occurrence
+            // of a region to a unique existential variable.
+            self.inner_fold_ty(t)
+        } else if let Some(&ty) = self.cache.get(&(self.binder_index, t)) {
             ty
         } else {
-            let res = self.cached_fold_ty(t);
+            let res = self.inner_fold_ty(t);
             let old = self.cache.insert((self.binder_index, t), res);
             assert_eq!(old, None);
             res
@@ -541,9 +571,9 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
 
     fn fold_clauses(&mut self, c: I::Clauses) -> I::Clauses {
         match self.canonicalize_mode {
-            CanonicalizeMode::Input { keep_static: true }
+            CanonicalizeMode::Input(CanonicalizeInputKind::ParamEnv)
             | CanonicalizeMode::Response { max_input_universe: _ } => {}
-            CanonicalizeMode::Input { keep_static: false } => {
+            CanonicalizeMode::Input(CanonicalizeInputKind::Predicate { .. }) => {
                 panic!("erasing 'static in env")
             }
         }
diff --git a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs
index 5ed316aa6b1..74c5b49ea92 100644
--- a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs
@@ -53,20 +53,19 @@ where
 {
     /// Canonicalizes the goal remembering the original values
     /// for each bound variable.
+    ///
+    /// This expects `goal` and `opaque_types` to be eager resolved.
     pub(super) fn canonicalize_goal(
         &self,
+        is_hir_typeck_root_goal: bool,
         goal: Goal<I, I::Predicate>,
+        opaque_types: Vec<(ty::OpaqueTypeKey<I>, I::Ty)>,
     ) -> (Vec<I::GenericArg>, CanonicalInput<I, I::Predicate>) {
-        // We only care about one entry per `OpaqueTypeKey` here,
-        // so we only canonicalize the lookup table and ignore
-        // duplicate entries.
-        let opaque_types = self.delegate.clone_opaque_types_lookup_table();
-        let (goal, opaque_types) = eager_resolve_vars(self.delegate, (goal, opaque_types));
-
         let mut orig_values = Default::default();
         let canonical = Canonicalizer::canonicalize_input(
             self.delegate,
             &mut orig_values,
+            is_hir_typeck_root_goal,
             QueryInput {
                 goal,
                 predefined_opaques_in_body: self
diff --git a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
index ce9b794d40d..8671cc7c3d3 100644
--- a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
@@ -20,6 +20,7 @@ use super::has_only_region_constraints;
 use crate::coherence;
 use crate::delegate::SolverDelegate;
 use crate::placeholder::BoundVarReplacer;
+use crate::resolve::eager_resolve_vars;
 use crate::solve::inspect::{self, ProofTreeBuilder};
 use crate::solve::search_graph::SearchGraph;
 use crate::solve::ty::may_use_unstable_feature;
@@ -440,6 +441,7 @@ where
             return Ok((
                 NestedNormalizationGoals::empty(),
                 GoalEvaluation {
+                    goal,
                     certainty: Certainty::Maybe(stalled_on.stalled_cause),
                     has_changed: HasChanged::No,
                     stalled_on: Some(stalled_on),
@@ -447,7 +449,16 @@ where
             ));
         }
 
-        let (orig_values, canonical_goal) = self.canonicalize_goal(goal);
+        // We only care about one entry per `OpaqueTypeKey` here,
+        // so we only canonicalize the lookup table and ignore
+        // duplicate entries.
+        let opaque_types = self.delegate.clone_opaque_types_lookup_table();
+        let (goal, opaque_types) = eager_resolve_vars(self.delegate, (goal, opaque_types));
+
+        let is_hir_typeck_root_goal = matches!(goal_evaluation_kind, GoalEvaluationKind::Root)
+            && self.delegate.in_hir_typeck();
+        let (orig_values, canonical_goal) =
+            self.canonicalize_goal(is_hir_typeck_root_goal, goal, opaque_types);
         let mut goal_evaluation =
             self.inspect.new_goal_evaluation(goal, &orig_values, goal_evaluation_kind);
         let canonical_result = self.search_graph.evaluate_goal(
@@ -525,7 +536,10 @@ where
             },
         };
 
-        Ok((normalization_nested_goals, GoalEvaluation { certainty, has_changed, stalled_on }))
+        Ok((
+            normalization_nested_goals,
+            GoalEvaluation { goal, certainty, has_changed, stalled_on },
+        ))
     }
 
     pub(super) fn compute_goal(&mut self, goal: Goal<I, I::Predicate>) -> QueryResult<I> {
@@ -661,7 +675,7 @@ where
 
                 let (
                     NestedNormalizationGoals(nested_goals),
-                    GoalEvaluation { certainty, stalled_on, has_changed: _ },
+                    GoalEvaluation { goal, certainty, stalled_on, has_changed: _ },
                 ) = self.evaluate_goal_raw(
                     GoalEvaluationKind::Nested,
                     source,
@@ -699,7 +713,15 @@ where
                 // FIXME: Do we need to eagerly resolve here? Or should we check
                 // if the cache key has any changed vars?
                 let with_resolved_vars = self.resolve_vars_if_possible(goal);
-                if pred.alias != goal.predicate.as_normalizes_to().unwrap().skip_binder().alias {
+                if pred.alias
+                    != with_resolved_vars
+                        .predicate
+                        .as_normalizes_to()
+                        .unwrap()
+                        .no_bound_vars()
+                        .unwrap()
+                        .alias
+                {
                     unchanged_certainty = None;
                 }
 
@@ -711,7 +733,7 @@ where
                     }
                 }
             } else {
-                let GoalEvaluation { certainty, has_changed, stalled_on } =
+                let GoalEvaluation { goal, certainty, has_changed, stalled_on } =
                     self.evaluate_goal(GoalEvaluationKind::Nested, source, goal, stalled_on)?;
                 if has_changed == HasChanged::Yes {
                     unchanged_certainty = None;
diff --git a/compiler/rustc_next_trait_solver/src/solve/mod.rs b/compiler/rustc_next_trait_solver/src/solve/mod.rs
index 5ea3f0d1061..aec9594b834 100644
--- a/compiler/rustc_next_trait_solver/src/solve/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/mod.rs
@@ -252,8 +252,6 @@ where
             return None;
         }
 
-        // FIXME(-Znext-solver): Add support to merge region constraints in
-        // responses to deal with trait-system-refactor-initiative#27.
         let one = responses[0];
         if responses[1..].iter().all(|&resp| resp == one) {
             return Some(one);
@@ -388,6 +386,23 @@ fn response_no_constraints_raw<I: Interner>(
 
 /// The result of evaluating a goal.
 pub struct GoalEvaluation<I: Interner> {
+    /// The goal we've evaluated. This is the input goal, but potentially with its
+    /// inference variables resolved. This never applies any inference constraints
+    /// from evaluating the goal.
+    ///
+    /// We rely on this to check whether root goals in HIR typeck had an unresolved
+    /// type inference variable in the input. We must not resolve this after evaluating
+    /// the goal as even if the inference variable has been resolved by evaluating the
+    /// goal itself, this goal may still end up failing due to region uniquification
+    /// later on.
+    ///
+    /// This is used as a minor optimization to avoid re-resolving inference variables
+    /// when reevaluating ambiguous goals. E.g. if we've got a goal `?x: Trait` with `?x`
+    /// already being constrained to `Vec<?y>`, then the first evaluation resolves it to
+    /// `Vec<?y>: Trait`. If this goal is still ambiguous and we later resolve `?y` to `u32`,
+    /// then reevaluating this goal now only needs to resolve `?y` while it would otherwise
+    /// have to resolve both `?x` and `?y`,
+    pub goal: Goal<I, I::Predicate>,
     pub certainty: Certainty,
     pub has_changed: HasChanged,
     /// If the [`Certainty`] was `Maybe`, then keep track of whether the goal has changed