Cleanup assoc parent utils

16 files changed, 79 insertions, 64 deletions

diff --git a/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs b/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs
index 5d9416b59fc..c0ca35f9ff8 100644
--- a/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs
+++ b/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs
@@ -677,12 +677,13 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
     ///  - is the trait from the local crate? If not, we can't suggest changing signatures
     ///  - `Span` of the argument in the trait definition
     fn is_error_in_trait(&self, local: Local) -> (bool, bool, Option<Span>) {
+        let tcx = self.infcx.tcx;
         if self.body.local_kind(local) != LocalKind::Arg {
             return (false, false, None);
         }
         let my_def = self.body.source.def_id();
         let Some(td) =
-            self.infcx.tcx.impl_of_assoc(my_def).and_then(|x| self.infcx.tcx.trait_id_of_impl(x))
+            tcx.trait_impl_of_assoc(my_def).and_then(|id| self.infcx.tcx.trait_id_of_impl(id))
         else {
             return (false, false, None);
         };
diff --git a/compiler/rustc_borrowck/src/type_check/mod.rs b/compiler/rustc_borrowck/src/type_check/mod.rs
index c3aa205d5aa..a960b96b91c 100644
--- a/compiler/rustc_borrowck/src/type_check/mod.rs
+++ b/compiler/rustc_borrowck/src/type_check/mod.rs
@@ -1773,10 +1773,7 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                     locations,
                 );
 
