Propagate the resolved type of assoc const bindings via query feeding

Author: fmease

compiler/rustc_hir_analysis/src/astconv/bounds.rs

@@ -404,8 +404,10 @@ impl<'tcx> dyn AstConv<'tcx> + '_ {
             let bound_vars = tcx.late_bound_vars(binding.hir_id);
             ty::Binder::bind_with_vars(subst_output, bound_vars)
         } else {
-            // Append the generic arguments of the associated type or const to the `trait_ref`.
-            candidate.map_bound(|trait_ref| {
+            // 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.
+            let alias_ty = candidate.map_bound(|trait_ref| {
                 let ident = Ident::new(assoc_item.name, binding.ident.span);
                 let item_segment = hir::PathSegment {
                     ident,
@@ -427,7 +429,20 @@ 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 {
+                // FIXME(fmease): Is this `map_bound` + `instantiate` correct? Consider:
+                // `const K<T>: for<'a> fn(&'a T);`, `P<K<for<'x> fn(&'x ())> = {…}>` or similar.
+                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 {

compiler/rustc_hir_analysis/src/collect.rs

@@ -60,6 +60,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,

compiler/rustc_hir_analysis/src/collect/type_of.rs

@@ -66,35 +66,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.hir_node(tcx.hir().parent_id(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
@@ -303,6 +278,21 @@ 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 guar = tcx.dcx().span_delayed_bug(
+        rustc_span::DUMMY_SP,
+        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, guar)))
+}
+
 fn get_path_containing_arg_in_pat<'hir>(
     pat: &'hir hir::Pat<'hir>,
     arg_id: HirId,

compiler/rustc_middle/src/query/erase.rs

@@ -184,6 +184,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>>>()];
 }

compiler/rustc_middle/src/query/mod.rs

@@ -255,6 +255,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",

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;
@@ -533,6 +534,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> {

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() {}

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() {}

tests/ui/generic-const-items/associated-const-equality.rs

@@ -6,15 +6,18 @@
 trait Owner {
     const C<const N: u32>: u32;
     const K<const N: u32>: u32;
+    const Q<T>: Option<T>;
 }
 
 impl Owner for () {
     const C<const N: u32>: u32 = N;
     const K<const N: u32>: u32 = N + 1;
+    const Q<T>: Option<T> = None;
 }
 
 fn take0<const N: u32>(_: impl Owner<C<N> = { N }>) {}
 fn take1(_: impl Owner<K<99> = 100>) {}
+fn take2(_: impl Owner<Q<()> = { Some(()) }>) {}
 
 fn main() {
     take0::<128>(());