Rollup merge of #162461 - jdonszelmann:restrict-fold-predicate, r=BoxyUwU

limit the api of `fold_predicate` and `visit_predicate`

r? @lcnr or anyone in @rust-lang/initiative-trait-system-refactor

In the future, we may want to start compressing clauses in the `ParamEnv`. One major problem was that we are leaking too many implementation details in `fold_predicate` and `visit_predicate`. Almost no code actually cares about dealing with an actual predicate there. Instead, what really matters is the type flags on a predicate for example. As such, this PR majorly limits the API that is exposed to a folder, hiding the underlying predicate data structure used.

As an example why this matters: in the case of compressed clauses, this will mean we won't need to "decompress" them. Instead, we can just fold over the self type, not telling folders whether the predicate was or was not compressed at all.

> [!NOTE]
> I've not used an LLM for any part of this PR, or any other PR I make. This includes any related work like research.

Author: JonathanBrouwer

compiler/rustc_hir_analysis/src/collect/item_bounds.rs

@@ -2,8 +2,8 @@ use rustc_data_structures::fx::{FxIndexMap, FxIndexSet};
 use rustc_hir as hir;
 use rustc_infer::traits::util;
 use rustc_middle::ty::{
-    self, GenericArgs, Ty, TyCtxt, TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitableExt,
-    Upcast, shift_vars,
+    self, GenericArgs, PredicateProxy, Ty, TyCtxt, TypeFoldable, TypeFolder, TypeSuperFoldable,
+    TypeVisitableExt, Upcast, shift_vars,
 };
 use rustc_middle::{bug, span_bug};
 use rustc_span::Span;
@@ -347,7 +347,7 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for MapAndCompressBoundVars<'tcx> {
         }
     }
 
-    fn fold_predicate(&mut self, p: ty::Predicate<'tcx>) -> ty::Predicate<'tcx> {
+    fn fold_predicate<P: PredicateProxy<TyCtxt<'tcx>>>(&mut self, p: P) -> P {
         if !p.has_bound_vars() { p } else { p.super_fold_with(self) }
     }
 }

compiler/rustc_hir_typeck/src/writeback.rs

@@ -21,9 +21,9 @@ use rustc_infer::traits::solve::Goal;
 use rustc_middle::traits::ObligationCause;
 use rustc_middle::ty::adjustment::{Adjust, Adjustment, PointerCoercion};
 use rustc_middle::ty::{
-    self, DefiningScopeKind, DefinitionSiteHiddenType, Flags, Ty, TyCtxt, TypeFoldable, TypeFolder,
-    TypeSuperFoldable, TypeSuperVisitable, TypeVisitable, TypeVisitableExt, TypeVisitor,
-    Unnormalized, fold_regions,
+    self, DefiningScopeKind, DefinitionSiteHiddenType, Flags, PredicateProxy, Ty, TyCtxt,
+    TypeFoldable, TypeFolder, TypeSuperFoldable, TypeSuperVisitable, TypeVisitable,
+    TypeVisitableExt, TypeVisitor, Unnormalized, fold_regions,
 };
 use rustc_span::Span;
 use rustc_trait_selection::error_reporting::infer::need_type_info::TypeAnnotationNeeded;
@@ -1032,7 +1032,7 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Resolver<'cx, 'tcx> {
         self.handle_term(ct, ty::Const::outer_exclusive_binder, ty::Const::new_error)
     }
 
-    fn fold_predicate(&mut self, predicate: ty::Predicate<'tcx>) -> ty::Predicate<'tcx> {
+    fn fold_predicate<P: PredicateProxy<TyCtxt<'tcx>>>(&mut self, predicate: P) -> P {
         assert!(
             !self.should_normalize,
             "normalizing predicates in writeback is not generally sound"

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

@@ -13,6 +13,7 @@ use rustc_middle::ty::{
     self, BoundVar, Flags, GenericArg, InferConst, List, Ty, TyCtxt, TypeFlags, TypeFoldable,
     TypeFolder, TypeSuperFoldable, TypeVisitableExt, TypingModeEqWrapper,
 };
+use rustc_type_ir::PredicateProxy;
 use smallvec::SmallVec;
 use tracing::debug;
 
@@ -483,7 +484,7 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
         }
     }
 
-    fn fold_predicate(&mut self, p: ty::Predicate<'tcx>) -> ty::Predicate<'tcx> {
+    fn fold_predicate<P: PredicateProxy<TyCtxt<'tcx>>>(&mut self, p: P) -> P {
         if p.flags().intersects(self.needs_canonical_flags) { p.super_fold_with(self) } else { p }
     }
 

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

@@ -11,6 +11,7 @@ use rustc_middle::ty::{
     self, DelayedMap, Ty, TyCtxt, TypeFlags, TypeFoldable, TypeFolder, TypeSuperFoldable,
     TypeSuperVisitable, TypeVisitable, TypeVisitableExt, TypeVisitor,
 };
+use rustc_type_ir::PredicateProxy;
 
 use crate::infer::canonical::{Canonical, CanonicalVarValues};
 
@@ -124,7 +125,7 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for CanonicalInstantiator<'tcx> {
         }
     }
 
