fast-reject: add cache

lcnr · lcnr · commit 39841a4b7c88 · 2024-11-29T10:41:11.000+01:00
diff --git a/compiler/rustc_trait_selection/src/traits/coherence.rs b/compiler/rustc_trait_selection/src/traits/coherence.rs
@@ -98,7 +98,7 @@ pub fn overlapping_impls(
     // Before doing expensive operations like entering an inference context, do
     // a quick check via fast_reject to tell if the impl headers could possibly
     // unify.
-    let drcx = DeepRejectCtxt::relate_infer_infer(tcx);
+    let mut drcx = DeepRejectCtxt::relate_infer_infer(tcx);
     let impl1_ref = tcx.impl_trait_ref(impl1_def_id);
     let impl2_ref = tcx.impl_trait_ref(impl2_def_id);
     let may_overlap = match (impl1_ref, impl2_ref) {
diff --git a/compiler/rustc_trait_selection/src/traits/effects.rs b/compiler/rustc_trait_selection/src/traits/effects.rs
@@ -79,7 +79,7 @@ fn evaluate_host_effect_from_bounds<'tcx>(
     obligation: &HostEffectObligation<'tcx>,
 ) -> Result<ThinVec<PredicateObligation<'tcx>>, EvaluationFailure> {
     let infcx = selcx.infcx;
-    let drcx = DeepRejectCtxt::relate_rigid_rigid(selcx.tcx());
+    let mut drcx = DeepRejectCtxt::relate_rigid_rigid(selcx.tcx());
     let mut candidate = None;
 
     for predicate in obligation.param_env.caller_bounds() {
diff --git a/compiler/rustc_trait_selection/src/traits/project.rs b/compiler/rustc_trait_selection/src/traits/project.rs
@@ -866,7 +866,7 @@ fn assemble_candidates_from_predicates<'cx, 'tcx>(
     potentially_unnormalized_candidates: bool,
 ) {
     let infcx = selcx.infcx;
-    let drcx = DeepRejectCtxt::relate_rigid_rigid(selcx.tcx());
+    let mut drcx = DeepRejectCtxt::relate_rigid_rigid(selcx.tcx());
     for predicate in env_predicates {
         let bound_predicate = predicate.kind();
         if let ty::ClauseKind::Projection(data) = predicate.kind().skip_binder() {
diff --git a/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs b/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs
@@ -232,7 +232,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
             .filter(|p| p.def_id() == stack.obligation.predicate.def_id())
             .filter(|p| p.polarity() == stack.obligation.predicate.polarity());
 
-        let drcx = DeepRejectCtxt::relate_rigid_rigid(self.tcx());
+        let mut drcx = DeepRejectCtxt::relate_rigid_rigid(self.tcx());
         let obligation_args = stack.obligation.predicate.skip_binder().trait_ref.args;
         // Keep only those bounds which may apply, and propagate overflow if it occurs.
         for bound in bounds {
@@ -548,7 +548,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
         obligation: &PolyTraitObligation<'tcx>,
         candidates: &mut SelectionCandidateSet<'tcx>,
     ) {
-        let drcx = DeepRejectCtxt::relate_rigid_infer(self.tcx());
+        let mut drcx = DeepRejectCtxt::relate_rigid_infer(self.tcx());
         let obligation_args = obligation.predicate.skip_binder().trait_ref.args;
         self.tcx().for_each_relevant_impl(
             obligation.predicate.def_id(),
diff --git a/compiler/rustc_type_ir/src/fast_reject.rs b/compiler/rustc_type_ir/src/fast_reject.rs
@@ -11,6 +11,7 @@ use rustc_data_structures::stable_hasher::{HashStable, StableHasher, ToStableHas
 #[cfg(feature = "nightly")]
 use rustc_macros::{HashStable_NoContext, TyDecodable, TyEncodable};
 
+use crate::data_structures::DelayedSet;
 use crate::inherent::*;
 use crate::visit::TypeVisitableExt as _;
 use crate::{self as ty, Interner};
@@ -181,41 +182,42 @@ impl<DefId> SimplifiedType<DefId> {
 /// We also use this function during coherence. For coherence the
 /// impls only have to overlap for some value, so we treat parameters
 /// on both sides like inference variables.
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug)]
 pub struct DeepRejectCtxt<
     I: Interner,
     const INSTANTIATE_LHS_WITH_INFER: bool,
     const INSTANTIATE_RHS_WITH_INFER: bool,
 > {
     _interner: PhantomData<I>,
+    cache: DelayedSet<(I::Ty, I::Ty)>,
 }
 
 impl<I: Interner> DeepRejectCtxt<I, false, false> {
     /// Treat parameters in both the lhs and the rhs as rigid.
     pub fn relate_rigid_rigid(_interner: I) -> DeepRejectCtxt<I, false, false> {
-        DeepRejectCtxt { _interner: PhantomData }
+        DeepRejectCtxt { _interner: PhantomData, cache: Default::default() }
     }
 }
 
 impl<I: Interner> DeepRejectCtxt<I, true, true> {
     /// Treat parameters in both the lhs and the rhs as infer vars.
     pub fn relate_infer_infer(_interner: I) -> DeepRejectCtxt<I, true, true> {
-        DeepRejectCtxt { _interner: PhantomData }
+        DeepRejectCtxt { _interner: PhantomData, cache: Default::default() }
     }
 }
 
 impl<I: Interner> DeepRejectCtxt<I, false, true> {
     /// Treat parameters in the lhs as rigid, and in rhs as infer vars.
     pub fn relate_rigid_infer(_interner: I) -> DeepRejectCtxt<I, false, true> {
-        DeepRejectCtxt { _interner: PhantomData }
+        DeepRejectCtxt { _interner: PhantomData, cache: Default::default() }
     }
 }
 
 impl<I: Interner, const INSTANTIATE_LHS_WITH_INFER: bool, const INSTANTIATE_RHS_WITH_INFER: bool>
     DeepRejectCtxt<I, INSTANTIATE_LHS_WITH_INFER, INSTANTIATE_RHS_WITH_INFER>
 {
     pub fn args_may_unify(
-        self,
+        &mut self,
         obligation_args: I::GenericArgs,
         impl_args: I::GenericArgs,
     ) -> bool {
@@ -234,7 +236,24 @@ impl<I: Interner, const INSTANTIATE_LHS_WITH_INFER: bool, const INSTANTIATE_RHS_
         })
     }
 
-    pub fn types_may_unify(self, lhs: I::Ty, rhs: I::Ty) -> bool {
+    /// We only cache types if they may be part of exponential blowup, i.e. recursing into
+    /// them may relate more than one types. Constants and regions bottom out, so we don't
+    /// need to worry about them.
+    ///
+    /// We use a cache here as exponentially large - but self-similar - types otherwise
+    /// cause hangs, e.g. when compiling itertools with the `-Znext-solver`.
+    fn relate_cached(&mut self, lhs: I::Ty, rhs: I::Ty, f: impl FnOnce(&mut Self) -> bool) -> bool {
+        if self.cache.contains(&(lhs, rhs)) {
+            true
+        } else if f(self) {
+            self.cache.insert((lhs, rhs));
+            true
+        } else {
+            false
+        }
+    }
+
+    pub fn types_may_unify(&mut self, lhs: I::Ty, rhs: I::Ty) -> bool {
         match rhs.kind() {
             // Start by checking whether the `rhs` type may unify with
             // pretty much everything. Just return `true` in that case.
@@ -283,8 +302,8 @@ impl<I: Interner, const INSTANTIATE_LHS_WITH_INFER: bool, const INSTANTIATE_RHS_
             },
 
             ty::Adt(lhs_def, lhs_args) => match rhs.kind() {
-                ty::Adt(rhs_def, rhs_args) => {
-                    lhs_def == rhs_def && self.args_may_unify(lhs_args, rhs_args)
+                ty::Adt(rhs_def, rhs_args) if lhs_def == rhs_def => {
+                    self.relate_cached(lhs, rhs, |this| this.args_may_unify(lhs_args, rhs_args))
                 }
                 _ => false,
             },
@@ -322,12 +341,10 @@ impl<I: Interner, const INSTANTIATE_LHS_WITH_INFER: bool, const INSTANTIATE_RHS_
             | ty::Never
             | ty::Foreign(_) => lhs == rhs,
 
-            ty::Tuple(lhs) => match rhs.kind() {
-                ty::Tuple(rhs) => {
-                    lhs.len() == rhs.len()
-                        && iter::zip(lhs.iter(), rhs.iter())
-                            .all(|(lhs, rhs)| self.types_may_unify(lhs, rhs))
-                }
+            ty::Tuple(l) => match rhs.kind() {
+                ty::Tuple(r) if l.len() == r.len() => self.relate_cached(lhs, rhs, |this| {
+                    iter::zip(l.iter(), r.iter()).all(|(lhs, rhs)| this.types_may_unify(lhs, rhs))
+                }),
                 _ => false,
             },
 
@@ -363,47 +380,51 @@ impl<I: Interner, const INSTANTIATE_LHS_WITH_INFER: bool, const INSTANTIATE_RHS_
                     let lhs_sig_tys = lhs_sig_tys.skip_binder().inputs_and_output;
                     let rhs_sig_tys = rhs_sig_tys.skip_binder().inputs_and_output;
 
-                    lhs_hdr == rhs_hdr
-                        && lhs_sig_tys.len() == rhs_sig_tys.len()
-                        && iter::zip(lhs_sig_tys.iter(), rhs_sig_tys.iter())
-                            .all(|(lhs, rhs)| self.types_may_unify(lhs, rhs))
+                    if lhs_hdr == rhs_hdr && lhs_sig_tys.len() == rhs_sig_tys.len() {
+                        self.relate_cached(lhs, rhs, |this| {
+                            iter::zip(lhs_sig_tys.iter(), rhs_sig_tys.iter())
+                                .all(|(lhs, rhs)| this.types_may_unify(lhs, rhs))
+                        })
+                    } else {
+                        false
+                    }
                 }
                 _ => false,
             },
 
             ty::Bound(..) => true,
 
             ty::FnDef(lhs_def_id, lhs_args) => match rhs.kind() {
-                ty::FnDef(rhs_def_id, rhs_args) => {
-                    lhs_def_id == rhs_def_id && self.args_may_unify(lhs_args, rhs_args)
+                ty::FnDef(rhs_def_id, rhs_args) if lhs_def_id == rhs_def_id => {
+                    self.relate_cached(lhs, rhs, |this| this.args_may_unify(lhs_args, rhs_args))
                 }
                 _ => false,
             },
 
             ty::Closure(lhs_def_id, lhs_args) => match rhs.kind() {
-                ty::Closure(rhs_def_id, rhs_args) => {
-                    lhs_def_id == rhs_def_id && self.args_may_unify(lhs_args, rhs_args)
+                ty::Closure(rhs_def_id, rhs_args) if lhs_def_id == rhs_def_id => {
+                    self.relate_cached(lhs, rhs, |this| this.args_may_unify(lhs_args, rhs_args))
                 }
                 _ => false,
             },
 
             ty::CoroutineClosure(lhs_def_id, lhs_args) => match rhs.kind() {
-                ty::CoroutineClosure(rhs_def_id, rhs_args) => {
-                    lhs_def_id == rhs_def_id && self.args_may_unify(lhs_args, rhs_args)
+                ty::CoroutineClosure(rhs_def_id, rhs_args) if lhs_def_id == rhs_def_id => {
+                    self.relate_cached(lhs, rhs, |this| this.args_may_unify(lhs_args, rhs_args))
                 }
                 _ => false,
             },
 
             ty::Coroutine(lhs_def_id, lhs_args) => match rhs.kind() {
-                ty::Coroutine(rhs_def_id, rhs_args) => {
-                    lhs_def_id == rhs_def_id && self.args_may_unify(lhs_args, rhs_args)
+                ty::Coroutine(rhs_def_id, rhs_args) if lhs_def_id == rhs_def_id => {
+                    self.relate_cached(lhs, rhs, |this| this.args_may_unify(lhs_args, rhs_args))
                 }
                 _ => false,
             },
 
             ty::CoroutineWitness(lhs_def_id, lhs_args) => match rhs.kind() {
-                ty::CoroutineWitness(rhs_def_id, rhs_args) => {
-                    lhs_def_id == rhs_def_id && self.args_may_unify(lhs_args, rhs_args)
+                ty::CoroutineWitness(rhs_def_id, rhs_args) if lhs_def_id == rhs_def_id => {
+                    self.relate_cached(lhs, rhs, |this| this.args_may_unify(lhs_args, rhs_args))
                 }
                 _ => false,
             },
@@ -417,7 +438,7 @@ impl<I: Interner, const INSTANTIATE_LHS_WITH_INFER: bool, const INSTANTIATE_RHS_
         }
     }
 
-    pub fn consts_may_unify(self, lhs: I::Const, rhs: I::Const) -> bool {
+    pub fn consts_may_unify(&mut self, lhs: I::Const, rhs: I::Const) -> bool {
         match rhs.kind() {
             ty::ConstKind::Param(_) => {
                 if INSTANTIATE_RHS_WITH_INFER {
@@ -465,7 +486,7 @@ impl<I: Interner, const INSTANTIATE_LHS_WITH_INFER: bool, const INSTANTIATE_RHS_
         }
     }
 
-    fn var_and_ty_may_unify(self, var: ty::InferTy, ty: I::Ty) -> bool {
+    fn var_and_ty_may_unify(&mut self, var: ty::InferTy, ty: I::Ty) -> bool {
         if !ty.is_known_rigid() {
             return true;
         }
diff --git a/src/librustdoc/html/render/write_shared.rs b/src/librustdoc/html/render/write_shared.rs
@@ -918,8 +918,8 @@ impl<'item> DocVisitor<'item> for TypeImplCollector<'_, '_, 'item> {
             // Be aware of `tests/rustdoc/type-alias/deeply-nested-112515.rs` which might regress.
             let Some(impl_did) = impl_item_id.as_def_id() else { continue };
             let for_ty = self.cx.tcx().type_of(impl_did).skip_binder();
-            let reject_cx = DeepRejectCtxt::relate_infer_infer(self.cx.tcx());
-            if !reject_cx.types_may_unify(aliased_ty, for_ty) {
+            let mut drcx = DeepRejectCtxt::relate_infer_infer(self.cx.tcx());
+            if !drcx.types_may_unify(aliased_ty, for_ty) {
                 continue;
             }
             // Avoid duplicates
