Auto merge of #97345 - lcnr:fast_reject, r=nnethercote

add a deep fast_reject routine

continues the work on #97136.

r? `@nnethercote`

Actually agree with you on the match structure 😆 let's see how that impacted perf 😅

Author: bors

compiler/rustc_middle/src/ty/fast_reject.rs

@@ -1,8 +1,10 @@
 use crate::mir::Mutability;
+use crate::ty::subst::GenericArgKind;
 use crate::ty::{self, Ty, TyCtxt, TypeFoldable};
 use rustc_hir::def_id::DefId;
 use std::fmt::Debug;
 use std::hash::Hash;
+use std::iter;
 
 use self::SimplifiedTypeGen::*;
 
@@ -72,6 +74,10 @@ pub enum TreatParams {
 
 /// Tries to simplify a type by only returning the outermost injective¹ layer, if one exists.
 ///
+/// **This function should only be used if you need to store or retrieve the type from some
+/// hashmap. If you want to quickly decide whether two types may unify, use the [DeepRejectCtxt]
+/// instead.**
+///
 /// The idea is to get something simple that we can use to quickly decide if two types could unify,
 /// for example during method lookup. If this function returns `Some(x)` it can only unify with
 /// types for which this method returns either `Some(x)` as well or `None`.
@@ -182,3 +188,218 @@ impl<D: Copy + Debug + Eq> SimplifiedTypeGen<D> {
         }
     }
 }
