Rollup merge of #92875 - BoxyUwU:infer_arg_opt_const_param_of, r=lcnr

Make `opt_const_param_of` work in the presence of `GenericArg::Infer`

highly recommend viewing the first and second commits on their own rather than looking at file changes :rofl:

Because we filtered args down to just const args we would ignore `GenericArg::Infer` which made us get a `arg_index` which was wrong by however many const `GenericArg::Infer` came previously

[example](https://play.rust-lang.org/?version=nightly&mode=debug&edition=2021&gist=46dba6a53aca6333028a10908ef16e0b) of the [bugs](https://play.rust-lang.org/?version=nightly&mode=debug&edition=2021&gist=a8eebced26eefa4119fc2e7ae0c76de6) fixed.

r? ```@lcnr```

6 files changed, 231 insertions, 219 deletions

diff --git a/compiler/rustc_hir/src/hir.rs b/compiler/rustc_hir/src/hir.rs
index f03d8eea40b..76f4df6ec2d 100644
--- a/compiler/rustc_hir/src/hir.rs
+++ b/compiler/rustc_hir/src/hir.rs
@@ -293,10 +293,6 @@ impl GenericArg<'_> {
         }
     }
 
-    pub fn is_const(&self) -> bool {
-        matches!(self, GenericArg::Const(_))
-    }
-
     pub fn is_synthetic(&self) -> bool {
         matches!(self, GenericArg::Lifetime(lifetime) if lifetime.name.ident() == Ident::empty())
     }
@@ -318,6 +314,13 @@ impl GenericArg<'_> {
             GenericArg::Infer(_) => ast::ParamKindOrd::Infer,
         }
     }
+
+    pub fn is_ty_or_const(&self) -> bool {
+        match self {
+            GenericArg::Lifetime(_) => false,
+            GenericArg::Type(_) | GenericArg::Const(_) | GenericArg::Infer(_) => true,
+        }
+    }
 }
 
 #[derive(Debug, HashStable_Generic)]
diff --git a/compiler/rustc_middle/src/ty/generics.rs b/compiler/rustc_middle/src/ty/generics.rs
index 1c3a01e2cfa..0bd96f8f865 100644
--- a/compiler/rustc_middle/src/ty/generics.rs
+++ b/compiler/rustc_middle/src/ty/generics.rs
@@ -31,6 +31,13 @@ impl GenericParamDefKind {
             GenericParamDefKind::Const { .. } => ast::ParamKindOrd::Const,
         }
     }
+
+    pub fn is_ty_or_const(&self) -> bool {
+        match self {
+            GenericParamDefKind::Lifetime => false,
+            GenericParamDefKind::Type { .. } | GenericParamDefKind::Const { .. } => true,
+        }
+    }
 }
 
 #[derive(Clone, Debug, TyEncodable, TyDecodable, HashStable)]
diff --git a/compiler/rustc_typeck/src/collect/type_of.rs b/compiler/rustc_typeck/src/collect/type_of.rs
index ae8d262fcf1..5ff2a747541 100644
--- a/compiler/rustc_typeck/src/collect/type_of.rs
+++ b/compiler/rustc_typeck/src/collect/type_of.rs
@@ -18,6 +18,7 @@ use super::{bad_placeholder, is_suggestable_infer_ty};
 /// Computes the relevant generic parameter for a potential generic const argument.
 ///
 /// This should be called using the query `tcx.opt_const_param_of`.
