Auto merge of #141581 - lcnr:fold-clauses, r=<try>

add additional `TypeFlags` fast paths

Some crates, e.g. `diesel`, have items with a lot of where-clauses (more than 150). In these cases checking the `TypeFlags` of the whole `param_env` can be very beneficial.

This adds `fn fold_clauses` to mirror the existing `fn visit_clauses` and then uses this in folders which fold `ParamEnv`s.

Split out from #141451, depends on #141500.

r? `@compiler-errors`

Author: bors

compiler/rustc_infer/src/infer/canonical/canonicalizer.rs

@@ -493,6 +493,14 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
             ct
         }
     }
+
+    fn fold_predicate(&mut self, p: ty::Predicate<'tcx>) -> ty::Predicate<'tcx> {
+        if p.flags().intersects(self.needs_canonical_flags) { p.super_fold_with(self) } else { p }
+    }
+
+    fn fold_clauses(&mut self, c: ty::Clauses<'tcx>) -> ty::Clauses<'tcx> {
+        if c.flags().intersects(self.needs_canonical_flags) { c.super_fold_with(self) } else { c }
+    }
 }
 
 impl<'cx, 'tcx> Canonicalizer<'cx, 'tcx> {

compiler/rustc_infer/src/infer/resolve.rs

@@ -55,6 +55,14 @@ impl<'a, 'tcx> TypeFolder<TyCtxt<'tcx>> for OpportunisticVarResolver<'a, 'tcx> {
             ct.super_fold_with(self)
         }
     }
+
+    fn fold_predicate(&mut self, p: ty::Predicate<'tcx>) -> ty::Predicate<'tcx> {
+        if !p.has_non_region_infer() { p } else { p.super_fold_with(self) }
+    }
+
+    fn fold_clauses(&mut self, c: ty::Clauses<'tcx>) -> ty::Clauses<'tcx> {
+        if !c.has_non_region_infer() { c } else { c.super_fold_with(self) }
+    }
 }
 
 /// The opportunistic region resolver opportunistically resolves regions

compiler/rustc_middle/src/ty/erase_regions.rs

@@ -86,4 +86,12 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for RegionEraserVisitor<'tcx> {
             p
         }
     }
+
+    fn fold_clauses(&mut self, c: ty::Clauses<'tcx>) -> ty::Clauses<'tcx> {
+        if c.has_type_flags(TypeFlags::HAS_BINDER_VARS | TypeFlags::HAS_FREE_REGIONS) {
+            c.super_fold_with(self)
+        } else {
+            c
+        }
+    }
 }

compiler/rustc_middle/src/ty/fold.rs

@@ -177,6 +177,10 @@ where
     fn fold_predicate(&mut self, p: ty::Predicate<'tcx>) -> ty::Predicate<'tcx> {
         if p.has_vars_bound_at_or_above(self.current_index) { p.super_fold_with(self) } else { p }
     }
+
+    fn fold_clauses(&mut self, c: ty::Clauses<'tcx>) -> ty::Clauses<'tcx> {
+        if c.has_vars_bound_at_or_above(self.current_index) { c.super_fold_with(self) } else { c }
+    }
 }
 
 impl<'tcx> TyCtxt<'tcx> {

compiler/rustc_middle/src/ty/predicate.rs

@@ -238,6 +238,8 @@ impl<'tcx> Clause<'tcx> {
     }
 }
 
+impl<'tcx> rustc_type_ir::inherent::Clauses<TyCtxt<'tcx>> for ty::Clauses<'tcx> {}
+
 #[extension(pub trait ExistentialPredicateStableCmpExt<'tcx>)]
 impl<'tcx> ExistentialPredicate<'tcx> {
     /// Compares via an ordering that will not change if modules are reordered or other changes are

compiler/rustc_middle/src/ty/structural_impls.rs

@@ -570,6 +570,19 @@ impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for ty::Clause<'tcx> {
     }
 }
 
+impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for ty::Clauses<'tcx> {
+    fn try_fold_with<F: FallibleTypeFolder<TyCtxt<'tcx>>>(
+        self,
+        folder: &mut F,
+    ) -> Result<Self, F::Error> {
+        folder.try_fold_clauses(self)
+    }
+
+    fn fold_with<F: TypeFolder<TyCtxt<'tcx>>>(self, folder: &mut F) -> Self {
+        folder.fold_clauses(self)
+    }
+}
+
 impl<'tcx> TypeVisitable<TyCtxt<'tcx>> for ty::Predicate<'tcx> {
     fn visit_with<V: TypeVisitor<TyCtxt<'tcx>>>(&self, visitor: &mut V) -> V::Result {
         visitor.visit_predicate(*self)
