uniquify root goals during HIR typeck

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");
             }
         }
     }