+#[instrument(level = "debug", skip(tcx))]
 pub(super) fn opt_const_param_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option<DefId> {
     // FIXME(generic_arg_infer): allow for returning DefIds of inference of
     // GenericArg::Infer below. This may require a change where GenericArg::Infer has some flag
@@ -25,231 +26,228 @@ pub(super) fn opt_const_param_of(tcx: TyCtxt<'_>, def_id: LocalDefId) -> Option<
     use hir::*;
     let hir_id = tcx.hir().local_def_id_to_hir_id(def_id);
 
-    if let Node::AnonConst(_) = tcx.hir().get(hir_id) {
-        let parent_node_id = tcx.hir().get_parent_node(hir_id);
-        let parent_node = tcx.hir().get(parent_node_id);
+    match tcx.hir().get(hir_id) {
+        Node::AnonConst(_) => (),
+        _ => return None,
+    };
 
-        match parent_node {
-            // This match arm is for when the def_id appears in a GAT whose
-            // path can't be resolved without typechecking e.g.
-            //
-            // trait Foo {
-            //   type Assoc<const N: usize>;
-            //   fn foo() -> Self::Assoc<3>;
-            // }
-            //
-            // In the above code we would call this query with the def_id of 3 and
-            // the parent_node we match on would be the hir node for Self::Assoc<3>
-            //
-            // `Self::Assoc<3>` cant be resolved without typchecking here as we
-            // didnt write <Self as Foo>::Assoc<3>. If we did then another match
-            // arm would handle this.
-            //
-            // 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)), .. }) => {
-                // Find the Item containing the associated type so we can create an ItemCtxt.
-                // Using the ItemCtxt convert the HIR for the unresolved assoc type into a
-                // ty which is a fully resolved projection.
-                // For the code example above, this would mean converting Self::Assoc<3>
-                // into a ty::Projection(<Self as Foo>::Assoc<3>)
-                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<'_>;
-                let ty = item_ctxt.ast_ty_to_ty(hir_ty);
-
-                // Iterate through the generics of the projection to find the one that corresponds to
-                // the def_id that this query was called with. We filter to only const args here as a
-                // precaution for if it's ever allowed to elide lifetimes in GAT's. It currently isn't
-                // but it can't hurt to be safe ^^
-                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");
-                        });
+    let parent_node_id = tcx.hir().get_parent_node(hir_id);
+    let parent_node = tcx.hir().get(parent_node_id);
 
-                    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)),
-                ..
-            }) => {
-                let body_owner = tcx.hir().local_def_id(tcx.hir().enclosing_body_owner(hir_id));
-                let tables = tcx.typeck(body_owner);
-                // This may fail in case the method/path does not actually exist.
-                // As there is no relevant param for `def_id`, we simply return
-                // `None` here.
-                let type_dependent_def = tables.type_dependent_def_id(parent_node_id)?;
-                let idx = segment
+    let (generics, arg_idx) = match parent_node {
+        // This match arm is for when the def_id appears in a GAT whose
+        // path can't be resolved without typechecking e.g.
+        //
+        // trait Foo {
+        //   type Assoc<const N: usize>;
+        //   fn foo() -> Self::Assoc<3>;
+        // }
+        //
+        // In the above code we would call this query with the def_id of 3 and
+        // the parent_node we match on would be the hir node for Self::Assoc<3>
+        //
+        // `Self::Assoc<3>` cant be resolved without typchecking here as we
+        // didnt write <Self as Foo>::Assoc<3>. If we did then another match
+        // arm would handle this.
+        //
+        // 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)), .. }) => {
+            // Find the Item containing the associated type so we can create an ItemCtxt.
+            // Using the ItemCtxt convert the HIR for the unresolved assoc type into a
+            // ty which is a fully resolved projection.
+            // For the code example above, this would mean converting Self::Assoc<3>
+            // into a ty::Projection(<Self as Foo>::Assoc<3>)
+            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<'_>;
+            let ty = item_ctxt.ast_ty_to_ty(hir_ty);
+
+            // Iterate through the generics of the projection to find the one that corresponds to
+            // the def_id that this query was called with. We filter to only const args here as a
+            // precaution for if it's ever allowed to elide lifetimes in GAT's. It currently isn't
+            // but it can't hurt to be safe ^^
+            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())
+                            .filter(|arg| arg.is_ty_or_const())
                             .position(|arg| arg.id() == hir_id)
                     })
                     .unwrap_or_else(|| {
                         bug!("no arg matching AnonConst in segment");
                     });
 