-    fn fold_predicate(&mut self, p: ty::Predicate<'tcx>) -> ty::Predicate<'tcx> {
+    fn fold_predicate<P: PredicateProxy<TyCtxt<'tcx>>>(&mut self, p: P) -> P {
         if p.has_type_flags(TypeFlags::HAS_CANONICAL_BOUND) { p.super_fold_with(self) } else { p }
     }
 

compiler/rustc_infer/src/infer/resolve.rs

@@ -3,6 +3,7 @@ use rustc_middle::ty::{
     self, Const, DelayedMap, FallibleTypeFolder, InferConst, Ty, TyCtxt, TypeFoldable, TypeFolder,
     TypeSuperFoldable, TypeVisitableExt,
 };
+use rustc_type_ir::PredicateProxy;
 
 use super::{FixupError, FixupResult, InferCtxt};
 use crate::infer::TyOrConstInferVar;
@@ -57,7 +58,7 @@ impl<'a, 'tcx> TypeFolder<TyCtxt<'tcx>> for OpportunisticVarResolver<'a, 'tcx> {
         }
     }
 
-    fn fold_predicate(&mut self, p: ty::Predicate<'tcx>) -> ty::Predicate<'tcx> {
+    fn fold_predicate<P: PredicateProxy<TyCtxt<'tcx>>>(&mut self, p: P) -> P {
         if !p.has_non_region_infer() { p } else { p.super_fold_with(self) }
     }
 

compiler/rustc_middle/src/ty/erase_regions.rs

@@ -1,3 +1,4 @@
+use rustc_type_ir::PredicateProxy;
 use tracing::debug;
 
 use crate::query::Providers;
@@ -79,7 +80,7 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for RegionEraserAndAnonymizerVisitor<'tcx> {
         }
     }
 