+
+/// Given generic arguments from an obligation and an impl,
+/// could these two be unified after replacing parameters in the
+/// the impl with inference variables.
+///
+/// For obligations, parameters won't be replaced by inference
+/// variables and only unify with themselves. We treat them
+/// the same way we treat placeholders.
+///
+/// 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. This behavior is toggled
+/// using the `treat_obligation_params` field.
+#[derive(Debug, Clone, Copy)]
+pub struct DeepRejectCtxt {
+    pub treat_obligation_params: TreatParams,
+}
+
+impl DeepRejectCtxt {
+    pub fn generic_args_may_unify(
+        self,
+        obligation_arg: ty::GenericArg<'_>,
+        impl_arg: ty::GenericArg<'_>,
+    ) -> bool {
+        match (obligation_arg.unpack(), impl_arg.unpack()) {
+            // We don't fast reject based on regions for now.
+            (GenericArgKind::Lifetime(_), GenericArgKind::Lifetime(_)) => true,
+            (GenericArgKind::Type(obl), GenericArgKind::Type(imp)) => {
+                self.types_may_unify(obl, imp)
+            }
+            (GenericArgKind::Const(obl), GenericArgKind::Const(imp)) => {
+                self.consts_may_unify(obl, imp)
+            }
+            _ => bug!("kind mismatch: {obligation_arg} {impl_arg}"),
+        }
+    }
+
+    pub fn types_may_unify(self, obligation_ty: Ty<'_>, impl_ty: Ty<'_>) -> bool {
+        match impl_ty.kind() {
+            // Start by checking whether the type in the impl may unify with
+            // pretty much everything. Just return `true` in that case.
+            ty::Param(_) | ty::Projection(_) | ty::Error(_) => return true,
+            // These types only unify with inference variables or their own
+            // variant.
+            ty::Bool
+            | ty::Char
+            | ty::Int(_)
+            | ty::Uint(_)
+            | ty::Float(_)
+            | ty::Adt(..)
+            | ty::Str
+            | ty::Array(..)
+            | ty::Slice(..)
+            | ty::RawPtr(..)
+            | ty::Dynamic(..)
+            | ty::Ref(..)
+            | ty::Never
+            | ty::Tuple(..)
+            | ty::FnPtr(..)
+            | ty::Foreign(..)
+            | ty::Opaque(..) => {}
+            ty::FnDef(..)
+            | ty::Closure(..)
+            | ty::Generator(..)
+            | ty::GeneratorWitness(..)
+            | ty::Placeholder(..)
+            | ty::Bound(..)
+            | ty::Infer(_) => bug!("unexpected impl_ty: {impl_ty}"),
+        }
+
+        let k = impl_ty.kind();
+        match *obligation_ty.kind() {
+            // Purely rigid types, use structural equivalence.
+            ty::Bool
+            | ty::Char
+            | ty::Int(_)
+            | ty::Uint(_)
+            | ty::Float(_)
+            | ty::Str
+            | ty::Never
+            | ty::Foreign(_) => obligation_ty == impl_ty,
+            ty::Ref(_, obl_ty, obl_mutbl) => match k {
+                &ty::Ref(_, impl_ty, impl_mutbl) => {
+                    obl_mutbl == impl_mutbl && self.types_may_unify(obl_ty, impl_ty)
+                }
+                _ => false,
+            },
+            ty::Adt(obl_def, obl_substs) => match k {
+                &ty::Adt(impl_def, impl_substs) => {
+                    obl_def == impl_def
+                        && iter::zip(obl_substs, impl_substs)
+                            .all(|(obl, imp)| self.generic_args_may_unify(obl, imp))
+                }
+                _ => false,
+            },
+            ty::Slice(obl_ty) => {
+                matches!(k, &ty::Slice(impl_ty) if self.types_may_unify(obl_ty, impl_ty))
+            }
+            ty::Array(obl_ty, obl_len) => match k {
+                &ty::Array(impl_ty, impl_len) => {
+                    self.types_may_unify(obl_ty, impl_ty)
+                        && self.consts_may_unify(obl_len, impl_len)
+                }
+                _ => false,
+            },
+            ty::Tuple(obl) => match k {
+                &ty::Tuple(imp) => {
+                    obl.len() == imp.len()
+                        && iter::zip(obl, imp).all(|(obl, imp)| self.types_may_unify(obl, imp))
+                }
+                _ => false,
+            },
+            ty::RawPtr(obl) => match k {
+                ty::RawPtr(imp) => obl.mutbl == imp.mutbl && self.types_may_unify(obl.ty, imp.ty),
+                _ => false,
+            },
+            ty::Dynamic(obl_preds, ..) => {
+                // Ideally we would walk the existential predicates here or at least
+                // compare their length. But considering that the relevant `Relate` impl
+                // actually sorts and deduplicates these, that doesn't work.
+                matches!(k, ty::Dynamic(impl_preds, ..) if
+                    obl_preds.principal_def_id() == impl_preds.principal_def_id()
+                )
+            }
+            ty::FnPtr(obl_sig) => match k {
+                ty::FnPtr(impl_sig) => {
+                    let ty::FnSig { inputs_and_output, c_variadic, unsafety, abi } =
+                        obl_sig.skip_binder();
+                    let impl_sig = impl_sig.skip_binder();
+
+                    abi == impl_sig.abi
+                        && c_variadic == impl_sig.c_variadic
+                        && unsafety == impl_sig.unsafety
+                        && inputs_and_output.len() == impl_sig.inputs_and_output.len()
+                        && iter::zip(inputs_and_output, impl_sig.inputs_and_output)
+                            .all(|(obl, imp)| self.types_may_unify(obl, imp))
+                }
+                _ => false,
+            },
+
+            // Opaque types in impls should be forbidden, but that doesn't
+            // stop compilation. So this match arm should never return true
+            // if compilation succeeds.
+            ty::Opaque(..) => matches!(k, ty::Opaque(..)),
+
+            // Impls cannot contain these types as these cannot be named directly.
+            ty::FnDef(..) | ty::Closure(..) | ty::Generator(..) => false,
+
+            ty::Placeholder(..) => false,
+
+            // Depending on the value of `treat_obligation_params`, we either
+            // treat generic parameters like placeholders or like inference variables.
+            ty::Param(_) => match self.treat_obligation_params {
+                TreatParams::AsPlaceholder => false,
+                TreatParams::AsInfer => true,
+            },
+
+            ty::Infer(_) => true,
+
+            // As we're walking the whole type, it may encounter projections
+            // inside of binders and what not, so we're just going to assume that
+            // projections can unify with other stuff.
+            //
+            // Looking forward to lazy normalization this is the safer strategy anyways.
+            ty::Projection(_) => true,
+
+            ty::Error(_) => true,
+
+            ty::GeneratorWitness(..) | ty::Bound(..) => {
+                bug!("unexpected obligation type: {:?}", obligation_ty)
+            }
+        }
+    }
+
+    pub fn consts_may_unify(self, obligation_ct: ty::Const<'_>, impl_ct: ty::Const<'_>) -> bool {
+        match impl_ct.val() {
+            ty::ConstKind::Param(_) | ty::ConstKind::Unevaluated(_) | ty::ConstKind::Error(_) => {
+                return true;
+            }
+            ty::ConstKind::Value(_) => {}
+            ty::ConstKind::Infer(_) | ty::ConstKind::Bound(..) | ty::ConstKind::Placeholder(_) => {
+                bug!("unexpected impl arg: {:?}", impl_ct)
+            }
+        }
+
+        let k = impl_ct.val();
+        match obligation_ct.val() {
+            ty::ConstKind::Param(_) => match self.treat_obligation_params {
+                TreatParams::AsPlaceholder => false,
+                TreatParams::AsInfer => true,
+            },
+
+            // As we don't necessarily eagerly evaluate constants,
+            // they might unify with any value.
+            ty::ConstKind::Unevaluated(_) | ty::ConstKind::Error(_) => true,
+            ty::ConstKind::Value(obl) => match k {
+                ty::ConstKind::Value(imp) => {
+                    // FIXME(valtrees): Once we have valtrees, we can just
+                    // compare them directly here.
+                    match (obl.try_to_scalar_int(), imp.try_to_scalar_int()) {
+                        (Some(obl), Some(imp)) => obl == imp,
+                        _ => true,
+                    }
+                }
+                _ => true,
+            },
+
+            ty::ConstKind::Infer(_) => true,
+
+            ty::ConstKind::Bound(..) | ty::ConstKind::Placeholder(_) => {
+                bug!("unexpected obl const: {:?}", obligation_ct)
+            }
+        }
+    }
+}