-                tcx.generics_of(type_dependent_def)
-                    .params
-                    .iter()
-                    .filter(|param| matches!(param.kind, ty::GenericParamDefKind::Const { .. }))
-                    .nth(idx)
-                    .map(|param| param.def_id)
+                (generics, arg_index)
+            } else {
+                // 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)),
+            ..
+        }) => {
+            let body_owner = tcx.hir().local_def_id(tcx.hir().enclosing_body_owner(hir_id));
+            let tables = tcx.typeck(body_owner);
+            // This may fail in case the method/path does not actually exist.
+            // As there is no relevant param for `def_id`, we simply return
+            // `None` here.
+            let type_dependent_def = tables.type_dependent_def_id(parent_node_id)?;
+            let idx = segment
+                .args
+                .and_then(|args| {
+                    args.args
+                        .iter()
+                        .filter(|arg| arg.is_ty_or_const())
+                        .position(|arg| arg.id() == hir_id)
+                })
+                .unwrap_or_else(|| {
+                    bug!("no arg matching AnonConst in segment");
+                });
 
-            Node::Ty(&Ty { kind: TyKind::Path(_), .. })
-            | Node::Expr(&Expr { kind: ExprKind::Path(_) | ExprKind::Struct(..), .. })
-            | Node::TraitRef(..)
-            | Node::Pat(_) => {
-                let path = match parent_node {
-                    Node::Ty(&Ty { kind: TyKind::Path(QPath::Resolved(_, path)), .. })
-                    | Node::TraitRef(&TraitRef { path, .. }) => &*path,
-                    Node::Expr(&Expr {
-                        kind:
-                            ExprKind::Path(QPath::Resolved(_, path))
-                            | ExprKind::Struct(&QPath::Resolved(_, path), ..),
-                        ..
-                    }) => {
-                        let body_owner =
-                            tcx.hir().local_def_id(tcx.hir().enclosing_body_owner(hir_id));
-                        let _tables = tcx.typeck(body_owner);
-                        &*path
-                    }
-                    Node::Pat(pat) => {
-                        if let Some(path) = get_path_containing_arg_in_pat(pat, hir_id) {
-                            path
-                        } else {
-                            tcx.sess.delay_span_bug(
-                                tcx.def_span(def_id),
-                                &format!(
-                                    "unable to find const parent for {} in pat {:?}",
-                                    hir_id, pat
-                                ),
-                            );
-                            return None;
-                        }
-                    }
-                    _ => {
-                        tcx.sess.delay_span_bug(
-                            tcx.def_span(def_id),
-                            &format!("unexpected const parent path {:?}", parent_node),
-                        );
-                        return None;
-                    }
-                };
-
-                // We've encountered an `AnonConst` in some path, so we need to
-                // figure out which generic parameter it corresponds to and return
-                // the relevant type.
-                let filtered = path
-                    .segments
-                    .iter()
-                    .filter_map(|seg| seg.args.map(|args| (args.args, seg)))
-                    .find_map(|(args, seg)| {
-                        args.iter()
-                            .filter(|arg| arg.is_const())
-                            .position(|arg| arg.id() == hir_id)
-                            .map(|index| (index, seg))
-                    });
-                let (arg_index, segment) = match filtered {
-                    None => {
-                        tcx.sess.delay_span_bug(
-                            tcx.def_span(def_id),
-                            "no arg matching AnonConst in path",
-                        );
-                        return None;
-                    }
-                    Some(inner) => inner,
-                };
-
-                // Try to use the segment resolution if it is valid, otherwise we
-                // default to the path resolution.
-                let res = segment.res.filter(|&r| r != Res::Err).unwrap_or(path.res);
-                use def::CtorOf;
-                let generics = match res {
-                    Res::Def(DefKind::Ctor(CtorOf::Variant, _), def_id) => tcx.generics_of(
-                        tcx.parent(def_id).and_then(|def_id| tcx.parent(def_id)).unwrap(),
-                    ),
-                    Res::Def(DefKind::Variant | DefKind::Ctor(CtorOf::Struct, _), def_id) => {
-                        tcx.generics_of(tcx.parent(def_id).unwrap())
-                    }
-                    // Other `DefKind`s don't have generics and would ICE when calling
-                    // `generics_of`.
-                    Res::Def(
-                        DefKind::Struct
-                        | DefKind::Union
-                        | DefKind::Enum
-                        | DefKind::Trait
-                        | DefKind::OpaqueTy
-                        | DefKind::TyAlias
-                        | DefKind::ForeignTy
-                        | DefKind::TraitAlias
-                        | DefKind::AssocTy
-                        | DefKind::Fn
-                        | DefKind::AssocFn
-                        | DefKind::AssocConst
-                        | DefKind::Impl,
-                        def_id,
-                    ) => tcx.generics_of(def_id),
-                    Res::Err => {
-                        tcx.sess.delay_span_bug(tcx.def_span(def_id), "anon const with Res::Err");
-                        return None;
-                    }
-                    _ => {
-                        // If the user tries to specify generics on a type that does not take them,
-                        // e.g. `usize<T>`, we may hit this branch, in which case we treat it as if
-                        // no arguments have been passed. An error should already have been emitted.
+            (tcx.generics_of(type_dependent_def), idx)
+        }
+
+        Node::Ty(&Ty { kind: TyKind::Path(_), .. })
+        | Node::Expr(&Expr { kind: ExprKind::Path(_) | ExprKind::Struct(..), .. })
+        | Node::TraitRef(..)
+        | Node::Pat(_) => {
+            let path = match parent_node {
+                Node::Ty(&Ty { kind: TyKind::Path(QPath::Resolved(_, path)), .. })
+                | Node::TraitRef(&TraitRef { path, .. }) => &*path,
+                Node::Expr(&Expr {
+                    kind:
+                        ExprKind::Path(QPath::Resolved(_, path))
+                        | ExprKind::Struct(&QPath::Resolved(_, path), ..),
+                    ..
+                }) => {
+                    let body_owner = tcx.hir().local_def_id(tcx.hir().enclosing_body_owner(hir_id));
+                    let _tables = tcx.typeck(body_owner);
+                    &*path
+                }
+                Node::Pat(pat) => {
+                    if let Some(path) = get_path_containing_arg_in_pat(pat, hir_id) {
+                        path
+                    } else {
                         tcx.sess.delay_span_bug(
                             tcx.def_span(def_id),
-                            &format!("unexpected anon const res {:?} in path: {:?}", res, path),
+                            &format!("unable to find const parent for {} in pat {:?}", hir_id, pat),
                         );
                         return None;
                     }
-                };
-
-                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),
+                        &format!("unexpected const parent path {:?}", parent_node),
+                    );
+                    return None;
+                }
+            };
+
+            // We've encountered an `AnonConst` in some path, so we need to
+            // figure out which generic parameter it corresponds to and return
+            // the relevant type.
+            let filtered = path
+                .segments
+                .iter()
+                .filter_map(|seg| seg.args.map(|args| (args.args, seg)))
+                .find_map(|(args, seg)| {
+                    args.iter()
+                        .filter(|arg| arg.is_ty_or_const())
+                        .position(|arg| arg.id() == hir_id)
+                        .map(|index| (index, seg))
+                });
+            let (arg_index, segment) = match filtered {
+                None => {
+                    tcx.sess
+                        .delay_span_bug(tcx.def_span(def_id), "no arg matching AnonConst in path");
+                    return None;
+                }
+                Some(inner) => inner,
+            };
+
+            // Try to use the segment resolution if it is valid, otherwise we
+            // default to the path resolution.
+            let res = segment.res.filter(|&r| r != Res::Err).unwrap_or(path.res);
+            use def::CtorOf;
+            let generics = match res {
+                Res::Def(DefKind::Ctor(CtorOf::Variant, _), def_id) => tcx
+                    .generics_of(tcx.parent(def_id).and_then(|def_id| tcx.parent(def_id)).unwrap()),
+                Res::Def(DefKind::Variant | DefKind::Ctor(CtorOf::Struct, _), def_id) => {
+                    tcx.generics_of(tcx.parent(def_id).unwrap())
+                }
+                // Other `DefKind`s don't have generics and would ICE when calling
+                // `generics_of`.
+                Res::Def(
+                    DefKind::Struct
+                    | DefKind::Union
+                    | DefKind::Enum
+                    | DefKind::Trait
+                    | DefKind::OpaqueTy
+                    | DefKind::TyAlias
+                    | DefKind::ForeignTy
+                    | DefKind::TraitAlias
+                    | DefKind::AssocTy
+                    | DefKind::Fn
+                    | DefKind::AssocFn
+                    | DefKind::AssocConst
+                    | DefKind::Impl,
+                    def_id,
+                ) => tcx.generics_of(def_id),
+                Res::Err => {
+                    tcx.sess.delay_span_bug(tcx.def_span(def_id), "anon const with Res::Err");
+                    return None;
+                }
+                _ => {
+                    // If the user tries to specify generics on a type that does not take them,
+                    // e.g. `usize<T>`, we may hit this branch, in which case we treat it as if
+                    // no arguments have been passed. An error should already have been emitted.
+                    tcx.sess.delay_span_bug(
+                        tcx.def_span(def_id),
+                        &format!("unexpected anon const res {:?} in path: {:?}", res, path),
+                    );
+                    return None;
+                }
+            };
+
+            (generics, arg_index)
+        }
+        _ => return None,
+    };
+
+    debug!(?parent_node);
+    debug!(?generics, ?arg_idx);
+    generics
+        .params
+        .iter()
+        .filter(|param| param.kind.is_ty_or_const())
+        .nth(match generics.has_self && generics.parent.is_none() {
+            true => arg_idx + 1,
+            false => arg_idx,
+        })
+        .and_then(|param| match param.kind {
+            ty::GenericParamDefKind::Const { .. } => {
+                debug!(?param);
+                Some(param.def_id)
             }
             _ => None,
-        }
-    } else {
-        None
-    }
+        })
 }
 
 fn get_path_containing_arg_in_pat<'hir>(
