Fix const generics in GAT

5 files changed, 114 insertions, 0 deletions

diff --git a/compiler/rustc_typeck/src/collect/type_of.rs b/compiler/rustc_typeck/src/collect/type_of.rs
index e4eabca9c3b..ef17ab9ccd1 100644
--- a/compiler/rustc_typeck/src/collect/type_of.rs
+++ b/compiler/rustc_typeck/src/collect/type_of.rs
@@ -29,6 +29,48 @@ 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 {
+            Node::Ty(hir_ty @ Ty { kind: TyKind::Path(QPath::TypeRelative(_, segment)), .. }) => {
+                let 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(id).to_def_id();
+                let item_ctxt = &ItemCtxt::new(tcx, item_did) as &dyn crate::astconv::AstConv<'_>;
+                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);
+                }
+
+                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)),
diff --git a/compiler/rustc_typeck/src/lib.rs b/compiler/rustc_typeck/src/lib.rs
index fd44bafab6f..2284767debb 100644
--- a/compiler/rustc_typeck/src/lib.rs
+++ b/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)]
diff --git a/src/test/ui/generic-associated-types/const-generics-gat-in-trait-return-type-1.rs b/src/test/ui/generic-associated-types/const-generics-gat-in-trait-return-type-1.rs
new file mode 100644
index 00000000000..ab33ef6f244
--- /dev/null
+++ b/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]);
+}
diff --git a/src/test/ui/generic-associated-types/const-generics-gat-in-trait-return-type-2.rs b/src/test/ui/generic-associated-types/const-generics-gat-in-trait-return-type-2.rs
new file mode 100644
index 00000000000..ba9a82ae721
--- /dev/null
+++ b/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]);
+}
diff --git a/src/test/ui/generic-associated-types/const-generics-gat-in-trait-return-type-3.rs b/src/test/ui/generic-associated-types/const-generics-gat-in-trait-return-type-3.rs
new file mode 100644
index 00000000000..9da5334056a
--- /dev/null
+++ b/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]);
+}