compiler/rustc_trait_selection/src/traits/coherence.rs

@@ -20,7 +20,7 @@ use rustc_hir::CRATE_HIR_ID;
 use rustc_infer::infer::{InferCtxt, TyCtxtInferExt};
 use rustc_infer::traits::{util, TraitEngine};
 use rustc_middle::traits::specialization_graph::OverlapMode;
-use rustc_middle::ty::fast_reject::{self, TreatParams};
+use rustc_middle::ty::fast_reject::{DeepRejectCtxt, TreatParams};
 use rustc_middle::ty::fold::TypeFoldable;
 use rustc_middle::ty::subst::Subst;
 use rustc_middle::ty::{self, ImplSubject, Ty, TyCtxt};
@@ -79,26 +79,21 @@ where
     // 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 { treat_obligation_params: TreatParams::AsInfer };
     let impl1_ref = tcx.impl_trait_ref(impl1_def_id);
     let impl2_ref = tcx.impl_trait_ref(impl2_def_id);
-
-    // Check if any of the input types definitely do not unify.
-    if iter::zip(
-        impl1_ref.iter().flat_map(|tref| tref.substs.types()),
-        impl2_ref.iter().flat_map(|tref| tref.substs.types()),
-    )
-    .any(|(ty1, ty2)| {
-        let t1 = fast_reject::simplify_type(tcx, ty1, TreatParams::AsInfer);
-        let t2 = fast_reject::simplify_type(tcx, ty2, TreatParams::AsInfer);
-
-        if let (Some(t1), Some(t2)) = (t1, t2) {
-            // Simplified successfully
-            t1 != t2
-        } else {
-            // Types might unify
-            false
+    let may_overlap = match (impl1_ref, impl2_ref) {
+        (Some(a), Some(b)) => iter::zip(a.substs, b.substs)
+            .all(|(arg1, arg2)| drcx.generic_args_may_unify(arg1, arg2)),
+        (None, None) => {
+            let self_ty1 = tcx.type_of(impl1_def_id);
+            let self_ty2 = tcx.type_of(impl2_def_id);
+            drcx.types_may_unify(self_ty1, self_ty2)
         }
-    }) {
+        _ => bug!("unexpected impls: {impl1_def_id:?} {impl2_def_id:?}"),
+    };
+
+    if !may_overlap {
         // Some types involved are definitely different, so the impls couldn't possibly overlap.
         debug!("overlapping_impls: fast_reject early-exit");
         return no_overlap();
@@ -519,7 +514,7 @@ pub fn orphan_check(tcx: TyCtxt<'_>, impl_def_id: DefId) -> Result<(), OrphanChe
 /// 3. Before this local type, no generic type parameter of the impl must
 ///    be reachable through fundamental types.
 ///     - e.g. `impl<T> Trait<LocalType> for Vec<T>` is fine, as `Vec` is not fundamental.
-///     - while `impl<T> Trait<LocalType for Box<T>` results in an error, as `T` is
+///     - while `impl<T> Trait<LocalType> for Box<T>` results in an error, as `T` is
 ///       reachable through the fundamental type `Box`.
 /// 4. Every type in the local key parameter not known in C, going
 ///    through the parameter's type tree, must appear only as a subtree of

compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs

@@ -539,8 +539,16 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
             obligation.predicate.def_id(),
             obligation.predicate.skip_binder().trait_ref.self_ty(),
             |impl_def_id| {
+                // Before we create the substitutions and everything, first
+                // consider a "quick reject". This avoids creating more types
+                // and so forth that we need to.
+                let impl_trait_ref = self.tcx().bound_impl_trait_ref(impl_def_id).unwrap();
+                if self.fast_reject_trait_refs(obligation, &impl_trait_ref.0) {
+                    return;
+                }
+
                 self.infcx.probe(|_| {
-                    if let Ok(_substs) = self.match_impl(impl_def_id, obligation) {
+                    if let Ok(_substs) = self.match_impl(impl_def_id, impl_trait_ref, obligation) {
                         candidates.vec.push(ImplCandidate(impl_def_id));
                     }
                 });