@@ -615,6 +628,19 @@ impl<'tcx> TypeSuperVisitable<TyCtxt<'tcx>> for ty::Clauses<'tcx> {
     }
 }
 
+impl<'tcx> TypeSuperFoldable<TyCtxt<'tcx>> for ty::Clauses<'tcx> {
+    fn try_super_fold_with<F: FallibleTypeFolder<TyCtxt<'tcx>>>(
+        self,
+        folder: &mut F,
+    ) -> Result<Self, F::Error> {
+        ty::util::try_fold_list(self, folder, |tcx, v| tcx.mk_clauses(v))
+    }
+
+    fn super_fold_with<F: TypeFolder<TyCtxt<'tcx>>>(self, folder: &mut F) -> Self {
+        ty::util::fold_list(self, folder, |tcx, v| tcx.mk_clauses(v))
+    }
+}
+
 impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for ty::Const<'tcx> {
     fn try_fold_with<F: FallibleTypeFolder<TyCtxt<'tcx>>>(
         self,
@@ -775,7 +801,6 @@ macro_rules! list_fold {
 }
 
 list_fold! {
-    ty::Clauses<'tcx> : mk_clauses,
     &'tcx ty::List<ty::PolyExistentialPredicate<'tcx>> : mk_poly_existential_predicates,
     &'tcx ty::List<PlaceElem<'tcx>> : mk_place_elems,
     &'tcx ty::List<ty::Pattern<'tcx>> : mk_patterns,

compiler/rustc_next_trait_solver/src/canonicalizer.rs

@@ -4,12 +4,22 @@ use rustc_type_ir::data_structures::{HashMap, ensure_sufficient_stack};
 use rustc_type_ir::inherent::*;
 use rustc_type_ir::solve::{Goal, QueryInput};
 use rustc_type_ir::{
-    self as ty, Canonical, CanonicalTyVarKind, CanonicalVarKind, InferCtxtLike, Interner,
-    TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitableExt,
+    self as ty, Canonical, CanonicalTyVarKind, CanonicalVarKind, Flags, InferCtxtLike, Interner,
+    TypeFlags, TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitableExt,
 };
 
 use crate::delegate::SolverDelegate;
 
+/// Does this have infer/placeholder/param, free regions or ReErased?
+const NEEDS_CANONICAL: TypeFlags = TypeFlags::from_bits(
+    TypeFlags::HAS_INFER.bits()
+        | TypeFlags::HAS_PLACEHOLDER.bits()
+        | TypeFlags::HAS_PARAM.bits()
+        | TypeFlags::HAS_FREE_REGIONS.bits()
+        | TypeFlags::HAS_RE_ERASED.bits(),
+)
+.unwrap();
+
 /// Whether we're canonicalizing a query input or the query response.
 ///
 /// When canonicalizing an input we're in the context of the caller
@@ -79,7 +89,11 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
             cache: Default::default(),
         };
 
-        let value = value.fold_with(&mut canonicalizer);
+        let value = if value.has_type_flags(NEEDS_CANONICAL) {
+            value.fold_with(&mut canonicalizer)
+        } else {
+            value
+        };
         assert!(!value.has_infer(), "unexpected infer in {value:?}");
         assert!(!value.has_placeholders(), "unexpected placeholders in {value:?}");
         let (max_universe, variables) = canonicalizer.finalize();
@@ -111,7 +125,14 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
 
             cache: Default::default(),
         };
-        let param_env = input.goal.param_env.fold_with(&mut env_canonicalizer);
+
+        let param_env = input.goal.param_env;
+        let param_env = if param_env.has_type_flags(NEEDS_CANONICAL) {
+            param_env.fold_with(&mut env_canonicalizer)
+        } else {
+            param_env
+        };
+
         debug_assert_eq!(env_canonicalizer.binder_index, ty::INNERMOST);
         // Then canonicalize the rest of the input without keeping `'static`
         // while *mostly* reusing the canonicalizer from above.
@@ -134,10 +155,22 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
             cache: Default::default(),
         };
 
-        let predicate = input.goal.predicate.fold_with(&mut rest_canonicalizer);
+        let predicate = input.goal.predicate;
+        let predicate = if predicate.has_type_flags(NEEDS_CANONICAL) {
+            predicate.fold_with(&mut rest_canonicalizer)
+        } else {
+            predicate
+        };
         let goal = Goal { param_env, predicate };
