Fix PartialEq args when #[const_trait] is enabled

5 files changed, 56 insertions, 37 deletions

diff --git a/compiler/rustc_hir_typeck/src/callee.rs b/compiler/rustc_hir_typeck/src/callee.rs
index 365170a6bbf..0c6c6926c4e 100644
--- a/compiler/rustc_hir_typeck/src/callee.rs
+++ b/compiler/rustc_hir_typeck/src/callee.rs
@@ -6,7 +6,7 @@ use crate::errors;
 use rustc_ast::util::parser::PREC_POSTFIX;
 use rustc_errors::{Applicability, Diagnostic, ErrorGuaranteed, StashKey};
 use rustc_hir as hir;
-use rustc_hir::def::{self, CtorKind, DefKind, Namespace, Res};
+use rustc_hir::def::{self, CtorKind, Namespace, Res};
 use rustc_hir::def_id::DefId;
 use rustc_hir::HirId;
 use rustc_hir_analysis::autoderef::Autoderef;
@@ -781,38 +781,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         let generics = tcx.generics_of(callee_did);
         let Some(host_effect_index) = generics.host_effect_index else { return };
 
-        // if the callee does have the param, we need to equate the param to some const
-        // value no matter whether the effects feature is enabled in the local crate,
-        // because inference will fail if we don't.
-        let mut host_always_on =
-            !tcx.features().effects || tcx.sess.opts.unstable_opts.unleash_the_miri_inside_of_you;
-
-        // Compute the constness required by the context.
-        let context = tcx.hir().enclosing_body_owner(call_expr_hir);
-        let const_context = tcx.hir().body_const_context(context);
-
-        let kind = tcx.def_kind(context.to_def_id());
-        debug_assert_ne!(kind, DefKind::ConstParam);
-
-        if tcx.has_attr(context.to_def_id(), sym::rustc_do_not_const_check) {
-            trace!("do not const check this context");
-            host_always_on = true;
-        }
-
-        let effect = match const_context {
-            _ if host_always_on => tcx.consts.true_,
-            Some(hir::ConstContext::Static(_) | hir::ConstContext::Const { .. }) => {
-                tcx.consts.false_
-            }
-            Some(hir::ConstContext::ConstFn) => {
-                let host_idx = tcx
-                    .generics_of(context)
-                    .host_effect_index
-                    .expect("ConstContext::Maybe must have host effect param");
-                ty::GenericArgs::identity_for_item(tcx, context).const_at(host_idx)
-            }
-            None => tcx.consts.true_,
-        };
+        let effect = tcx.expected_const_effect_param_for_body(self.body_id);
 
         trace!(?effect, ?generics, ?callee_args);
 
diff --git a/compiler/rustc_middle/src/ty/util.rs b/compiler/rustc_middle/src/ty/util.rs
index cbf1a9900d9..afa6eac4e2f 100644
--- a/compiler/rustc_middle/src/ty/util.rs
+++ b/compiler/rustc_middle/src/ty/util.rs
@@ -779,6 +779,40 @@ impl<'tcx> TyCtxt<'tcx> {
             // the language.
             || self.extern_crate(key.as_def_id()).is_some_and(|e| e.is_direct())
     }
+
+    pub fn expected_const_effect_param_for_body(self, def_id: LocalDefId) -> ty::Const<'tcx> {
+        // if the callee does have the param, we need to equate the param to some const
+        // value no matter whether the effects feature is enabled in the local crate,
+        // because inference will fail if we don't.
+        let mut host_always_on =
+            !self.features().effects || self.sess.opts.unstable_opts.unleash_the_miri_inside_of_you;
+
+        // Compute the constness required by the context.
+        let const_context = self.hir().body_const_context(def_id);
+
+        let kind = self.def_kind(def_id);
+        debug_assert_ne!(kind, DefKind::ConstParam);
+
+        if self.has_attr(def_id, sym::rustc_do_not_const_check) {
+            trace!("do not const check this context");
+            host_always_on = true;
+        }
+
+        match const_context {
+            _ if host_always_on => self.consts.true_,
+            Some(hir::ConstContext::Static(_) | hir::ConstContext::Const { .. }) => {
+                self.consts.false_
+            }
+            Some(hir::ConstContext::ConstFn) => {
+                let host_idx = self
+                    .generics_of(def_id)
+                    .host_effect_index
+                    .expect("ConstContext::Maybe must have host effect param");
+                ty::GenericArgs::identity_for_item(self, def_id).const_at(host_idx)
+            }
+            None => self.consts.true_,
+        }
+    }
 }
 
 struct OpaqueTypeExpander<'tcx> {
