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 rust-lang/trait-system-refactor-initiative#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 rust-lang/trait-system-refactor-initiative#27 and #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 #127670. The complete list of changes here:
- #107981
- #110180
- #114117
- #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 rust-lang/trait-system-refactor-initiative#127.

This was originally part of #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 #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`

Author: bors

compiler/rustc_hir_typeck/src/lib.rs

@@ -53,7 +53,7 @@ use rustc_hir_analysis::check::{check_abi, check_custom_abi};
 use rustc_hir_analysis::hir_ty_lowering::HirTyLowerer;
 use rustc_infer::traits::{ObligationCauseCode, ObligationInspector, WellFormedLoc};
 use rustc_middle::query::Providers;
-use rustc_middle::ty::{self, Ty, TyCtxt};
+use rustc_middle::ty::{self, Ty, TyCtxt, TypeVisitableExt};
 use rustc_middle::{bug, span_bug};
 use rustc_session::config;
 use rustc_span::Span;
@@ -259,6 +259,21 @@ fn typeck_with_inspect<'tcx>(
 
     let typeck_results = fcx.resolve_type_vars_in_body(body);
 
+    // Handle potentially region dependent goals, see `InferCtxt::in_hir_typeck`.
+    if let None = fcx.infcx.tainted_by_errors() {
+        for obligation in fcx.take_hir_typeck_potentially_region_dependent_goals() {
+            let obligation = fcx.resolve_vars_if_possible(obligation);
+            if obligation.has_non_region_infer() {
+                bug!("unexpected inference variable after writeback: {obligation:?}");
+            }
+            fcx.register_predicate(obligation);
+        }
+        fcx.select_obligations_where_possible(|_| {});
+        if let None = fcx.infcx.tainted_by_errors() {
+            fcx.report_ambiguity_errors();
+        }
+    }
+
     fcx.detect_opaque_types_added_during_writeback();
 
     // Consistency check our TypeckResults instance can hold all ItemLocalIds

compiler/rustc_hir_typeck/src/typeck_root_ctxt.rs