-                assert!(!matches!(
-                    tcx.impl_of_assoc(def_id).map(|imp| tcx.def_kind(imp)),
-                    Some(DefKind::Impl { of_trait: true })
-                ));
+                assert_eq!(tcx.trait_impl_of_assoc(def_id), None);
                 self.prove_predicates(
                     args.types().map(|ty| ty::ClauseKind::WellFormed(ty.into())),
                     locations,
diff --git a/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs b/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs
index 6cbf2dbf7d3..2c3a84499ac 100644
--- a/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs
+++ b/compiler/rustc_codegen_llvm/src/debuginfo/mod.rs
@@ -533,31 +533,26 @@ impl<'ll, 'tcx> DebugInfoCodegenMethods<'tcx> for CodegenCx<'ll, 'tcx> {
             // First, let's see if this is a method within an inherent impl. Because
             // if yes, we want to make the result subroutine DIE a child of the
             // subroutine's self-type.
-            if let Some(impl_def_id) = cx.tcx.impl_of_assoc(instance.def_id()) {
-                // If the method does *not* belong to a trait, proceed
-                if cx.tcx.trait_id_of_impl(impl_def_id).is_none() {
-                    let impl_self_ty = cx.tcx.instantiate_and_normalize_erasing_regions(
-                        instance.args,
-                        cx.typing_env(),
-                        cx.tcx.type_of(impl_def_id),
-                    );
-
-                    // Only "class" methods are generally understood by LLVM,
-                    // so avoid methods on other types (e.g., `<*mut T>::null`).
-                    if let ty::Adt(def, ..) = impl_self_ty.kind()
-                        && !def.is_box()
-                    {
-                        // Again, only create type information if full debuginfo is enabled
-                        if cx.sess().opts.debuginfo == DebugInfo::Full && !impl_self_ty.has_param()
-                        {
-                            return (type_di_node(cx, impl_self_ty), true);
-                        } else {
-                            return (namespace::item_namespace(cx, def.did()), false);
-                        }
+            // For trait method impls we still use the "parallel namespace"
+            // strategy
+            if let Some(imp_def_id) = cx.tcx.inherent_impl_of_assoc(instance.def_id()) {
+                let impl_self_ty = cx.tcx.instantiate_and_normalize_erasing_regions(
+                    instance.args,
+                    cx.typing_env(),
+                    cx.tcx.type_of(imp_def_id),
+                );
+
+                // Only "class" methods are generally understood by LLVM,
+                // so avoid methods on other types (e.g., `<*mut T>::null`).
+                if let ty::Adt(def, ..) = impl_self_ty.kind()
+                    && !def.is_box()
+                {
+                    // Again, only create type information if full debuginfo is enabled
+                    if cx.sess().opts.debuginfo == DebugInfo::Full && !impl_self_ty.has_param() {
+                        return (type_di_node(cx, impl_self_ty), true);
+                    } else {
+                        return (namespace::item_namespace(cx, def.did()), false);
                     }
-                } else {
-                    // For trait method impls we still use the "parallel namespace"
-                    // strategy
                 }
             }
 
diff --git a/compiler/rustc_const_eval/src/interpret/call.rs b/compiler/rustc_const_eval/src/interpret/call.rs
index b1cc0cc2878..21237cfe492 100644
--- a/compiler/rustc_const_eval/src/interpret/call.rs
+++ b/compiler/rustc_const_eval/src/interpret/call.rs
@@ -731,7 +731,7 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
     ) {
         let tcx = *self.tcx;
 
-        let trait_def_id = tcx.trait_of_assoc(def_id).unwrap();
+        let trait_def_id = tcx.parent(def_id);
         let virtual_trait_ref = ty::TraitRef::from_assoc(tcx, trait_def_id, virtual_instance.args);
         let existential_trait_ref = ty::ExistentialTraitRef::erase_self_ty(tcx, virtual_trait_ref);
         let concrete_trait_ref = existential_trait_ref.with_self_ty(tcx, dyn_ty);
diff --git a/compiler/rustc_hir_analysis/src/check/compare_impl_item.rs b/compiler/rustc_hir_analysis/src/check/compare_impl_item.rs
index 6767e5ed88d..e4827256193 100644
--- a/compiler/rustc_hir_analysis/src/check/compare_impl_item.rs
+++ b/compiler/rustc_hir_analysis/src/check/compare_impl_item.rs
@@ -445,10 +445,10 @@ pub(super) fn collect_return_position_impl_trait_in_trait_tys<'tcx>(
     tcx: TyCtxt<'tcx>,
     impl_m_def_id: LocalDefId,
 ) -> Result<&'tcx DefIdMap<ty::EarlyBinder<'tcx, Ty<'tcx>>>, ErrorGuaranteed> {
-    let impl_m = tcx.opt_associated_item(impl_m_def_id.to_def_id()).unwrap();
-    let trait_m = tcx.opt_associated_item(impl_m.trait_item_def_id.unwrap()).unwrap();
+    let impl_m = tcx.associated_item(impl_m_def_id.to_def_id());
+    let trait_m = tcx.associated_item(impl_m.trait_item_def_id.unwrap());
     let impl_trait_ref =
-        tcx.impl_trait_ref(impl_m.impl_container(tcx).unwrap()).unwrap().instantiate_identity();
+        tcx.impl_trait_ref(tcx.parent(impl_m_def_id.to_def_id())).unwrap().instantiate_identity();
     // First, check a few of the same things as `compare_impl_method`,
     // just so we don't ICE during instantiation later.
     check_method_is_structurally_compatible(tcx, impl_m, trait_m, impl_trait_ref, true)?;
diff --git a/compiler/rustc_hir_analysis/src/check/wfcheck.rs b/compiler/rustc_hir_analysis/src/check/wfcheck.rs
index c642435b989..e6a1f6d8d8b 100644
--- a/compiler/rustc_hir_analysis/src/check/wfcheck.rs
+++ b/compiler/rustc_hir_analysis/src/check/wfcheck.rs
@@ -2287,8 +2287,7 @@ fn lint_redundant_lifetimes<'tcx>(
             // Proceed
         }
         DefKind::AssocFn | DefKind::AssocTy | DefKind::AssocConst => {
-            let parent_def_id = tcx.local_parent(owner_id);
-            if matches!(tcx.def_kind(parent_def_id), DefKind::Impl { of_trait: true }) {
+            if tcx.trait_impl_of_assoc(owner_id.to_def_id()).is_some() {
                 // Don't check for redundant lifetimes for associated items of trait
                 // implementations, since the signature is required to be compatible
                 // with the trait, even if the implementation implies some lifetimes
diff --git a/compiler/rustc_lint/src/pass_by_value.rs b/compiler/rustc_lint/src/pass_by_value.rs
index 4f65acd8001..29006732aad 100644
--- a/compiler/rustc_lint/src/pass_by_value.rs
+++ b/compiler/rustc_lint/src/pass_by_value.rs
@@ -24,10 +24,8 @@ impl<'tcx> LateLintPass<'tcx> for PassByValue {
     fn check_ty(&mut self, cx: &LateContext<'_>, ty: &'tcx hir::Ty<'tcx, AmbigArg>) {
         match &ty.kind {
             TyKind::Ref(_, hir::MutTy { ty: inner_ty, mutbl: hir::Mutability::Not }) => {
-                if let Some(impl_did) = cx.tcx.impl_of_assoc(ty.hir_id.owner.to_def_id()) {
-                    if cx.tcx.impl_trait_ref(impl_did).is_some() {
-                        return;
-                    }
+                if cx.tcx.trait_impl_of_assoc(ty.hir_id.owner.to_def_id()).is_some() {
+                    return;
                 }
                 if let Some(t) = path_for_pass_by_value(cx, inner_ty) {
                     cx.emit_span_lint(
diff --git a/compiler/rustc_lint/src/types.rs b/compiler/rustc_lint/src/types.rs
index b0afc333ebe..f8a692313f0 100644
--- a/compiler/rustc_lint/src/types.rs
+++ b/compiler/rustc_lint/src/types.rs
@@ -1904,10 +1904,9 @@ impl InvalidAtomicOrdering {
         if let ExprKind::MethodCall(method_path, _, args, _) = &expr.kind
             && recognized_names.contains(&method_path.ident.name)
             && let Some(m_def_id) = cx.typeck_results().type_dependent_def_id(expr.hir_id)
-            && let Some(impl_did) = cx.tcx.impl_of_assoc(m_def_id)
-            && let Some(adt) = cx.tcx.type_of(impl_did).instantiate_identity().ty_adt_def()
             // skip extension traits, only lint functions from the standard library
-            && cx.tcx.trait_id_of_impl(impl_did).is_none()
+            && let Some(impl_did) = cx.tcx.inherent_impl_of_assoc(m_def_id)
+            && let Some(adt) = cx.tcx.type_of(impl_did).instantiate_identity().ty_adt_def()
             && let parent = cx.tcx.parent(adt.did())
             && cx.tcx.is_diagnostic_item(sym::atomic_mod, parent)
             && ATOMIC_TYPES.contains(&cx.tcx.item_name(adt.did()))
diff --git a/compiler/rustc_middle/src/middle/privacy.rs b/compiler/rustc_middle/src/middle/privacy.rs
index 785ddd1ee29..e3e04c9d180 100644
--- a/compiler/rustc_middle/src/middle/privacy.rs
+++ b/compiler/rustc_middle/src/middle/privacy.rs
@@ -181,11 +181,7 @@ impl EffectiveVisibilities {
             // nominal visibility. For some items nominal visibility doesn't make sense so we
             // don't check this condition for them.
             let is_impl = matches!(tcx.def_kind(def_id), DefKind::Impl { .. });
-            let is_associated_item_in_trait_impl = tcx
-                .impl_of_assoc(def_id.to_def_id())
-                .and_then(|impl_id| tcx.trait_id_of_impl(impl_id))
-                .is_some();
-            if !is_impl && !is_associated_item_in_trait_impl {
+            if !is_impl && tcx.trait_impl_of_assoc(def_id.to_def_id()).is_none() {
                 let nominal_vis = tcx.visibility(def_id);
                 if !nominal_vis.is_at_least(ev.reachable, tcx) {
                     span_bug!(
diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
index 73e1661106e..e70c98ab704 100644
--- a/compiler/rustc_middle/src/ty/mod.rs
+++ b/compiler/rustc_middle/src/ty/mod.rs
@@ -1925,21 +1925,50 @@ impl<'tcx> TyCtxt<'tcx> {
         self.impl_trait_ref(def_id).map(|tr| tr.skip_binder().def_id)
     }
 
-    /// If the given `DefId` is an associated item, returns the `DefId` of the parent trait or impl.
-    pub fn assoc_parent(self, def_id: DefId) -> Option<DefId> {
-        self.def_kind(def_id).is_assoc().then(|| self.parent(def_id))
+    /// If the given `DefId` is an associated item, returns the `DefId` and `DefKind` of the parent trait or impl.
+    pub fn assoc_parent(self, def_id: DefId) -> Option<(DefId, DefKind)> {
+        if !self.def_kind(def_id).is_assoc() {
+            return None;
+        }
+        let parent = self.parent(def_id);
+        let def_kind = self.def_kind(parent);
+        Some((parent, def_kind))
     }
 
     /// 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> {
-        self.assoc_parent(def_id).filter(|id| self.def_kind(id) == DefKind::Trait)
+        match self.assoc_parent(def_id) {
+            Some((id, DefKind::Trait)) => Some(id),
+            _ => None,
+        }
     }
 
     /// If the given `DefId` is an associated item of an impl,
     /// returns the `DefId` of the impl; otherwise returns `None`.
     pub fn impl_of_assoc(self, def_id: DefId) -> Option<DefId> {
-        self.assoc_parent(def_id).filter(|id| matches!(self.def_kind(id), DefKind::Impl { .. }))
+        match self.assoc_parent(def_id) {
+            Some((id, DefKind::Impl { .. })) => Some(id),
+            _ => None,
+        }
+    }
+
+    /// If the given `DefId` is an associated item of an inherent impl,
+    /// returns the `DefId` of the impl; otherwise, returns `None`.
+    pub fn inherent_impl_of_assoc(self, def_id: DefId) -> Option<DefId> {
+        match self.assoc_parent(def_id) {
+            Some((id, DefKind::Impl { of_trait: false })) => Some(id),
+            _ => None,
+        }
+    }
+
+    /// If the given `DefId` is an associated item of a trait impl,
+    /// returns the `DefId` of the impl; otherwise, returns `None`.
+    pub fn trait_impl_of_assoc(self, def_id: DefId) -> Option<DefId> {
+        match self.assoc_parent(def_id) {
+            Some((id, DefKind::Impl { of_trait: true })) => Some(id),
+            _ => None,
+        }
     }
 
     pub fn is_exportable(self, def_id: DefId) -> bool {
diff --git a/compiler/rustc_mir_transform/src/check_call_recursion.rs b/compiler/rustc_mir_transform/src/check_call_recursion.rs
index 6d61ac2dd80..a9acb1da5a3 100644
--- a/compiler/rustc_mir_transform/src/check_call_recursion.rs
+++ b/compiler/rustc_mir_transform/src/check_call_recursion.rs
@@ -43,8 +43,8 @@ impl<'tcx> MirLint<'tcx> for CheckDropRecursion {
 
         // First check if `body` is an `fn drop()` of `Drop`
         if let DefKind::AssocFn = tcx.def_kind(def_id)
-        && let Some(trait_ref) =
-            tcx.impl_of_assoc(def_id.to_def_id()).and_then(|def_id| tcx.impl_trait_ref(def_id))
+        && let Some(impl_id) = tcx.trait_impl_of_assoc(def_id.to_def_id())
+        && let trait_ref = tcx.impl_trait_ref(impl_id).unwrap()
         && tcx.is_lang_item(trait_ref.instantiate_identity().def_id, LangItem::Drop)
         // avoid erroneous `Drop` impls from causing ICEs below
         && let sig = tcx.fn_sig(def_id).instantiate_identity()
diff --git a/compiler/rustc_mir_transform/src/check_packed_ref.rs b/compiler/rustc_mir_transform/src/check_packed_ref.rs
index dcb812c7899..100104e9de0 100644
--- a/compiler/rustc_mir_transform/src/check_packed_ref.rs
+++ b/compiler/rustc_mir_transform/src/check_packed_ref.rs
@@ -40,7 +40,7 @@ impl<'tcx> Visitor<'tcx> for PackedRefChecker<'_, 'tcx> {
         if context.is_borrow() && util::is_disaligned(self.tcx, self.body, self.typing_env, *place)
         {
             let def_id = self.body.source.instance.def_id();
-            if let Some(impl_def_id) = self.tcx.impl_of_assoc(def_id)
+            if let Some(impl_def_id) = self.tcx.trait_impl_of_assoc(def_id)
                 && self.tcx.is_builtin_derived(impl_def_id)
             {
                 // If we ever reach here it means that the generated derive
diff --git a/compiler/rustc_mir_transform/src/shim.rs b/compiler/rustc_mir_transform/src/shim.rs
index c687036f544..c6760b3583f 100644
--- a/compiler/rustc_mir_transform/src/shim.rs
+++ b/compiler/rustc_mir_transform/src/shim.rs
@@ -75,7 +75,7 @@ fn make_shim<'tcx>(tcx: TyCtxt<'tcx>, instance: ty::InstanceKind<'tcx>) -> Body<
             build_call_shim(tcx, instance, Some(adjustment), CallKind::Direct(def_id))
         }
         ty::InstanceKind::FnPtrShim(def_id, ty) => {
-            let trait_ = tcx.trait_of_assoc(def_id).unwrap();
+            let trait_ = tcx.parent(def_id);
             // Supports `Fn` or `async Fn` traits.
             let adjustment = match tcx
                 .fn_trait_kind_from_def_id(trait_)
diff --git a/compiler/rustc_monomorphize/src/partitioning.rs b/compiler/rustc_monomorphize/src/partitioning.rs
index d76b27d9970..54d77b182f8 100644
--- a/compiler/rustc_monomorphize/src/partitioning.rs
+++ b/compiler/rustc_monomorphize/src/partitioning.rs
@@ -650,17 +650,18 @@ fn characteristic_def_id_of_mono_item<'tcx>(
             // its self-type. If the self-type does not provide a characteristic
             // DefId, we use the location of the impl after all.
 
-            if tcx.trait_of_assoc(def_id).is_some() {
+            let assoc_parent = tcx.assoc_parent(def_id);
+
+            if let Some((_, DefKind::Trait)) = assoc_parent {
                 let self_ty = instance.args.type_at(0);
                 // This is a default implementation of a trait method.
                 return characteristic_def_id_of_type(self_ty).or(Some(def_id));
             }
 
-            if let Some(impl_def_id) = tcx.impl_of_assoc(def_id) {
-                if tcx.sess.opts.incremental.is_some()
-                    && tcx
-                        .trait_id_of_impl(impl_def_id)
-                        .is_some_and(|def_id| tcx.is_lang_item(def_id, LangItem::Drop))
+            if let Some((impl_def_id, DefKind::Impl { of_trait })) = assoc_parent {
+                if of_trait
+                    && tcx.sess.opts.incremental.is_some()
+                    && tcx.is_lang_item(tcx.trait_id_of_impl(impl_def_id).unwrap(), LangItem::Drop)
                 {
                     // Put `Drop::drop` into the same cgu as `drop_in_place`
                     // since `drop_in_place` is the only thing that can
diff --git a/compiler/rustc_passes/src/dead.rs b/compiler/rustc_passes/src/dead.rs
index de52973acbb..08d06402000 100644
--- a/compiler/rustc_passes/src/dead.rs
+++ b/compiler/rustc_passes/src/dead.rs
@@ -371,7 +371,7 @@ impl<'tcx> MarkSymbolVisitor<'tcx> {
     /// will be ignored for the purposes of dead code analysis (see PR #85200
     /// for discussion).
     fn should_ignore_item(&mut self, def_id: DefId) -> bool {
-        if let Some(impl_of) = self.tcx.impl_of_assoc(def_id) {
+        if let Some(impl_of) = self.tcx.trait_impl_of_assoc(def_id) {
             if !self.tcx.is_automatically_derived(impl_of) {
                 return false;
             }
diff --git a/compiler/rustc_trait_selection/src/traits/project.rs b/compiler/rustc_trait_selection/src/traits/project.rs
index 581191b2036..884d53732fe 100644
--- a/compiler/rustc_trait_selection/src/traits/project.rs
+++ b/compiler/rustc_trait_selection/src/traits/project.rs
@@ -1507,7 +1507,7 @@ fn confirm_builtin_candidate<'cx, 'tcx>(
     let tcx = selcx.tcx();
     let self_ty = obligation.predicate.self_ty();
     let item_def_id = obligation.predicate.def_id;
-    let trait_def_id = tcx.trait_of_assoc(item_def_id).unwrap();
+    let trait_def_id = tcx.parent(item_def_id);
     let args = tcx.mk_args(&[self_ty.into()]);
     let (term, obligations) = if tcx.is_lang_item(trait_def_id, LangItem::DiscriminantKind) {
         let discriminant_def_id =