-    fn fold_predicate(&mut self, p: ty::Predicate<'tcx>) -> ty::Predicate<'tcx> {
+    fn fold_predicate<P: PredicateProxy<TyCtxt<'tcx>>>(&mut self, p: P) -> P {
         if p.has_type_flags(TypeFlags::HAS_BINDER_VARS | TypeFlags::HAS_FREE_REGIONS) {
             p.super_fold_with(self)
         } else {

compiler/rustc_middle/src/ty/fold.rs

@@ -1,5 +1,6 @@
 use rustc_data_structures::fx::FxIndexMap;
 use rustc_hir::def_id::DefId;
+use rustc_type_ir::PredicateProxy;
 use rustc_type_ir::data_structures::DelayedMap;
 
 use crate::ty::{
@@ -180,7 +181,7 @@ where
         }
     }
 
-    fn fold_predicate(&mut self, p: ty::Predicate<'tcx>) -> ty::Predicate<'tcx> {
+    fn fold_predicate<P: PredicateProxy<TyCtxt<'tcx>>>(&mut self, p: P) -> P {
         if p.has_vars_bound_at_or_above(self.current_index) { p.super_fold_with(self) } else { p }
     }
 

compiler/rustc_middle/src/ty/structural_impls.rs

@@ -9,14 +9,14 @@ use rustc_abi::TyAndLayout;
 use rustc_hir::def::Namespace;
 use rustc_hir::def_id::LocalDefId;
 use rustc_span::Spanned;
-use rustc_type_ir::{ConstKind, TypeFolder, VisitorResult, try_visit};
+use rustc_type_ir::{ConstKind, PredicateProxy, TypeFolder, Upcast, VisitorResult, try_visit};
 
 use super::{GenericArg, GenericArgKind, Pattern};
 use crate::mir::PlaceElem;
 use crate::ty::print::{FmtPrinter, Printer, with_no_trimmed_paths};
 use crate::ty::{
-    self, FallibleTypeFolder, Lift, Term, TermKind, Ty, TyCtxt, TypeFoldable, TypeSuperFoldable,
-    TypeSuperVisitable, TypeVisitable, TypeVisitor,
+    self, Binder, FallibleTypeFolder, Lift, ProjectionClause, Term, TermKind, Ty, TyCtxt,
+    TypeFoldable, TypeSuperFoldable, TypeSuperVisitable, TypeVisitable, TypeVisitor,
 };
 
 impl fmt::Debug for ty::TraitDef {
@@ -491,17 +491,84 @@ impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for ty::Predicate<'tcx> {
     }
 }
 
+impl<'tcx> PredicateProxy<TyCtxt<'tcx>> for ty::Predicate<'tcx> {
+    fn allow_normalization(&self) -> bool {
+        rustc_type_ir::inherent::Predicate::allow_normalization(*self)
+    }
+
+    fn map_projection(
+        self,
+        tcx: TyCtxt<'tcx>,
+        f: impl FnOnce(Binder<'tcx, ProjectionClause<'tcx>>) -> Binder<'tcx, ProjectionClause<'tcx>>,
+    ) -> Option<Self> {
+        self.as_projection_clause().map(|kind| f(kind).upcast(tcx))
+    }
+
+    fn clause_kind_unchecked(&self) -> Option<ty::Binder<'tcx, ty::ClauseKind<'tcx>>> {
+        self.as_clause().map(|clause| clause.kind())
+    }
+}
+
 // FIXME(clause): This is wonky
 impl<'tcx> TypeFoldable<TyCtxt<'tcx>> for ty::Clause<'tcx> {
     fn try_fold_with<F: FallibleTypeFolder<TyCtxt<'tcx>>>(
         self,
         folder: &mut F,
     ) -> Result<Self, F::Error> {
-        Ok(folder.try_fold_predicate(self.as_predicate())?.expect_clause())
+        Ok(folder.try_fold_predicate(self)?)
     }
 
     fn fold_with<F: TypeFolder<TyCtxt<'tcx>>>(self, folder: &mut F) -> Self {
-        folder.fold_predicate(self.as_predicate()).expect_clause()
+        folder.fold_predicate(self)
+    }
+}
+
+// follow `Predicate`'s implementation (by deferring to it)
+impl<'tcx> TypeSuperFoldable<TyCtxt<'tcx>> for ty::Clause<'tcx> {
+    fn try_super_fold_with<F: FallibleTypeFolder<TyCtxt<'tcx>>>(
+        self,
+        folder: &mut F,
+    ) -> Result<Self, F::Error> {
+        <ty::Predicate<'_> as TypeSuperFoldable<TyCtxt<'tcx>>>::try_super_fold_with(
+            self.as_predicate(),
+            folder,
+        )
+        .map(|i| i.expect_clause())
+    }
+
+    fn super_fold_with<F: TypeFolder<TyCtxt<'tcx>>>(self, folder: &mut F) -> Self {
+        <ty::Predicate<'_> as TypeSuperFoldable<TyCtxt<'tcx>>>::super_fold_with(
+            self.as_predicate(),
+            folder,
+        )
+        .expect_clause()
+    }
+}
+
+impl<'tcx> TypeSuperVisitable<TyCtxt<'tcx>> for ty::Clause<'tcx> {
+    fn super_visit_with<V: TypeVisitor<TyCtxt<'tcx>>>(&self, visitor: &mut V) -> V::Result {
+        <ty::Predicate<'_> as TypeSuperVisitable<TyCtxt<'tcx>>>::super_visit_with(
+            &self.as_predicate(),
+            visitor,
+        )
+    }
+}
+
+impl<'tcx> PredicateProxy<TyCtxt<'tcx>> for ty::Clause<'tcx> {
+    fn allow_normalization(&self) -> bool {
+        self.as_predicate().allow_normalization()
+    }
+
+    fn map_projection(
+        self,
+        tcx: TyCtxt<'tcx>,
+        f: impl FnOnce(Binder<'tcx, ProjectionClause<'tcx>>) -> Binder<'tcx, ProjectionClause<'tcx>>,
+    ) -> Option<Self> {
+        self.as_projection_clause().map(|kind| f(kind).upcast(tcx))
+    }
+
+    fn clause_kind_unchecked(&self) -> Option<ty::Binder<'tcx, ty::ClauseKind<'tcx>>> {
+        Some(self.kind())
     }
 }
 

