Use the opaque_types_defined_by query to cheaply check for whether a hidden type may be registered for an opaque type

9 files changed, 204 insertions, 153 deletions

diff --git a/compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs b/compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs
index b7253957edc..9e78e6acba5 100644
--- a/compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs
+++ b/compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs
@@ -735,7 +735,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 if let ty::subst::GenericArgKind::Type(ty) = ty.unpack()
                     && let ty::Alias(ty::Opaque, ty::AliasTy { def_id, .. }) = *ty.kind()
                     && let Some(def_id) = def_id.as_local()
-                    && self.opaque_type_origin(def_id, self.param_env).is_some() {
+                    && self.opaque_type_origin(def_id).is_some() {
                     return None;
                 }
             }
diff --git a/compiler/rustc_infer/src/infer/error_reporting/note_and_explain.rs b/compiler/rustc_infer/src/infer/error_reporting/note_and_explain.rs
index e9c3726f8c3..421eb807a14 100644
--- a/compiler/rustc_infer/src/infer/error_reporting/note_and_explain.rs
+++ b/compiler/rustc_infer/src/infer/error_reporting/note_and_explain.rs
@@ -1,10 +1,9 @@
-use crate::infer::opaque_types::may_define_impl_trait_in_assoc_ty_modulo_sig;
-
 use super::TypeErrCtxt;
 use rustc_errors::Applicability::{MachineApplicable, MaybeIncorrect};
 use rustc_errors::{pluralize, Diagnostic, MultiSpan};
 use rustc_hir as hir;
-use rustc_middle::traits::ObligationCauseCode::{self, MiscObligation};
+use rustc_hir::def::DefKind;
+use rustc_middle::traits::ObligationCauseCode;
 use rustc_middle::ty::error::ExpectedFound;
 use rustc_middle::ty::print::Printer;
 use rustc_middle::{
@@ -258,9 +257,9 @@ impl<T> Trait<T> for X {
                             );
                         }
                     }
