GAT/const_generics: Allow with_opt_const_param to return GAT param def_id

compiler/rustc_typeck/src/collect/type_of.rs

@@ -29,6 +29,64 @@ pub(super) fn opt_const_param_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option<
         let parent_node = tcx.hir().get(parent_node_id);
 
         match parent_node {
+            // This matches on types who's paths couldn't be resolved without typeck'ing e.g.
+            //
+            // trait Foo {
+            //   type Assoc<const N1: usize>;
+            //   fn foo() -> Self::Assoc<3>;
+            //   // note: if the def_id argument is the 3 then in this example
+            //   // parent_node would be the node for Self::Assoc<_>
+            // }
+            // We didnt write <Self as Foo>::Assoc so the Self::Assoc<_> is lowered to QPath::TypeRelative.
+            // I believe this match arm is only needed for GAT but I am not 100% sure - BoxyUwU
+            Node::Ty(hir_ty @ Ty { kind: TyKind::Path(QPath::TypeRelative(_, segment)), .. }) => {
+                // Walk up from the parent_node to find an item so that
+                // we can resolve the relative path to an actual associated type.
+                // For the code example above, this item would be the Foo trait.
+                let item_hir_id = tcx
+                    .hir()
+                    .parent_iter(hir_id)
+                    .filter(|(_, node)| matches!(node, Node::Item(_)))
+                    .map(|(id, _)| id)
+                    .next()
+                    .unwrap();
+                let item_did = tcx.hir().local_def_id(item_hir_id).to_def_id();
+                let item_ctxt = &ItemCtxt::new(tcx, item_did) as &dyn crate::astconv::AstConv<'_>;
+
+                // This ty will be the actual associated type so that we can
+                // go through its generics to find which param our def_id corresponds to.
+                // For the code example above, this ty would be the Assoc<const N1: usize>.
+                let ty = item_ctxt.ast_ty_to_ty(hir_ty);
+                if let ty::Projection(projection) = ty.kind() {
+                    let generics = tcx.generics_of(projection.item_def_id);
+
+                    let arg_index = segment
+                        .args
+                        .and_then(|args| {
+                            args.args
+                                .iter()
+                                .filter(|arg| arg.is_const())
+                                .position(|arg| arg.id() == hir_id)
+                        })
+                        .unwrap_or_else(|| {
+                            bug!("no arg matching AnonConst in segment");
+                        });
+
+                    return generics
+                        .params
+                        .iter()
+                        .filter(|param| matches!(param.kind, ty::GenericParamDefKind::Const))
+                        .nth(arg_index)
+                        .map(|param| param.def_id);
+                }
+
+                // I dont think it's possible to reach this but I'm not 100% sure - BoxyUwU
+                tcx.sess.delay_span_bug(
+                    tcx.def_span(def_id),
+                    "unexpected non-GAT usage of an anon const",
+                );
+                return None;
+            }
             Node::Expr(&Expr {
                 kind:
                     ExprKind::MethodCall(segment, ..) | ExprKind::Path(QPath::TypeRelative(_, segment)),

compiler/rustc_typeck/src/lib.rs

@@ -56,6 +56,7 @@ This API is completely unstable and subject to change.
 */
 
 #![doc(html_root_url = "https://doc.rust-lang.org/nightly/nightly-rustc/")]
+#![feature(bindings_after_at)]
 #![feature(bool_to_option)]
 #![feature(box_syntax)]
 #![feature(crate_visibility_modifier)]

src/test/ui/generic-associated-types/const-generics-gat-in-trait-return-type-1.rs

@@ -0,0 +1,22 @@
+// run-pass
+#![feature(generic_associated_types)]
+#![allow(incomplete_features)]
+
+// This test unsures that with_opt_const_param returns the
+// def_id of the N param in the Foo::Assoc GAT.
+
+trait Foo {
+    type Assoc<const N: usize>;
+    fn foo(&self) -> Self::Assoc<3>;
+}
+
+impl Foo for () {
+    type Assoc<const N: usize> = [(); N];
+    fn foo(&self) -> Self::Assoc<3> {
+        [(); 3]
+    }
+}
+
+fn main() {
+    assert_eq!(().foo(), [(); 3]);
+}

src/test/ui/generic-associated-types/const-generics-gat-in-trait-return-type-2.rs

@@ -0,0 +1,22 @@
+// run-pass
+#![feature(generic_associated_types)]
+#![allow(incomplete_features)]
+
+// This test unsures that with_opt_const_param returns the
+// def_id of the N param in the Foo::Assoc GAT.
+
+trait Foo {
+    type Assoc<const N: usize>;
+    fn foo<const N: usize>(&self) -> Self::Assoc<N>;
+}
+
+impl Foo for () {
+    type Assoc<const N: usize> = [(); N];
+    fn foo<const N: usize>(&self) -> Self::Assoc<N> {
+        [(); N]
+    }
+}
+
+fn main() {
+    assert_eq!(().foo::<10>(), [(); 10]);
+}

src/test/ui/generic-associated-types/const-generics-gat-in-trait-return-type-3.rs

@@ -0,0 +1,27 @@
+// run-pass
+#![feature(generic_associated_types)]
+#![allow(incomplete_features)]
+
+// This test unsures that with_opt_const_param returns the
+// def_id of the N param in the Bar::Assoc GAT.
+
+trait Bar {
+    type Assoc<const N: usize>;
+}
+trait Foo: Bar {
+    fn foo(&self) -> Self::Assoc<3>;
+}
+
+impl Bar for () {
+    type Assoc<const N: usize> = [(); N];
+}
+
+impl Foo for () {
+    fn foo(&self) -> Self::Assoc<3> {
+        [(); 3]
+    }
+}
+
+fn main() {
+    assert_eq!(().foo(), [(); 3]);
+}