Propagate the resolved type of assoc const bindings via query feeding

9 files changed, 92 insertions, 38 deletions

diff --git a/compiler/rustc_hir_analysis/src/astconv/bounds.rs b/compiler/rustc_hir_analysis/src/astconv/bounds.rs
index 6940b4a5045..a9b4dca08d4 100644
--- a/compiler/rustc_hir_analysis/src/astconv/bounds.rs
+++ b/compiler/rustc_hir_analysis/src/astconv/bounds.rs
@@ -408,7 +408,7 @@ impl<'tcx> dyn AstConv<'tcx> + '_ {
             // Create the generic arguments for the associated type or constant by joining the
             // parent arguments (the arguments of the trait) and the own arguments (the ones of
             // the associated item itself) and construct an alias type using them.
-            candidate.map_bound(|trait_ref| {
+            let alias_ty = candidate.map_bound(|trait_ref| {
                 let ident = Ident::new(assoc_item.name, binding.ident.span);
                 let item_segment = hir::PathSegment {
                     ident,
@@ -430,7 +430,18 @@ impl<'tcx> dyn AstConv<'tcx> + '_ {
                 // *constants* to represent *const projections*. Alias *term* would be a more
                 // appropriate name but alas.
                 ty::AliasTy::new(tcx, assoc_item.def_id, alias_args)
-            })
+            });
+
+            // Provide the resolved type of the associated constant to `type_of(AnonConst)`.
+            if !speculative && let ty::AssocKind::Const = assoc_kind {
+                let ty = alias_ty.map_bound(|ty| tcx.type_of(ty.def_id).instantiate(tcx, ty.args));
+                // Since the arguments passed to the alias type above may contain early-bound
+                // generic parameters, the instantiated type may contain some as well.
+                // Therefore wrap it in `EarlyBinder`.
+                tcx.feed_type_of_assoc_const_binding(binding.hir_id, ty::EarlyBinder::bind(ty));
+            }
+
+            alias_ty
         };
 
         match binding.kind {
diff --git a/compiler/rustc_hir_analysis/src/collect.rs b/compiler/rustc_hir_analysis/src/collect.rs
index a5ef1490bce..d82d3eccfc6 100644
--- a/compiler/rustc_hir_analysis/src/collect.rs
+++ b/compiler/rustc_hir_analysis/src/collect.rs
@@ -63,6 +63,7 @@ pub fn provide(providers: &mut Providers) {
     *providers = Providers {
         type_of: type_of::type_of,
         type_of_opaque: type_of::type_of_opaque,
+        type_of_assoc_const_binding: type_of::type_of_assoc_const_binding,
         type_alias_is_lazy: type_of::type_alias_is_lazy,
         item_bounds: item_bounds::item_bounds,
         explicit_item_bounds: item_bounds::explicit_item_bounds,
diff --git a/compiler/rustc_hir_analysis/src/collect/type_of.rs b/compiler/rustc_hir_analysis/src/collect/type_of.rs
index c0128afe2bf..57bcf7602ea 100644
--- a/compiler/rustc_hir_analysis/src/collect/type_of.rs
+++ b/compiler/rustc_hir_analysis/src/collect/type_of.rs
@@ -78,35 +78,10 @@ fn anon_const_type_of<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> Ty<'tcx> {
                 .expect("const parameter types cannot be generic");
         }
 
-        Node::TypeBinding(binding @ &TypeBinding { hir_id: binding_id, .. })
-            if let Node::TraitRef(trait_ref) = tcx.parent_hir_node(binding_id) =>
-        {
-            let Some(trait_def_id) = trait_ref.trait_def_id() else {
-                return Ty::new_error_with_message(
-                    tcx,
-                    tcx.def_span(def_id),
-                    "Could not find trait",
-                );
-            };
-            let assoc_items = tcx.associated_items(trait_def_id);
-            let assoc_item = assoc_items.find_by_name_and_kind(
-                tcx,
-                binding.ident,
-                ty::AssocKind::Const,
-                def_id.to_def_id(),
-            );
-            return if let Some(assoc_item) = assoc_item {
-                tcx.type_of(assoc_item.def_id)
-                    .no_bound_vars()
-                    .expect("const parameter types cannot be generic")
-            } else {
-                // FIXME(associated_const_equality): add a useful error message here.
-                Ty::new_error_with_message(
-                    tcx,
-                    tcx.def_span(def_id),
-                    "Could not find associated const on trait",
-                )
-            };
+        Node::TypeBinding(&TypeBinding { hir_id, .. }) => {
+            // FIXME(fmease): Reject “escaping” early-bound generic parameters.
+            // FIXME(fmease): Reject escaping late-bound vars.
+            return tcx.type_of_assoc_const_binding(hir_id).skip_binder().skip_binder();
         }
 
         // This match arm is for when the def_id appears in a GAT whose
@@ -315,6 +290,18 @@ fn anon_const_type_of<'tcx>(tcx: TyCtxt<'tcx>, def_id: LocalDefId) -> Ty<'tcx> {
     }
 }
 
