Introduce trait_item_of

8 files changed, 25 insertions, 28 deletions

diff --git a/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs b/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs
index ea264c8064a..6d69040c711 100644
--- a/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs
+++ b/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs
@@ -716,7 +716,7 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
             return (false, false, None);
         };
 
-        let implemented_trait_item = self.infcx.tcx.associated_item(my_def).trait_item_def_id;
+        let implemented_trait_item = self.infcx.tcx.trait_item_of(my_def);
 
         (
             true,
diff --git a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
index 008340e614d..dc500c363f4 100644
--- a/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
+++ b/compiler/rustc_codegen_ssa/src/codegen_attrs.rs
@@ -562,15 +562,6 @@ fn codegen_fn_attrs(tcx: TyCtxt<'_>, did: LocalDefId) -> CodegenFnAttrs {
     codegen_fn_attrs
 }
 
-/// If the provided DefId is a method in a trait impl, return the DefId of the method prototype.
-fn opt_trait_item(tcx: TyCtxt<'_>, def_id: DefId) -> Option<DefId> {
-    let impl_item = tcx.opt_associated_item(def_id)?;
-    match impl_item.container {
-        ty::AssocItemContainer::Impl => impl_item.trait_item_def_id,
-        _ => None,
-    }
-}
-
 fn disabled_sanitizers_for(tcx: TyCtxt<'_>, did: LocalDefId) -> SanitizerSet {
     // Backtrack to the crate root.
     let mut disabled = match tcx.opt_local_parent(did) {
@@ -600,14 +591,15 @@ fn disabled_sanitizers_for(tcx: TyCtxt<'_>, did: LocalDefId) -> SanitizerSet {
 /// Checks if the provided DefId is a method in a trait impl for a trait which has track_caller
 /// applied to the method prototype.
 fn should_inherit_track_caller(tcx: TyCtxt<'_>, def_id: DefId) -> bool {
-    let Some(trait_item) = opt_trait_item(tcx, def_id) else { return false };
-    tcx.codegen_fn_attrs(trait_item).flags.intersects(CodegenFnAttrFlags::TRACK_CALLER)
+    tcx.trait_item_of(def_id).is_some_and(|id| {
+        tcx.codegen_fn_attrs(id).flags.intersects(CodegenFnAttrFlags::TRACK_CALLER)
+    })
 }
 
 /// If the provided DefId is a method in a trait impl, return the value of the `#[align]`
 /// attribute on the method prototype (if any).
 fn inherited_align<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId) -> Option<Align> {
-    tcx.codegen_fn_attrs(opt_trait_item(tcx, def_id)?).alignment
+    tcx.codegen_fn_attrs(tcx.trait_item_of(def_id)?).alignment
 }
 
 /// We now check the #\[rustc_autodiff\] attributes which we generated from the #[autodiff(...)]
diff --git a/compiler/rustc_hir_analysis/src/collect/predicates_of.rs b/compiler/rustc_hir_analysis/src/collect/predicates_of.rs
index 126ffabd448..dd3590f9ac5 100644
--- a/compiler/rustc_hir_analysis/src/collect/predicates_of.rs
+++ b/compiler/rustc_hir_analysis/src/collect/predicates_of.rs
@@ -111,9 +111,8 @@ fn gather_explicit_predicates_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::Gen
         }
 
         Some(ImplTraitInTraitData::Impl { fn_def_id }) => {
-            let assoc_item = tcx.associated_item(def_id);
-            let trait_assoc_predicates =
-                tcx.explicit_predicates_of(assoc_item.trait_item_def_id.unwrap());
+            let trait_item_def_id = tcx.trait_item_of(def_id).unwrap();
+            let trait_assoc_predicates = tcx.explicit_predicates_of(trait_item_def_id);
 
             let impl_assoc_identity_args = ty::GenericArgs::identity_for_item(tcx, def_id);
             let impl_def_id = tcx.parent(fn_def_id);
diff --git a/compiler/rustc_hir_analysis/src/collect/type_of.rs b/compiler/rustc_hir_analysis/src/collect/type_of.rs
index 62125c99d80..a8af73717cd 100644
--- a/compiler/rustc_hir_analysis/src/collect/type_of.rs
+++ b/compiler/rustc_hir_analysis/src/collect/type_of.rs
@@ -125,8 +125,8 @@ pub(super) fn type_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> ty::EarlyBinder<'_
         Some(ty::ImplTraitInTraitData::Impl { fn_def_id }) => {
             match tcx.collect_return_position_impl_trait_in_trait_tys(fn_def_id) {
                 Ok(map) => {
-                    let assoc_item = tcx.associated_item(def_id);
-                    return map[&assoc_item.trait_item_def_id.unwrap()];
+                    let trait_item_def_id = tcx.trait_item_of(def_id).unwrap();
+                    return map[&trait_item_def_id];
                 }
                 Err(_) => {
                     return ty::EarlyBinder::bind(Ty::new_error_with_message(
diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
index da17ec1f9f3..5a4ebfe28a4 100644
--- a/compiler/rustc_middle/src/ty/mod.rs
+++ b/compiler/rustc_middle/src/ty/mod.rs
@@ -1934,6 +1934,15 @@ impl<'tcx> TyCtxt<'tcx> {
         Some((parent, def_kind))
     }
 
+    /// Returns the trait item that is implemented by the given item `DefId`.
+    pub fn trait_item_of(self, def_id: impl IntoQueryParam<DefId>) -> Option<DefId> {
+        let assoc = self.opt_associated_item(def_id.into_query_param())?;
+        if assoc.container != AssocItemContainer::Impl {
+            return None;
+        }
+        assoc.trait_item_def_id
+    }
+
     /// If the given `DefId` is an associated item of a trait,
     /// returns the `DefId` of the trait; otherwise, returns `None`.
     pub fn trait_of_assoc(self, def_id: DefId) -> Option<DefId> {
diff --git a/compiler/rustc_passes/src/dead.rs b/compiler/rustc_passes/src/dead.rs
index fc33405d455..f36a7a5b184 100644
--- a/compiler/rustc_passes/src/dead.rs
+++ b/compiler/rustc_passes/src/dead.rs
@@ -484,13 +484,11 @@ impl<'tcx> MarkSymbolVisitor<'tcx> {
     fn check_impl_or_impl_item_live(&mut self, local_def_id: LocalDefId) -> bool {
         let (impl_block_id, trait_def_id) = match self.tcx.def_kind(local_def_id) {
             // assoc impl items of traits are live if the corresponding trait items are live
-            DefKind::AssocConst | DefKind::AssocTy | DefKind::AssocFn => (
-                self.tcx.local_parent(local_def_id),
-                self.tcx
-                    .associated_item(local_def_id)
-                    .trait_item_def_id
-                    .and_then(|def_id| def_id.as_local()),
-            ),
+            DefKind::AssocConst | DefKind::AssocTy | DefKind::AssocFn => {
+                let trait_item_id =
+                    self.tcx.trait_item_of(local_def_id).and_then(|def_id| def_id.as_local());
+                (self.tcx.local_parent(local_def_id), trait_item_id)
+            }
             // impl items are live if the corresponding traits are live
             DefKind::Impl { of_trait: true } => (
                 local_def_id,
diff --git a/compiler/rustc_sanitizers/src/cfi/typeid/itanium_cxx_abi/transform.rs b/compiler/rustc_sanitizers/src/cfi/typeid/itanium_cxx_abi/transform.rs
index d81fa062e01..577a16a0d25 100644
--- a/compiler/rustc_sanitizers/src/cfi/typeid/itanium_cxx_abi/transform.rs
+++ b/compiler/rustc_sanitizers/src/cfi/typeid/itanium_cxx_abi/transform.rs
@@ -469,8 +469,7 @@ fn implemented_method<'tcx>(
     let ancestor = if let Some(impl_id) = tcx.impl_of_assoc(instance.def_id()) {
         // Implementation in an `impl` block
         trait_ref = tcx.impl_trait_ref(impl_id)?;
-        let impl_method = tcx.associated_item(instance.def_id());
-        method_id = impl_method.trait_item_def_id?;
+        method_id = tcx.trait_item_of(instance.def_id())?;
         trait_method = tcx.associated_item(method_id);
         trait_id = trait_ref.skip_binder().def_id;
         impl_id
diff --git a/compiler/rustc_ty_utils/src/implied_bounds.rs b/compiler/rustc_ty_utils/src/implied_bounds.rs
index cdfb93c4e7d..543f6a3ccf7 100644
--- a/compiler/rustc_ty_utils/src/implied_bounds.rs
+++ b/compiler/rustc_ty_utils/src/implied_bounds.rs
@@ -107,7 +107,7 @@ fn assumed_wf_types<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> &'tcx [(Ty<'
                 // the assumed wf types of the trait's RPITIT GAT.
                 ty::ImplTraitInTraitData::Impl { .. } => {
                     let impl_def_id = tcx.local_parent(def_id);
-                    let rpitit_def_id = tcx.associated_item(def_id).trait_item_def_id.unwrap();
+                    let rpitit_def_id = tcx.trait_item_of(def_id).unwrap();
                     let args = ty::GenericArgs::identity_for_item(tcx, def_id).rebase_onto(
                         tcx,
                         impl_def_id.to_def_id(),