diff --git a/compiler/rustc_mir_build/src/build/matches/test.rs b/compiler/rustc_mir_build/src/build/matches/test.rs
index 6bd60972c8b..0b7fad26f69 100644
--- a/compiler/rustc_mir_build/src/build/matches/test.rs
+++ b/compiler/rustc_mir_build/src/build/matches/test.rs
@@ -494,7 +494,14 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         }
 
         let eq_def_id = self.tcx.require_lang_item(LangItem::PartialEq, Some(source_info.span));
-        let method = trait_method(self.tcx, eq_def_id, sym::eq, [ty, ty]);
+
+        let mut args: Vec<ty::GenericArg<'tcx>> = vec![ty.into(), ty.into()];
+        // If `PartialEq` is `#[const_trait]`, then add a const effect param
+        if self.tcx.generics_of(eq_def_id).host_effect_index.is_some() {
+            args.push(self.tcx.expected_const_effect_param_for_body(self.def_id).into());
+        }
+
+        let method = trait_method(self.tcx, eq_def_id, sym::eq, args);
 
         let bool_ty = self.tcx.types.bool;
         let eq_result = self.temp(bool_ty, source_info.span);
diff --git a/compiler/rustc_mir_build/src/thir/pattern/const_to_pat.rs b/compiler/rustc_mir_build/src/thir/pattern/const_to_pat.rs
index c4bf2f23899..4db9e71ed25 100644
--- a/compiler/rustc_mir_build/src/thir/pattern/const_to_pat.rs
+++ b/compiler/rustc_mir_build/src/thir/pattern/const_to_pat.rs
@@ -258,18 +258,27 @@ impl<'tcx> ConstToPat<'tcx> {
 
     #[instrument(level = "trace", skip(self), ret)]
     fn type_has_partial_eq_impl(&self, ty: Ty<'tcx>) -> bool {
+        let tcx = self.tcx();
         // double-check there even *is* a semantic `PartialEq` to dispatch to.
         //
         // (If there isn't, then we can safely issue a hard
         // error, because that's never worked, due to compiler
         // using `PartialEq::eq` in this scenario in the past.)
-        let tcx = self.tcx();
         let partial_eq_trait_id = tcx.require_lang_item(hir::LangItem::PartialEq, Some(self.span));
+        let mut args: Vec<ty::GenericArg<'tcx>> = vec![ty.into(), ty.into()];
+        // If `PartialEq` is `#[const_trait]`, then add a const effect param
+        if tcx.generics_of(partial_eq_trait_id).host_effect_index.is_some() {
+            args.push(
+                tcx.expected_const_effect_param_for_body(tcx.hir().enclosing_body_owner(self.id))
+                    .into(),
+            );
+        }
+
         let partial_eq_obligation = Obligation::new(
             tcx,
             ObligationCause::dummy(),
             self.param_env,
-            ty::TraitRef::new(tcx, partial_eq_trait_id, [ty, ty]),
+            ty::TraitRef::new(tcx, partial_eq_trait_id, args),
         );
 
         // This *could* accept a type that isn't actually `PartialEq`, because region bounds get
diff --git a/tests/ui/rfcs/rfc-2632-const-trait-impl/effects/minicore.rs b/tests/ui/rfcs/rfc-2632-const-trait-impl/effects/minicore.rs
index 67da5f9533c..d79cf372b87 100644
--- a/tests/ui/rfcs/rfc-2632-const-trait-impl/effects/minicore.rs
+++ b/tests/ui/rfcs/rfc-2632-const-trait-impl/effects/minicore.rs
@@ -336,7 +336,7 @@ fn from_str(s: &str) -> Result<bool, ()> {
 }
 
 #[lang = "eq"]
-// FIXME #[const_trait]
+#[const_trait]
 trait PartialEq<Rhs: ?Sized = Self> {
     fn eq(&self, other: &Rhs) -> bool;
     fn ne(&self, other: &Rhs) -> bool {