+
+        let predefined_opaques_in_body = input.predefined_opaques_in_body;
         let predefined_opaques_in_body =
-            input.predefined_opaques_in_body.fold_with(&mut rest_canonicalizer);
+            if input.predefined_opaques_in_body.has_type_flags(NEEDS_CANONICAL) {
+                predefined_opaques_in_body.fold_with(&mut rest_canonicalizer)
+            } else {
+                predefined_opaques_in_body
+            };
+
         let value = QueryInput { goal, predefined_opaques_in_body };
 
         assert!(!value.has_infer(), "unexpected infer in {value:?}");
@@ -387,7 +420,11 @@ impl<'a, D: SolverDelegate<Interner = I>, I: Interner> Canonicalizer<'a, D, I> {
             | ty::Alias(_, _)
             | ty::Bound(_, _)
             | ty::Error(_) => {
-                return ensure_sufficient_stack(|| t.super_fold_with(self));
+                return if t.has_type_flags(NEEDS_CANONICAL) {
+                    ensure_sufficient_stack(|| t.super_fold_with(self))
+                } else {
+                    t
+                };
             }
         };
 
@@ -522,11 +559,28 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
             | ty::ConstKind::Unevaluated(_)
             | ty::ConstKind::Value(_)
             | ty::ConstKind::Error(_)
-            | ty::ConstKind::Expr(_) => return c.super_fold_with(self),
+            | ty::ConstKind::Expr(_) => {
+                return if c.has_type_flags(NEEDS_CANONICAL) { c.super_fold_with(self) } else { c };
+            }
         };
 
         let var = self.get_or_insert_bound_var(c, kind);
 
         Const::new_anon_bound(self.cx(), self.binder_index, var)
     }
+
+    fn fold_predicate(&mut self, p: I::Predicate) -> I::Predicate {
+        if p.flags().intersects(NEEDS_CANONICAL) { p.super_fold_with(self) } else { p }
+    }
+
+    fn fold_clauses(&mut self, c: I::Clauses) -> I::Clauses {
+        match self.canonicalize_mode {
+            CanonicalizeMode::Input { keep_static: true }
+            | CanonicalizeMode::Response { max_input_universe: _ } => {}
+            CanonicalizeMode::Input { keep_static: false } => {
+                panic!("erasing 'static in env")
+            }
+        }
+        if c.flags().intersects(NEEDS_CANONICAL) { c.super_fold_with(self) } else { c }
+    }
 }

compiler/rustc_next_trait_solver/src/resolve.rs

@@ -11,7 +11,7 @@ use crate::delegate::SolverDelegate;
 // EAGER RESOLUTION
 
 /// Resolves ty, region, and const vars to their inferred values or their root vars.
-pub struct EagerResolver<'a, D, I = <D as SolverDelegate>::Interner>
+struct EagerResolver<'a, D, I = <D as SolverDelegate>::Interner>
 where
     D: SolverDelegate<Interner = I>,
     I: Interner,
@@ -22,8 +22,20 @@ where
     cache: DelayedMap<I::Ty, I::Ty>,
 }
 
+pub fn eager_resolve_vars<D: SolverDelegate, T: TypeFoldable<D::Interner>>(
+    delegate: &D,
+    value: T,
+) -> T {
+    if value.has_infer() {
+        let mut folder = EagerResolver::new(delegate);
+        value.fold_with(&mut folder)
+    } else {
+        value
+    }
+}
+
 impl<'a, D: SolverDelegate> EagerResolver<'a, D> {
-    pub fn new(delegate: &'a D) -> Self {
+    fn new(delegate: &'a D) -> Self {
         EagerResolver { delegate, cache: Default::default() }
     }
 }
@@ -86,4 +98,12 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for EagerResolv
             }
         }
     }
+
+    fn fold_predicate(&mut self, p: I::Predicate) -> I::Predicate {
+        if p.has_infer() { p.super_fold_with(self) } else { p }
+    }
+
+    fn fold_clauses(&mut self, c: I::Clauses) -> I::Clauses {
+        if c.has_infer() { c.super_fold_with(self) } else { c }
+    }
 }

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

@@ -22,7 +22,7 @@ use tracing::{debug, instrument, trace};
 
 use crate::canonicalizer::Canonicalizer;
 use crate::delegate::SolverDelegate;