-                    (ty::Alias(ty::Opaque, alias), _) | (_, ty::Alias(ty::Opaque, alias)) if matches!(cause.code(), MiscObligation) => {
-                        if let Some(def_id) = alias.def_id.as_local() {
-                            if may_define_impl_trait_in_assoc_ty_modulo_sig(tcx, body_owner_def_id.expect_local(), def_id).is_some() {
+                    (ty::Alias(ty::Opaque, alias), _) | (_, ty::Alias(ty::Opaque, alias)) if alias.def_id.is_local() && matches!(tcx.def_kind(body_owner_def_id), DefKind::AssocFn | DefKind::AssocConst) => {
+                        if tcx.is_type_alias_impl_trait(alias.def_id) {
+                            if !tcx.opaque_types_defined_by(body_owner_def_id.expect_local()).contains(&alias.def_id.expect_local()) {
                                 diag.span_note(tcx.def_span(body_owner_def_id), "\
                                     this item must have the opaque type in its signature \
                                     in order to be able to register hidden types");
diff --git a/compiler/rustc_infer/src/infer/opaque_types.rs b/compiler/rustc_infer/src/infer/opaque_types.rs
index 0b6f51652fd..545310ad351 100644
--- a/compiler/rustc_infer/src/infer/opaque_types.rs
+++ b/compiler/rustc_infer/src/infer/opaque_types.rs
@@ -3,10 +3,9 @@ use super::{DefineOpaqueTypes, InferResult};
 use crate::errors::OpaqueHiddenTypeDiag;
 use crate::infer::{DefiningAnchor, InferCtxt, InferOk};
 use crate::traits;
-use hir::def::DefKind;
 use hir::def_id::{DefId, LocalDefId};
 use hir::OpaqueTyOrigin;
-use rustc_data_structures::fx::{FxHashSet, FxIndexMap};
+use rustc_data_structures::fx::FxIndexMap;
 use rustc_data_structures::sync::Lrc;
 use rustc_hir as hir;
 use rustc_middle::traits::ObligationCause;
@@ -54,9 +53,7 @@ impl<'tcx> InferCtxt<'tcx> {
         }
         let mut obligations = vec![];
         let replace_opaque_type = |def_id: DefId| {
-            def_id
-                .as_local()
-                .map_or(false, |def_id| self.opaque_type_origin(def_id, param_env).is_some())
+            def_id.as_local().map_or(false, |def_id| self.opaque_type_origin(def_id).is_some())
         };
         let value = value.fold_with(&mut BottomUpFolder {
             tcx: self.tcx,
@@ -141,7 +138,7 @@ impl<'tcx> InferCtxt<'tcx> {
                         //     let x = || foo(); // returns the Opaque assoc with `foo`
                         // }
                         // ```
-                        self.opaque_type_origin(def_id, param_env)?
+                        self.opaque_type_origin(def_id)?
                     }
                     DefiningAnchor::Bubble => self.opaque_type_origin_unchecked(def_id),
                     DefiningAnchor::Error => return None,
@@ -152,9 +149,8 @@ impl<'tcx> InferCtxt<'tcx> {
                     // no one encounters it in practice.
                     // It does occur however in `fn fut() -> impl Future<Output = i32> { async { 42 } }`,
                     // where it is of no concern, so we only check for TAITs.
-                    if let Some(OpaqueTyOrigin::TyAlias { .. }) = b_def_id
-                        .as_local()
-                        .and_then(|b_def_id| self.opaque_type_origin(b_def_id, param_env))
+                    if let Some(OpaqueTyOrigin::TyAlias { .. }) =
+                        b_def_id.as_local().and_then(|b_def_id| self.opaque_type_origin(b_def_id))
                     {
                         self.tcx.sess.emit_err(OpaqueHiddenTypeDiag {
                             span: cause.span,
@@ -370,12 +366,8 @@ impl<'tcx> InferCtxt<'tcx> {
 
     /// Returns the origin of the opaque type `def_id` if we're currently
     /// in its defining scope.
-    #[instrument(skip(self, param_env), level = "trace", ret)]
-    pub fn opaque_type_origin(
-        &self,
-        def_id: LocalDefId,
-        param_env: ty::ParamEnv<'tcx>,
-    ) -> Option<OpaqueTyOrigin> {
+    #[instrument(skip(self), level = "trace", ret)]
+    pub fn opaque_type_origin(&self, def_id: LocalDefId) -> Option<OpaqueTyOrigin> {
         let opaque_hir_id = self.tcx.hir().local_def_id_to_hir_id(def_id);
         let parent_def_id = match self.defining_use_anchor {
             DefiningAnchor::Bubble | DefiningAnchor::Error => return None,
@@ -391,7 +383,7 @@ impl<'tcx> InferCtxt<'tcx> {
             // Named `type Foo = impl Bar;`
             hir::OpaqueTyOrigin::TyAlias { in_assoc_ty } => {
                 if in_assoc_ty {
-                    may_define_impl_trait_in_assoc_ty(self.tcx, parent_def_id, def_id, param_env)
+                    self.tcx.opaque_types_defined_by(parent_def_id).contains(&def_id)
                 } else {
                     may_define_opaque_type(self.tcx, parent_def_id, opaque_hir_id)
                 }
@@ -654,105 +646,3 @@ fn may_define_opaque_type(tcx: TyCtxt<'_>, def_id: LocalDefId, opaque_hir_id: hi
     );
     res
 }
-
-#[derive(Debug, TypeVisitable, Clone)]
-/// Helper datastructure containing the signature
-/// that the opaque type extraction logic uses for determining
-/// whether an opaque type may have its hidden types registered
-/// by an item.
-enum FnSigOrTy<'tcx> {
-    FnSig(ty::PolyFnSig<'tcx>),
-    Ty(Ty<'tcx>),
-}
-
-/// Checks that the item may register hidden types for the
-/// opaque type, if the opaque type shows up in its signature.
-#[instrument(level = "debug", skip(tcx), ret)]
-pub fn may_define_impl_trait_in_assoc_ty_modulo_sig<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    def_id: LocalDefId,
-    opaque_def_id: LocalDefId,
-) -> Option<impl TypeVisitable<TyCtxt<'tcx>>> {
-    let sig = match tcx.def_kind(def_id) {
-        DefKind::AssocFn => FnSigOrTy::FnSig(tcx.fn_sig(def_id).subst_identity()),
-        DefKind::AssocConst | DefKind::AssocTy => {
-            FnSigOrTy::Ty(tcx.type_of(def_id).subst_identity())
-        }
-        _ => return None,
-    };
-    let impl_id = tcx.local_parent(def_id);
-    trace!(?impl_id);
-    let mut assoc_id = opaque_def_id;
-    // Peel nested opaque types.
-    while let DefKind::OpaqueTy = tcx.def_kind(assoc_id) {
-        trace!(?assoc_id);
-        assoc_id = tcx.local_parent(assoc_id);
-    }
-    trace!(?assoc_id);
-    if !matches!(tcx.def_kind(assoc_id), DefKind::AssocTy) {
-        tcx.sess
-            .delay_span_bug(tcx.def_span(opaque_def_id), format!("{:?}", tcx.def_kind(assoc_id)));
-    }
-    let assoc_impl_id = tcx.local_parent(assoc_id);
-    trace!(?assoc_impl_id);
-
-    if impl_id != assoc_impl_id {
-        return None;
-    }
-
-    Some(sig)
-}
-
-#[instrument(level = "debug", skip(tcx, param_env), ret)]
-fn may_define_impl_trait_in_assoc_ty<'tcx>(
-    tcx: TyCtxt<'tcx>,
-    def_id: LocalDefId,
-    opaque_def_id: LocalDefId,
-    param_env: ty::ParamEnv<'tcx>,
-) -> bool {
-    let Some(sig) = may_define_impl_trait_in_assoc_ty_modulo_sig(tcx, def_id, opaque_def_id) else {
-        return false;
-    };
-
-    struct Visitor<'tcx> {
-        opaque_def_id: LocalDefId,
-        param_env: ty::ParamEnv<'tcx>,
-        tcx: TyCtxt<'tcx>,
-        seen: FxHashSet<LocalDefId>,
-    }
-
-    impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for Visitor<'tcx> {
-        type BreakTy = ();
-        #[instrument(level = "trace", skip(self), ret)]
-        fn visit_ty(&mut self, ty: Ty<'tcx>) -> ControlFlow<()> {
-            // FIXME(oli-obk): We should be checking if the associated type
-            // is mentioned instead of normalizing to find the opaque type.
-            // But that requires a way to figure out that a projection refers
-            // to a specific opaque type. That is probably doable by checking for
-            // `Self` as the `substs[0]`.
-            let normalized_ty = self.tcx.normalize_erasing_regions(self.param_env, ty);
-            if let ty::Alias(ty::Opaque, alias) = normalized_ty.kind() {
-                if let Some(def_id) = alias.def_id.as_local() {
-                    trace!(?alias.def_id);
-                    if def_id == self.opaque_def_id {
-                        return ControlFlow::Break(());
-                    }
-
-                    if self.seen.insert(def_id) {
-                        // Look into nested obligations like `impl Trait<Assoc = impl OtherTrait>`.
-                        for (pred, _) in self
-                            .tcx
-                            .explicit_item_bounds(alias.def_id)
-                            .subst_iter_copied(self.tcx, alias.substs)
-                        {
-                            pred.visit_with(self)?;
-                        }
-                    }
-                }
-            }
-            normalized_ty.super_visit_with(self)
-        }
-    }
-    sig.visit_with(&mut Visitor { opaque_def_id, param_env, tcx, seen: Default::default() })
-        .is_break()
-}
diff --git a/compiler/rustc_middle/src/query/erase.rs b/compiler/rustc_middle/src/query/erase.rs
index 28a9c1eef1a..b45f7caaabe 100644
--- a/compiler/rustc_middle/src/query/erase.rs
+++ b/compiler/rustc_middle/src/query/erase.rs
@@ -172,6 +172,10 @@ impl EraseType for ty::Binder<'_, ty::FnSig<'_>> {
     type Result = [u8; size_of::<ty::Binder<'static, ty::FnSig<'static>>>()];
 }
 
+impl EraseType for ty::Binder<'_, &'_ ty::List<Ty<'_>>> {
+    type Result = [u8; size_of::<ty::Binder<'static, &'static ty::List<Ty<'static>>>>()];
+}
+
 impl<T0, T1> EraseType for (&'_ T0, &'_ T1) {
     type Result = [u8; size_of::<(&'static (), &'static ())>()];
 }
diff --git a/compiler/rustc_middle/src/ty/sty.rs b/compiler/rustc_middle/src/ty/sty.rs
index 8d0737e1eee..31929e4f4f6 100644
--- a/compiler/rustc_middle/src/ty/sty.rs
+++ b/compiler/rustc_middle/src/ty/sty.rs
@@ -1265,7 +1265,7 @@ impl<'tcx> AliasTy<'tcx> {
 
     /// Extracts the underlying trait reference and own substs from this projection.
     /// For example, if this is a projection of `<T as StreamingIterator>::Item<'a>`,
-    /// then this function would return a `T: Iterator` trait reference and `['a]` as the own substs
+    /// then this function would return a `T: StreamingIterator` trait reference and `['a]` as the own substs
     pub fn trait_ref_and_own_substs(
         self,
         tcx: TyCtxt<'tcx>,
diff --git a/compiler/rustc_middle/src/ty/util.rs b/compiler/rustc_middle/src/ty/util.rs
index 97e1f69057e..9bab693156b 100644
--- a/compiler/rustc_middle/src/ty/util.rs
+++ b/compiler/rustc_middle/src/ty/util.rs
@@ -36,7 +36,12 @@ pub struct Discr<'tcx> {
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 pub enum CheckRegions {
     No,
+    /// Only permit early bound regions. This is useful for Adts which
+    /// can never have late bound regions.
     OnlyEarlyBound,
+    /// Permit both late bound and early bound regions. Use this for functions,
+    /// which frequently have late bound regions.
+    Bound,
 }
 
 #[derive(Copy, Clone, Debug)]
@@ -471,15 +476,21 @@ impl<'tcx> TyCtxt<'tcx> {
         ignore_regions: CheckRegions,
     ) -> Result<(), NotUniqueParam<'tcx>> {
         let mut seen = GrowableBitSet::default();
+        let mut seen_late = FxHashSet::default();
         for arg in substs {
             match arg.unpack() {
                 GenericArgKind::Lifetime(lt) => match (ignore_regions, lt.kind()) {
-                    (CheckRegions::OnlyEarlyBound, ty::ReEarlyBound(p)) => {
+                    (CheckRegions::Bound, ty::ReLateBound(di, reg)) => {
+                        if !seen_late.insert((di, reg)) {
+                            return Err(NotUniqueParam::DuplicateParam(lt.into()));
+                        }
+                    }
+                    (CheckRegions::OnlyEarlyBound | CheckRegions::Bound, ty::ReEarlyBound(p)) => {
                         if !seen.insert(p.index) {
                             return Err(NotUniqueParam::DuplicateParam(lt.into()));
                         }
                     }
-                    (CheckRegions::OnlyEarlyBound, _) => {
+                    (CheckRegions::OnlyEarlyBound | CheckRegions::Bound, _) => {
                         return Err(NotUniqueParam::NotParam(lt.into()));
                     }
                     (CheckRegions::No, _) => {}
diff --git a/compiler/rustc_ty_utils/messages.ftl b/compiler/rustc_ty_utils/messages.ftl
index 15a14112f4a..5bc3e3c00c9 100644
--- a/compiler/rustc_ty_utils/messages.ftl
+++ b/compiler/rustc_ty_utils/messages.ftl
@@ -55,3 +55,11 @@ ty_utils_multiple_array_fields_simd_type = monomorphising SIMD type `{$ty}` with
 ty_utils_oversized_simd_type = monomorphising SIMD type `{$ty}` of length greater than {$max_lanes}
 
 ty_utils_non_primitive_simd_type = monomorphising SIMD type `{$ty}` with a non-primitive-scalar (integer/float/pointer) element type `{$e_ty}`
+
+ty_utils_impl_trait_duplicate_arg = non-defining opaque type use in defining scope
+    .label = generic argument `{$arg}` used twice
+    .note = for this opaque type
+
+ty_utils_impl_trait_not_param = non-defining opaque type use in defining scope
+    .label = argument `{$arg}` is not a generic parameter
+    .note = for this opaque type
diff --git a/compiler/rustc_ty_utils/src/errors.rs b/compiler/rustc_ty_utils/src/errors.rs
index 3d3fc50e6e5..553bf40ef3a 100644
--- a/compiler/rustc_ty_utils/src/errors.rs
+++ b/compiler/rustc_ty_utils/src/errors.rs
@@ -1,7 +1,7 @@
 //! Errors emitted by ty_utils
 
 use rustc_macros::{Diagnostic, Subdiagnostic};
-use rustc_middle::ty::Ty;
+use rustc_middle::ty::{GenericArg, Ty};
 use rustc_span::Span;
 
 #[derive(Diagnostic)]
@@ -100,3 +100,25 @@ pub struct NonPrimitiveSimdType<'tcx> {
     pub ty: Ty<'tcx>,
     pub e_ty: Ty<'tcx>,
 }
+
+#[derive(Diagnostic)]
+#[diag(ty_utils_impl_trait_duplicate_arg)]
+pub struct DuplicateArg<'tcx> {
+    pub arg: GenericArg<'tcx>,
+    #[primary_span]
+    #[label]
+    pub span: Span,
+    #[note]
+    pub opaque_span: Span,
+}
+
+#[derive(Diagnostic)]
+#[diag(ty_utils_impl_trait_not_param)]
+pub struct NotParam<'tcx> {
+    pub arg: GenericArg<'tcx>,
+    #[primary_span]
+    #[label]
+    pub span: Span,
+    #[note]
+    pub opaque_span: Span,
+}
diff --git a/compiler/rustc_ty_utils/src/opaque_types.rs b/compiler/rustc_ty_utils/src/opaque_types.rs
index cdb8e4a439d..25ebb333bf7 100644
--- a/compiler/rustc_ty_utils/src/opaque_types.rs
+++ b/compiler/rustc_ty_utils/src/opaque_types.rs
@@ -1,46 +1,165 @@
+use rustc_data_structures::fx::FxHashSet;
+use rustc_errors::ErrorGuaranteed;
 use rustc_hir::{def::DefKind, def_id::LocalDefId};
-use rustc_middle::ty::util::CheckRegions;
+use rustc_middle::ty::util::{CheckRegions, NotUniqueParam};
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_middle::ty::{TypeSuperVisitable, TypeVisitable, TypeVisitor};
+use rustc_span::Span;
 use rustc_type_ir::AliasKind;
 use std::ops::ControlFlow;
 
+use crate::errors::{DuplicateArg, NotParam};
+
 struct OpaqueTypeCollector<'tcx> {
     tcx: TyCtxt<'tcx>,
     opaques: Vec<LocalDefId>,
+    /// The `DefId` of the item which we are collecting opaque types for.
+    item: LocalDefId,
+
+    /// Avoid infinite recursion due to recursive declarations.
+    seen: FxHashSet<LocalDefId>,
+}
+
+impl<'tcx> OpaqueTypeCollector<'tcx> {
+    fn collect(
+        tcx: TyCtxt<'tcx>,
+        item: LocalDefId,
+        val: ty::Binder<'tcx, impl TypeVisitable<TyCtxt<'tcx>>>,
+    ) -> Vec<LocalDefId> {
+        let mut collector = Self { tcx, opaques: Vec::new(), item, seen: Default::default() };
+        val.skip_binder().visit_with(&mut collector);
+        collector.opaques
+    }
+
+    fn span(&self) -> Span {
+        self.tcx.def_span(self.item)
+    }
+
+    fn parent(&self) -> Option<LocalDefId> {
+        match self.tcx.def_kind(self.item) {
+            DefKind::Fn => None,
+            DefKind::AssocFn | DefKind::AssocTy | DefKind::AssocConst => {
+                Some(self.tcx.local_parent(self.item))
+            }
+            other => span_bug!(
+                self.tcx.def_span(self.item),
+                "unhandled item with opaque types: {other:?}"
+            ),
+        }
+    }
 }
 
 impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for OpaqueTypeCollector<'tcx> {
-    fn visit_ty(&mut self, t: Ty<'tcx>) -> ControlFlow<!> {
+    type BreakTy = ErrorGuaranteed;
+
+    #[instrument(skip(self), ret, level = "trace")]
+    fn visit_ty(&mut self, t: Ty<'tcx>) -> ControlFlow<ErrorGuaranteed> {
         match t.kind() {
-            ty::Alias(AliasKind::Opaque, alias_ty) => {
-                if let Some(def_id) = alias_ty.def_id.as_local() {
-                    if self
-                        .tcx
-                        .uses_unique_generic_params(alias_ty.substs, CheckRegions::OnlyEarlyBound)
-                        .is_ok()
-                    {
-                        self.opaques.push(def_id);
-                        return ControlFlow::Continue(());
-                    } else {
-                        warn!(?t, "opaque types with non-unique params in sig: {t:?}");
+            ty::Alias(AliasKind::Opaque, alias_ty) if alias_ty.def_id.is_local() => {
+                if !self.seen.insert(alias_ty.def_id.expect_local()) {
+                    return ControlFlow::Continue(());
+                }
+                match self.tcx.uses_unique_generic_params(alias_ty.substs, CheckRegions::Bound) {
+                    Ok(()) => {
+                        // FIXME: implement higher kinded lifetime bounds on nested opaque types. They are not
+                        // supported at all, so this is sound to do, but once we want to support them, you'll
+                        // start seeing the error below.
+
+                        self.opaques.push(alias_ty.def_id.expect_local());
+
+                        // Collect opaque types nested within the associated type bounds of this opaque type.
+                        for (pred, _span) in self
+                            .tcx
+                            .explicit_item_bounds(alias_ty.def_id)
+                            .subst_iter_copied(self.tcx, alias_ty.substs)
+                        {
+                            trace!(?pred);
+                            pred.visit_with(self)?;
+                        }
+
+                        ControlFlow::Continue(())
+                    }
+                    Err(NotUniqueParam::NotParam(arg)) => {
+                        let err = self.tcx.sess.emit_err(NotParam {
+                            arg,
+                            span: self.span(),
+                            opaque_span: self.tcx.def_span(alias_ty.def_id),
+                        });
+                        ControlFlow::Break(err)
+                    }
+                    Err(NotUniqueParam::DuplicateParam(arg)) => {
+                        let err = self.tcx.sess.emit_err(DuplicateArg {
+                            arg,
+                            span: self.span(),
+                            opaque_span: self.tcx.def_span(alias_ty.def_id),
+                        });
+                        ControlFlow::Break(err)
                     }
                 }
             }
-            _ => {}
+            ty::Alias(AliasKind::Projection, alias_ty) => {
+                if let Some(parent) = self.parent() {
+                    trace!(?alias_ty);
+                    let (trait_ref, own_substs) = alias_ty.trait_ref_and_own_substs(self.tcx);
+
+                    trace!(?trait_ref, ?own_substs);
+                    // This avoids having to do normalization of `Self::AssocTy` by only
+                    // supporting the case of a method defining opaque types from assoc types
+                    // in the same impl block.
+                    if trait_ref.self_ty() == self.tcx.type_of(parent).subst_identity() {
+                        for assoc in self.tcx.associated_items(parent).in_definition_order() {
+                            trace!(?assoc);
+                            if assoc.trait_item_def_id == Some(alias_ty.def_id) {
+                                // We reconstruct the generic args of the associated type within the impl
+                                // from the impl's generics and the generic args passed to the type via the
+                                // projection.
+                                let substs = ty::InternalSubsts::identity_for_item(
+                                    self.tcx,
+                                    parent.to_def_id(),
+                                );
+                                trace!(?substs);
+                                let substs: Vec<_> =
+                                    substs.iter().chain(own_substs.iter().copied()).collect();
+                                trace!(?substs);
+                                // Find opaque types in this associated type.
+                                return self
+                                    .tcx
+                                    .type_of(assoc.def_id)
+                                    .subst(self.tcx, &substs)
+                                    .visit_with(self);
+                            }
+                        }
+                    }
+                }
+                t.super_visit_with(self)
+            }
+            _ => t.super_visit_with(self),
         }
-        t.super_visit_with(self)
     }
 }
 
 fn opaque_types_defined_by<'tcx>(tcx: TyCtxt<'tcx>, item: LocalDefId) -> &'tcx [LocalDefId] {
-    // FIXME(type_alias_impl_trait): This is definitely still wrong except for RPIT.
-    match tcx.def_kind(item) {
-        DefKind::Fn | DefKind::AssocFn => {
-            let sig = tcx.fn_sig(item).subst_identity();
-            let mut collector = OpaqueTypeCollector { tcx, opaques: Vec::new() };
-            sig.visit_with(&mut collector);
-            tcx.arena.alloc_from_iter(collector.opaques)
+    let kind = tcx.def_kind(item);
+    trace!(?kind);
+    // FIXME(type_alias_impl_trait): This is definitely still wrong except for RPIT and impl trait in assoc types.
+    match kind {
+        // We're also doing this for `AssocTy` for the wf checks in `check_opaque_meets_bounds`
+        DefKind::Fn | DefKind::AssocFn | DefKind::AssocTy | DefKind::AssocConst => {
+            let defined_opaques = match kind {
+                DefKind::Fn => {
+                    OpaqueTypeCollector::collect(tcx, item, tcx.fn_sig(item).subst_identity())
+                }
+                DefKind::AssocFn => {
+                    OpaqueTypeCollector::collect(tcx, item, tcx.fn_sig(item).subst_identity())
+                }
+                DefKind::AssocTy | DefKind::AssocConst => OpaqueTypeCollector::collect(
+                    tcx,
+                    item,
+                    ty::Binder::dummy(tcx.type_of(item).subst_identity()),
+                ),
+                _ => unreachable!(),
+            };
+            tcx.arena.alloc_from_iter(defined_opaques)
         }
         DefKind::Mod
         | DefKind::Struct
@@ -51,13 +170,11 @@ fn opaque_types_defined_by<'tcx>(tcx: TyCtxt<'tcx>, item: LocalDefId) -> &'tcx [
         | DefKind::TyAlias
         | DefKind::ForeignTy
         | DefKind::TraitAlias
-        | DefKind::AssocTy
         | DefKind::TyParam
         | DefKind::Const
         | DefKind::ConstParam
         | DefKind::Static(_)
         | DefKind::Ctor(_, _)
-        | DefKind::AssocConst
         | DefKind::Macro(_)
         | DefKind::ExternCrate
         | DefKind::Use