Auto merge of #144258 - lcnr:stupid-fast-path, r=<try>

`compute_trait_goal` structural equality fast path

Fixes the hang in rayon and rust-lang/trait-system-refactor-initiative#210.

This changes means we now have a far stronger reliance on structural equality making rust-lang/trait-system-refactor-initiative#127 a significantly larger problem. However, I've come to believe that fast paths based on structural equality are just too desirable.

Fixes #139409.

r? `@compiler-errors`

Author: rust-bors[bot]

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

@@ -247,6 +247,7 @@ pub struct InferCtxt<'tcx> {
     /// the root universe. Most notably, this is used during hir typeck as region
     /// solving is left to borrowck instead.
     pub considering_regions: bool,
+    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 +507,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 +520,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 +538,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 +576,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),

compiler/rustc_next_trait_solver/src/canonicalizer.rs

@@ -19,6 +19,12 @@ const NEEDS_CANONICAL: TypeFlags = TypeFlags::from_bits(
 )
 .unwrap();
 
+#[derive(Debug, Clone, Copy)]
+enum CanonicalizeInputKind {
+    ParamEnv,
+    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
@@ -29,7 +35,7 @@ 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 +128,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 +160,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 +186,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 +196,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,
@@ -413,10 +422,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,20 +437,20 @@ 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:?}")
                 }
             },
 
             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,23 +468,36 @@ 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())
                     }
                 }
             }
         };
 
-        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
+        {
+            self.cached_fold_ty(t)
+        } else if let Some(&ty) = self.cache.get(&(self.binder_index, t)) {
             ty
         } else {
             let res = self.cached_fold_ty(t);
@@ -541,9 +563,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")
             }
         }

compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs

@@ -55,6 +55,7 @@ where
     /// for each bound variable.
     pub(super) fn canonicalize_goal(
         &self,
+        is_hir_typeck_root_goal: bool,
         goal: Goal<I, I::Predicate>,
     ) -> (Vec<I::GenericArg>, CanonicalInput<I, I::Predicate>) {
         // We only care about one entry per `OpaqueTypeKey` here,
@@ -67,6 +68,7 @@ where
         let canonical = Canonicalizer::canonicalize_input(
             self.delegate,
             &mut orig_values,
+            is_hir_typeck_root_goal,
             QueryInput {
                 goal,
                 predefined_opaques_in_body: self

compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs

@@ -197,7 +197,14 @@ where
             self.cx().recursion_limit(),
             GenerateProofTree::No,
             span,
-            |ecx| ecx.evaluate_goal(GoalEvaluationKind::Root, GoalSource::Misc, goal, stalled_on),
+            |ecx| {
+                ecx.evaluate_goal(
+                    GoalEvaluationKind::Root { in_hir_typeck: self.in_hir_typeck() },
+                    GoalSource::Misc,
+                    goal,
+                    stalled_on,
+                )
+            },
         )
         .0
     }
@@ -209,7 +216,12 @@ where
     ) -> bool {
         self.probe(|| {
             EvalCtxt::enter_root(self, root_depth, GenerateProofTree::No, I::Span::dummy(), |ecx| {
-                ecx.evaluate_goal(GoalEvaluationKind::Root, GoalSource::Misc, goal, None)
+                ecx.evaluate_goal(
+                    GoalEvaluationKind::Root { in_hir_typeck: self.in_hir_typeck() },
+                    GoalSource::Misc,
+                    goal,
+                    None,
+                )
             })
             .0
         })
@@ -230,7 +242,14 @@ where
             self.cx().recursion_limit(),
             GenerateProofTree::Yes,
             span,
-            |ecx| ecx.evaluate_goal_raw(GoalEvaluationKind::Root, GoalSource::Misc, goal, None),
+            |ecx| {
+                ecx.evaluate_goal_raw(
+                    GoalEvaluationKind::Root { in_hir_typeck: self.in_hir_typeck() },
+                    GoalSource::Misc,
+                    goal,
+                    None,
+                )
+            },
         );
         (result, proof_tree.unwrap())
     }
@@ -447,7 +466,10 @@ where
             ));
         }
 
-        let (orig_values, canonical_goal) = self.canonicalize_goal(goal);
+        let is_hir_typeck_root_goal =
+            matches!(goal_evaluation_kind, GoalEvaluationKind::Root { in_hir_typeck: true });
+
+        let (orig_values, canonical_goal) = self.canonicalize_goal(is_hir_typeck_root_goal, goal);
         let mut goal_evaluation =
             self.inspect.new_goal_evaluation(goal, &orig_values, goal_evaluation_kind);
         let canonical_result = self.search_graph.evaluate_goal(

compiler/rustc_next_trait_solver/src/solve/inspect/build.rs

@@ -237,7 +237,7 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> ProofTreeBuilder<D> {
         kind: GoalEvaluationKind,
     ) -> ProofTreeBuilder<D> {
         self.opt_nested(|| match kind {
-            GoalEvaluationKind::Root => Some(WipGoalEvaluation {
+            GoalEvaluationKind::Root { in_hir_typeck: _ } => Some(WipGoalEvaluation {
                 uncanonicalized_goal,
                 orig_values: orig_values.to_vec(),
                 encountered_overflow: false,

compiler/rustc_next_trait_solver/src/solve/mod.rs

@@ -43,7 +43,7 @@ const FIXPOINT_STEP_LIMIT: usize = 8;
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]
 enum GoalEvaluationKind {
-    Root,
+    Root { in_hir_typeck: bool },
     Nested,
 }
 

compiler/rustc_next_trait_solver/src/solve/trait_goals.rs

@@ -1428,6 +1428,21 @@ where
         &mut self,
         goal: Goal<I, TraitPredicate<I>>,
     ) -> Result<(CanonicalResponse<I>, Option<TraitGoalProvenVia>), NoSolution> {
+        if goal
+            .param_env
+            .caller_bounds()
+            .iter()
+            .filter_map(|c| c.as_trait_clause())
+            .filter_map(|c| c.no_bound_vars())
+            .any(|p| p == goal.predicate)
+        {
+            let candidate = self
+                .probe_trait_candidate(CandidateSource::ParamEnv(ParamEnvSource::NonGlobal))
+                .enter(|ecx| ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes))
+                .unwrap();
+            return Ok((candidate.result, Some(TraitGoalProvenVia::ParamEnv)));
+        }
+
         let candidates = self.assemble_and_evaluate_candidates(goal, AssembleCandidatesFrom::All);
         self.merge_trait_candidates(candidates)
     }