+pub(super) fn type_of_assoc_const_binding<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    hir_id: HirId,
+) -> ty::EarlyBinder<ty::Binder<'tcx, Ty<'tcx>>> {
+    let reported = tcx.dcx().delayed_bug(format!(
+        "attempt to obtain type of assoc const binding `{hir_id}` before \
+         it was resolved by `add_predicates_for_ast_type_binding`"
+    ));
+
+    ty::EarlyBinder::bind(ty::Binder::dummy(Ty::new_error(tcx, reported)))
+}
+
 fn get_path_containing_arg_in_pat<'hir>(
     pat: &'hir hir::Pat<'hir>,
     arg_id: HirId,
diff --git a/compiler/rustc_middle/src/query/erase.rs b/compiler/rustc_middle/src/query/erase.rs
index 2cdcdcb1492..82e0844a85c 100644
--- a/compiler/rustc_middle/src/query/erase.rs
+++ b/compiler/rustc_middle/src/query/erase.rs
@@ -193,6 +193,10 @@ impl<T: EraseType> EraseType for ty::EarlyBinder<T> {
     type Result = T::Result;
 }
 
+impl EraseType for ty::Binder<'_, Ty<'_>> {
+    type Result = [u8; size_of::<ty::Binder<'static, Ty<'static>>>()];
+}
+
 impl EraseType for ty::Binder<'_, ty::FnSig<'_>> {
     type Result = [u8; size_of::<ty::Binder<'static, ty::FnSig<'static>>>()];
 }
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
index 5be45c33e11..ce1e3484d69 100644
--- a/compiler/rustc_middle/src/query/mod.rs
+++ b/compiler/rustc_middle/src/query/mod.rs
@@ -248,6 +248,11 @@ rustc_queries! {
         cycle_stash
     }
 