diff --git a/src/test/ui/const-generics/generic_arg_infer/infer_arg_and_const_arg.rs b/src/test/ui/const-generics/generic_arg_infer/infer_arg_and_const_arg.rs
new file mode 100644
index 00000000000..23c8d753752
--- /dev/null
+++ b/src/test/ui/const-generics/generic_arg_infer/infer_arg_and_const_arg.rs
@@ -0,0 +1,12 @@
+// check-pass
+#![feature(generic_arg_infer)]
+
+struct Foo<const N: bool, const M: u8>;
+struct Bar<const N: u8, const M: u32>;
+
+fn main() {
+    let _: Foo<true, _> = Foo::<_, 1>;
+    let _: Foo<_, 1> = Foo::<true, _>;
+    let _: Bar<1, _> = Bar::<_, 300>;
+    let _: Bar<_, 300> = Bar::<1, _>;
+}
diff --git a/src/test/ui/const-generics/issues/issue-62878.full.stderr b/src/test/ui/const-generics/issues/issue-62878.full.stderr
index f074a65313f..3a2b291d7ba 100644
--- a/src/test/ui/const-generics/issues/issue-62878.full.stderr
+++ b/src/test/ui/const-generics/issues/issue-62878.full.stderr
@@ -4,13 +4,6 @@ error[E0770]: the type of const parameters must not depend on other generic para
 LL | fn foo<const N: usize, const A: [u8; N]>() {}
    |                                      ^ the type must not depend on the parameter `N`
 
-error[E0308]: mismatched types
-  --> $DIR/issue-62878.rs:10:15
-   |
-LL |     foo::<_, {[1]}>();
-   |               ^^^ expected `usize`, found array `[{integer}; 1]`
-
-error: aborting due to 2 previous errors
+error: aborting due to previous error
 
-Some errors have detailed explanations: E0308, E0770.
-For more information about an error, try `rustc --explain E0308`.
+For more information about this error, try `rustc --explain E0770`.
diff --git a/src/test/ui/const-generics/issues/issue-62878.rs b/src/test/ui/const-generics/issues/issue-62878.rs
index 38f5ff77b56..578ce765b2f 100644
--- a/src/test/ui/const-generics/issues/issue-62878.rs
+++ b/src/test/ui/const-generics/issues/issue-62878.rs
@@ -7,6 +7,5 @@ fn foo<const N: usize, const A: [u8; N]>() {}
 //[min]~| ERROR `[u8; _]` is forbidden as the type of a const generic parameter
 
 fn main() {
-    foo::<_, {[1]}>();
-    //[full]~^ ERROR mismatched types
+    foo::<_, { [1] }>();
 }