-use crate::resolve::EagerResolver;
+use crate::resolve::eager_resolve_vars;
 use crate::solve::eval_ctxt::CurrentGoalKind;
 use crate::solve::{
     CanonicalInput, CanonicalResponse, Certainty, EvalCtxt, ExternalConstraintsData, Goal,
@@ -61,8 +61,7 @@ where
         // 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) =
-            (goal, opaque_types).fold_with(&mut EagerResolver::new(self.delegate));
+        let (goal, opaque_types) = eager_resolve_vars(self.delegate, (goal, opaque_types));
 
         let mut orig_values = Default::default();
         let canonical = Canonicalizer::canonicalize_input(
@@ -157,8 +156,8 @@ where
 
         let external_constraints =
             self.compute_external_query_constraints(certainty, normalization_nested_goals);
-        let (var_values, mut external_constraints) = (self.var_values, external_constraints)
-            .fold_with(&mut EagerResolver::new(self.delegate));
+        let (var_values, mut external_constraints) =
+            eager_resolve_vars(self.delegate, (self.var_values, external_constraints));
 
         // Remove any trivial or duplicated region constraints once we've resolved regions
         let mut unique = HashSet::default();
@@ -469,7 +468,7 @@ where
 {
     let var_values = CanonicalVarValues { var_values: delegate.cx().mk_args(var_values) };
     let state = inspect::State { var_values, data };
-    let state = state.fold_with(&mut EagerResolver::new(delegate));
+    let state = eager_resolve_vars(delegate, state);
     Canonicalizer::canonicalize_response(delegate, max_input_universe, &mut vec![], state)
 }
 

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

@@ -848,6 +848,22 @@ where
                     }
                 }
             }
+
+            fn visit_predicate(&mut self, p: I::Predicate) -> Self::Result {
+                if p.has_non_region_infer() || p.has_placeholders() {
+                    p.super_visit_with(self)
+                } else {
+                    ControlFlow::Continue(())
+                }
+            }
+
+            fn visit_clauses(&mut self, c: I::Clauses) -> Self::Result {
+                if c.has_non_region_infer() || c.has_placeholders() {
+                    c.super_visit_with(self)
+                } else {
+                    ControlFlow::Continue(())
+                }
+            }
         }
 
         let mut visitor = ContainsTermOrNotNameable {

compiler/rustc_trait_selection/src/solve/inspect/analyse.rs

@@ -15,9 +15,9 @@ use rustc_infer::infer::{DefineOpaqueTypes, InferCtxt, InferOk};
 use rustc_macros::extension;
 use rustc_middle::traits::ObligationCause;
 use rustc_middle::traits::solve::{Certainty, Goal, GoalSource, NoSolution, QueryResult};
-use rustc_middle::ty::{TyCtxt, TypeFoldable, VisitorResult, try_visit};
+use rustc_middle::ty::{TyCtxt, VisitorResult, try_visit};
 use rustc_middle::{bug, ty};
-use rustc_next_trait_solver::resolve::EagerResolver;
+use rustc_next_trait_solver::resolve::eager_resolve_vars;
 use rustc_next_trait_solver::solve::inspect::{self, instantiate_canonical_state};
 use rustc_next_trait_solver::solve::{GenerateProofTree, MaybeCause, SolverDelegateEvalExt as _};
 use rustc_span::{DUMMY_SP, Span};
@@ -187,8 +187,7 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
             let _ = term_hack.constrain(infcx, span, param_env);
         }
 
-        let opt_impl_args =
-            opt_impl_args.map(|impl_args| impl_args.fold_with(&mut EagerResolver::new(infcx)));
+        let opt_impl_args = opt_impl_args.map(|impl_args| eager_resolve_vars(infcx, impl_args));
 
         let goals = instantiated_goals
             .into_iter()
@@ -392,7 +391,7 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
             infcx,
             depth,
             orig_values,
-            goal: uncanonicalized_goal.fold_with(&mut EagerResolver::new(infcx)),
+            goal: eager_resolve_vars(infcx, uncanonicalized_goal),
             result,
             evaluation_kind: evaluation.kind,
             normalizes_to_term_hack,

compiler/rustc_type_ir/src/binder.rs

@@ -711,6 +711,14 @@ impl<'a, I: Interner> TypeFolder<I> for ArgFolder<'a, I> {
             c.super_fold_with(self)
         }
     }
+
+    fn fold_predicate(&mut self, p: I::Predicate) -> I::Predicate {
+        if p.has_param() { p.super_fold_with(self) } else { p }
+    }
+
+    fn fold_clauses(&mut self, c: I::Clauses) -> I::Clauses {
+        if c.has_param() { c.super_fold_with(self) } else { c }
+    }
 }
 
 impl<'a, I: Interner> ArgFolder<'a, I> {