+    query type_of_assoc_const_binding(key: hir::HirId) -> ty::EarlyBinder<ty::Binder<'tcx, Ty<'tcx>>> {
+        desc { |tcx| "getting type of associated constant binding `{key:?}`" }
+        feedable
+    }
+
     query type_alias_is_lazy(key: DefId) -> bool {
         desc { |tcx|
             "computing whether `{path}` is a lazy type alias",
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
index cc734e7157f..acd68a5bd87 100644
--- a/compiler/rustc_middle/src/ty/context.rs
+++ b/compiler/rustc_middle/src/ty/context.rs
@@ -71,6 +71,7 @@ use rustc_type_ir::TyKind::*;
 use rustc_type_ir::WithCachedTypeInfo;
 use rustc_type_ir::{CollectAndApply, Interner, TypeFlags};
 
+use std::assert_matches::debug_assert_matches;
 use std::borrow::Borrow;
 use std::cmp::Ordering;
 use std::fmt;
@@ -540,6 +541,22 @@ impl<'tcx> TyCtxt<'tcx> {
         debug_assert_eq!(self.def_kind(key), DefKind::AnonConst);
         TyCtxtFeed { tcx: self, key }.type_of(value)
     }
+
+    pub fn feed_type_of_assoc_const_binding(
+        self,
+        key: hir::HirId,
+        value: ty::EarlyBinder<ty::Binder<'tcx, Ty<'tcx>>>,
+    ) {
+        debug_assert_matches!(
+            self.hir_node(key),
+            hir::Node::TypeBinding(hir::TypeBinding {
+                kind: hir::TypeBindingKind::Equality { term: hir::Term::Const(_) },
+                ..
+            })
+        );
+
+        TyCtxtFeed { tcx: self, key }.type_of_assoc_const_binding(value)
+    }
 }
 
 impl<'tcx, KEY: Copy> TyCtxtFeed<'tcx, KEY> {
diff --git a/tests/ui/associated-consts/assoc-const-eq-supertraits.rs b/tests/ui/associated-consts/assoc-const-eq-supertraits.rs
new file mode 100644
index 00000000000..d5d724c9b15
--- /dev/null
+++ b/tests/ui/associated-consts/assoc-const-eq-supertraits.rs
@@ -0,0 +1,16 @@
+// Regression test for issue #118040.
+// Ensure that we support assoc const eq bounds where the assoc const comes from a supertrait.
+
+//@ check-pass
+
+#![feature(associated_const_equality)]
+
+trait Trait: SuperTrait {}
+trait SuperTrait: SuperSuperTrait<i32> {}
+trait SuperSuperTrait<T> {
+    const K: T;
+}
+
+fn take(_: impl Trait<K = 0>) {}
+
+fn main() {}
diff --git a/tests/ui/associated-consts/assoc-const-eq-ty-alias-noninteracting.rs b/tests/ui/associated-consts/assoc-const-eq-ty-alias-noninteracting.rs
index 4ff05112897..76df014ccd9 100644
--- a/tests/ui/associated-consts/assoc-const-eq-ty-alias-noninteracting.rs
+++ b/tests/ui/associated-consts/assoc-const-eq-ty-alias-noninteracting.rs
@@ -1,21 +1,25 @@
 // Regression test for issue #112560.
 // Respect the fact that (associated) types and constants live in different namespaces and
 // therefore equality bounds involving identically named associated items don't conflict if
-// their kind (type vs. const) differs.
-
-// FIXME(fmease): Extend this test to cover supertraits again
-// once #118040 is fixed. See initial version of PR #118360.
+// their kind (type vs. const) differs. This obviously extends to supertraits.
 
 //@ check-pass
 
 #![feature(associated_const_equality)]
 
-trait Trait {
+trait Trait: SuperTrait {
     type N;
+    type Q;
 
     const N: usize;
 }
 
-fn take(_: impl Trait<N = 0, N = ()>) {}
+trait SuperTrait {
+    const Q: &'static str;
+}
+
+fn take0(_: impl Trait<N = 0, N = ()>) {}
+
+fn take1(_: impl Trait<Q = "...", Q = [()]>) {}
 
 fn main() {}
diff --git a/tests/ui/generic-const-items/associated-const-equality.rs b/tests/ui/generic-const-items/associated-const-equality.rs
index 3c727097e2b..c0179f02fd2 100644
--- a/tests/ui/generic-const-items/associated-const-equality.rs
+++ b/tests/ui/generic-const-items/associated-const-equality.rs
@@ -1,22 +1,31 @@
 //@ check-pass
 
-#![feature(generic_const_items, associated_const_equality)]
+#![feature(generic_const_items, associated_const_equality, adt_const_params)]
 #![allow(incomplete_features)]
 
 trait Owner {
     const C<const N: u32>: u32;
     const K<const N: u32>: u32;
+    const Q<T>: Maybe<T>;
 }
 
 impl Owner for () {
     const C<const N: u32>: u32 = N;
     const K<const N: u32>: u32 = N + 1;
+    const Q<T>: Maybe<T> = Maybe::Nothing;
 }
 
 fn take0<const N: u32>(_: impl Owner<C<N> = { N }>) {}
 fn take1(_: impl Owner<K<99> = 100>) {}
+fn take2(_: impl Owner<Q<()> = { Maybe::Just(()) }>) {}
 
 fn main() {
     take0::<128>(());
     take1(());
 }
+
+#[derive(PartialEq, Eq, std::marker::ConstParamTy)]
+enum Maybe<T> {
+    Nothing,
+    Just(T),
+}