compiler/rustc_middle/src/ty/util.rs

@@ -15,6 +15,7 @@ use rustc_index::bit_set::GrowableBitSet;
 use rustc_macros::{StableHash, TyDecodable, TyEncodable, extension};
 use rustc_span::sym;
 use rustc_structures::Limit;
+use rustc_type_ir::PredicateProxy;
 use rustc_type_ir::solve::SizedTraitKind;
 use smallvec::{SmallVec, smallvec};
 use tracing::{debug, instrument};
@@ -27,7 +28,7 @@ use crate::traits::ObligationCause;
 use crate::ty::layout::{FloatExt, IntegerExt};
 use crate::ty::{
     self, Asyncness, FallibleTypeFolder, GenericArgKind, GenericArgsRef, Ty, TyCtxt, TypeFoldable,
-    TypeFolder, TypeSuperFoldable, TypeVisitableExt, Unnormalized, Upcast,
+    TypeFolder, TypeSuperFoldable, TypeVisitableExt, Unnormalized,
 };
 
 #[derive(Copy, Clone, Debug)]
@@ -1037,24 +1038,23 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for OpaqueTypeExpander<'tcx> {
         }
     }
 
-    fn fold_predicate(&mut self, p: ty::Predicate<'tcx>) -> ty::Predicate<'tcx> {
-        if let ty::PredicateKind::Clause(clause) = p.kind().skip_binder()
-            && let ty::ClauseKind::Projection(projection_pred) = clause
-        {
-            p.kind()
-                .rebind(ty::ProjectionClause {
-                    projection_term: projection_pred.projection_term.fold_with(self),
-                    // Don't fold the term on the RHS of the projection predicate.
-                    // This is because for default trait methods with RPITITs, we
-                    // install a `NormalizesTo(Projection(RPITIT) -> Opaque(RPITIT))`
-                    // predicate, which would trivially cause a cycle when we do
-                    // anything that requires `TypingEnv::with_post_analysis_normalized`.
-                    term: projection_pred.term,
-                })
-                .upcast(self.tcx)
-        } else {
-            p.super_fold_with(self)
-        }
+    fn fold_predicate<P: PredicateProxy<TyCtxt<'tcx>>>(&mut self, p: P) -> P {
+        // We use `map_projection` to execute the closure only if `p` is a projection clause,
+        // to implement the logic described below (i.e. avoid folding the `term`).
+        // In all other cases, fold recursively, as normal.
+        p.map_projection(self.tcx, |bound_clause| {
+            let projection_clause = bound_clause.skip_binder();
+            bound_clause.rebind(ty::ProjectionClause {
+                projection_term: projection_clause.projection_term.fold_with(self),
+                // Don't fold the term on the RHS of the projection predicate.
+                // This is because for default trait methods with RPITITs, we
+                // install a `NormalizesTo(Projection(RPITIT) -> Opaque(RPITIT))`
+                // predicate, which would trivially cause a cycle when we do
+                // anything that requires `TypingEnv::with_post_analysis_normalized`.
+                term: projection_clause.term,
+            })
+        })
+        .unwrap_or_else(|| p.super_fold_with(self))
     }
 }
 

compiler/rustc_next_trait_solver/src/canonical/canonicalizer.rs

@@ -5,8 +5,8 @@ use rustc_type_ir::inherent::*;
 use rustc_type_ir::solve::{Goal, QueryInput};
 use rustc_type_ir::{
     self as ty, Canonical, CanonicalParamEnvCacheEntry, CanonicalVarKind, CanonicalizerState,
-    Flags, InferCtxtLike, Interner, PlaceholderConst, PlaceholderType, Region, TypeFlags,
-    TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitableExt,
+    Flags, InferCtxtLike, Interner, PlaceholderConst, PlaceholderType, PredicateProxy, Region,
+    TypeFlags, TypeFoldable, TypeFolder, TypeSuperFoldable, TypeVisitableExt,
 };
 use thin_vec::ThinVec;
 
@@ -583,7 +583,7 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> TypeFolder<I> for Canonicaliz
         Const::new_canonical_bound(self.cx(), var)
     }
 
-    fn fold_predicate(&mut self, p: I::Predicate) -> I::Predicate {
+    fn fold_predicate<P: PredicateProxy<I>>(&mut self, p: P) -> P {
         if !p.flags().intersects(NEEDS_CANONICAL) { p } else { p.super_fold_with(self) }
     }