@@ -85,8 +85,11 @@ impl<'tcx> TypeckRootCtxt<'tcx> {
     pub(crate) fn new(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> Self {
         let hir_owner = tcx.local_def_id_to_hir_id(def_id).owner;
 
-        let infcx =
-            tcx.infer_ctxt().ignoring_regions().build(TypingMode::typeck_for_body(tcx, def_id));
+        let infcx = tcx
+            .infer_ctxt()
+            .ignoring_regions()
+            .in_hir_typeck()
+            .build(TypingMode::typeck_for_body(tcx, def_id));
         let typeck_results = RefCell::new(ty::TypeckResults::new(hir_owner));
         let fulfillment_cx = RefCell::new(<dyn TraitEngine<'_, _>>::new(&infcx));
 

compiler/rustc_infer/src/infer/at.rs

@@ -71,6 +71,7 @@ impl<'tcx> InferCtxt<'tcx> {
             tcx: self.tcx,
             typing_mode: self.typing_mode,
             considering_regions: self.considering_regions,
+            in_hir_typeck: self.in_hir_typeck,
             skip_leak_check: self.skip_leak_check,
             inner: self.inner.clone(),
             lexical_region_resolutions: self.lexical_region_resolutions.clone(),
@@ -95,6 +96,7 @@ impl<'tcx> InferCtxt<'tcx> {
             tcx: self.tcx,
             typing_mode,
             considering_regions: self.considering_regions,
+            in_hir_typeck: self.in_hir_typeck,
             skip_leak_check: self.skip_leak_check,
             inner: self.inner.clone(),
             lexical_region_resolutions: self.lexical_region_resolutions.clone(),

compiler/rustc_infer/src/infer/context.rs

@@ -22,6 +22,10 @@ impl<'tcx> rustc_type_ir::InferCtxtLike for InferCtxt<'tcx> {
         self.next_trait_solver
     }
 
+    fn in_hir_typeck(&self) -> bool {
+        self.in_hir_typeck
+    }
+
     fn typing_mode(&self) -> ty::TypingMode<'tcx> {
         self.typing_mode()
     }

compiler/rustc_infer/src/infer/mod.rs

@@ -37,10 +37,11 @@ use snapshot::undo_log::InferCtxtUndoLogs;
 use tracing::{debug, instrument};
 use type_variable::TypeVariableOrigin;
 
-use crate::infer::region_constraints::UndoLog;
+use crate::infer::snapshot::undo_log::UndoLog;
 use crate::infer::unify_key::{ConstVariableOrigin, ConstVariableValue, ConstVidKey};
 use crate::traits::{
-    self, ObligationCause, ObligationInspector, PredicateObligations, TraitEngine,
+    self, ObligationCause, ObligationInspector, PredicateObligation, PredicateObligations,
+    TraitEngine,
 };
 
 pub mod at;
@@ -156,6 +157,12 @@ pub struct InferCtxtInner<'tcx> {
     /// which may cause types to no longer be considered well-formed.
     region_assumptions: Vec<ty::ArgOutlivesPredicate<'tcx>>,
 
+    /// `-Znext-solver`: Successfully proven goals during HIR typeck which
+    /// reference inference variables and get reproven after writeback.
+    ///
+    /// See the documentation of `InferCtxt::in_hir_typeck` for more details.
+    hir_typeck_potentially_region_dependent_goals: Vec<PredicateObligation<'tcx>>,
+
     /// Caches for opaque type inference.
     opaque_type_storage: OpaqueTypeStorage<'tcx>,
 }
@@ -173,6 +180,7 @@ impl<'tcx> InferCtxtInner<'tcx> {
             region_constraint_storage: Some(Default::default()),
             region_obligations: Default::default(),
             region_assumptions: Default::default(),
+            hir_typeck_potentially_region_dependent_goals: Default::default(),
             opaque_type_storage: Default::default(),
         }
     }
@@ -244,9 +252,29 @@ pub struct InferCtxt<'tcx> {
     typing_mode: TypingMode<'tcx>,
 
     /// Whether this inference context should care about region obligations in
-    /// the root universe. Most notably, this is used during hir typeck as region
+    /// the root universe. Most notably, this is used during HIR typeck as region
     /// solving is left to borrowck instead.
     pub considering_regions: bool,
+    /// `-Znext-solver`: Whether this inference context is used by HIR typeck. If so, we
+    /// need to make sure we don't rely on region identity in the trait solver or when
+    /// relating types. This is necessary as borrowck starts by replacing each occurrence of a
+    /// free region with a unique inference variable. If HIR typeck ends up depending on two
+    /// regions being equal we'd get unexpected mismatches between HIR typeck and MIR typeck,
+    /// resulting in an ICE.
+    ///
+    /// The trait solver sometimes depends on regions being identical. As a concrete example
+    /// the trait solver ignores other candidates if one candidate exists without any constraints.
+    /// The goal `&'a u32: Equals<&'a u32>` has no constraints right now. If we replace each
+    /// occurrence of `'a` with a unique region the goal now equates these regions. See
+    /// the tests in trait-system-refactor-initiative#27 for concrete examples.
+    ///
+    /// We handle this by *uniquifying* region when canonicalizing root goals during HIR typeck.
+    /// This is still insufficient as inference variables may *hide* region variables, so e.g.
+    /// `dyn TwoSuper<?x, ?x>: Super<?x>` may hold but MIR typeck could end up having to prove
+    /// `dyn TwoSuper<&'0 (), &'1 ()>: Super<&'2 ()>` which is now ambiguous. Because of this we
+    /// stash all successfully proven goals which reference inference variables and then reprove
+    /// them after writeback.
+    pub in_hir_typeck: bool,
 
     /// If set, this flag causes us to skip the 'leak check' during
     /// higher-ranked subtyping operations. This flag is a temporary one used
@@ -506,6 +534,7 @@ pub struct TypeOutlivesConstraint<'tcx> {
 pub struct InferCtxtBuilder<'tcx> {
     tcx: TyCtxt<'tcx>,
     considering_regions: bool,
+    in_hir_typeck: bool,
     skip_leak_check: bool,
     /// Whether we should use the new trait solver in the local inference context,
     /// which affects things like which solver is used in `predicate_may_hold`.
@@ -518,6 +547,7 @@ impl<'tcx> TyCtxt<'tcx> {
         InferCtxtBuilder {
             tcx: self,
             considering_regions: true,
+            in_hir_typeck: false,
             skip_leak_check: false,
             next_trait_solver: self.next_trait_solver_globally(),
         }
@@ -535,6 +565,11 @@ impl<'tcx> InferCtxtBuilder<'tcx> {
         self
     }
 
+    pub fn in_hir_typeck(mut self) -> Self {
+        self.in_hir_typeck = true;
+        self
+    }
+
     pub fn skip_leak_check(mut self, skip_leak_check: bool) -> Self {
         self.skip_leak_check = skip_leak_check;
         self
@@ -568,12 +603,18 @@ impl<'tcx> InferCtxtBuilder<'tcx> {
     }
 
     pub fn build(&mut self, typing_mode: TypingMode<'tcx>) -> InferCtxt<'tcx> {
-        let InferCtxtBuilder { tcx, considering_regions, skip_leak_check, next_trait_solver } =
-            *self;
+        let InferCtxtBuilder {
+            tcx,
+            considering_regions,
+            in_hir_typeck,
+            skip_leak_check,
+            next_trait_solver,
+        } = *self;
         InferCtxt {
             tcx,
             typing_mode,
             considering_regions,
+            in_hir_typeck,
             skip_leak_check,
             inner: RefCell::new(InferCtxtInner::new()),
             lexical_region_resolutions: RefCell::new(None),
@@ -978,6 +1019,22 @@ impl<'tcx> InferCtxt<'tcx> {
         }
     }
 
+    pub fn push_hir_typeck_potentially_region_dependent_goal(
+        &self,
+        goal: PredicateObligation<'tcx>,
+    ) {
+        let mut inner = self.inner.borrow_mut();
+        inner.undo_log.push(UndoLog::PushHirTypeckPotentiallyRegionDependentGoal);
+        inner.hir_typeck_potentially_region_dependent_goals.push(goal);
+    }
+
+    pub fn take_hir_typeck_potentially_region_dependent_goals(
+        &self,
+    ) -> Vec<PredicateObligation<'tcx>> {
+        assert!(!self.in_snapshot(), "cannot take goals in a snapshot");
+        std::mem::take(&mut self.inner.borrow_mut().hir_typeck_potentially_region_dependent_goals)
+    }
+
     pub fn ty_to_string(&self, t: Ty<'tcx>) -> String {
         self.resolve_vars_if_possible(t).to_string()
     }

compiler/rustc_infer/src/infer/outlives/obligations.rs

@@ -177,6 +177,7 @@ impl<'tcx> InferCtxt<'tcx> {
     }
 
     pub fn take_registered_region_assumptions(&self) -> Vec<ty::ArgOutlivesPredicate<'tcx>> {
+        assert!(!self.in_snapshot(), "cannot take registered region assumptions in a snapshot");
         std::mem::take(&mut self.inner.borrow_mut().region_assumptions)
     }
 

compiler/rustc_infer/src/infer/snapshot/undo_log.rs

@@ -29,6 +29,7 @@ pub(crate) enum UndoLog<'tcx> {
     ProjectionCache(traits::UndoLog<'tcx>),
     PushTypeOutlivesConstraint,
     PushRegionAssumption,
+    PushHirTypeckPotentiallyRegionDependentGoal,
 }
 
 macro_rules! impl_from {
@@ -79,7 +80,12 @@ impl<'tcx> Rollback<UndoLog<'tcx>> for InferCtxtInner<'tcx> {
                 assert_matches!(popped, Some(_), "pushed region constraint but could not pop it");
             }
             UndoLog::PushRegionAssumption => {
-                self.region_assumptions.pop();
+                let popped = self.region_assumptions.pop();
+                assert_matches!(popped, Some(_), "pushed region assumption but could not pop it");
+            }
+            UndoLog::PushHirTypeckPotentiallyRegionDependentGoal => {
+                let popped = self.hir_typeck_potentially_region_dependent_goals.pop();
+                assert_matches!(popped, Some(_), "pushed goal but could not pop it");
             }
         }
     }