Improve checking for attribute

1 files changed, 56 insertions, 47 deletions

diff --git a/compiler/rustc_passes/src/check_attr.rs b/compiler/rustc_passes/src/check_attr.rs
index 022f1858878..cf08881cd47 100644
--- a/compiler/rustc_passes/src/check_attr.rs
+++ b/compiler/rustc_passes/src/check_attr.rs
@@ -760,7 +760,7 @@ impl CheckAttrVisitor<'tcx> {
         target: Target,
         item: Option<ItemLike<'_>>,
     ) -> bool {
-        let is_function = matches!(target, Target::Fn | Target::Method(..) | Target::ForeignFn);
+        let is_function = matches!(target, Target::Fn | Target::Method(..));
         if !is_function {
             self.tcx
                 .sess
@@ -776,55 +776,64 @@ impl CheckAttrVisitor<'tcx> {
             Some(it) => it,
         };
 
+        let (decl, generics) = match item {
+            Some(ItemLike::Item(Item {
+                kind: ItemKind::Fn(FnSig { decl, .. }, generics, _),
+                ..
+            })) => (decl, generics),
+            _ => bug!("should be a function item"),
+        };
+
+        for param in generics.params {
+            match param.kind {
+                hir::GenericParamKind::Const { .. } => {}
+                _ => {
+                    self.tcx
+                        .sess
+                        .struct_span_err(
+                            attr.span,
+                            "#[rustc_legacy_const_generics] functions must \
+                             only have const generics",
+                        )
+                        .span_label(param.span, "non-const generic parameter")
+                        .emit();
+                    return false;
+                }
+            }
+        }
+
+        if list.len() != generics.params.len() {
+            self.tcx
+                .sess
+                .struct_span_err(
+                    attr.span,
+                    "#[rustc_legacy_const_generics] must have one index for each generic parameter",
+                )
+                .span_label(generics.span, "generic parameters")
+                .emit();
+            return false;
+        }
+
+        let arg_count = decl.inputs.len() as u128 + generics.params.len() as u128;
         let mut invalid_args = vec![];
         for meta in list {
             if let Some(LitKind::Int(val, _)) = meta.literal().map(|lit| &lit.kind) {
-                if let Some(ItemLike::Item(Item {
-                    kind: ItemKind::Fn(FnSig { decl, .. }, generics, _),
-                    ..
-                }))
-                | Some(ItemLike::ForeignItem(ForeignItem {
-                    kind: ForeignItemKind::Fn(decl, _, generics),
-                    ..
-                })) = item
-                {
-                    let arg_count = decl.inputs.len() as u128 + generics.params.len() as u128;
-                    for param in generics.params {
-                        match param.kind {
-                            hir::GenericParamKind::Const { .. } => {}
-                            _ => {
-                                self.tcx
-                                    .sess
-                                    .struct_span_err(
-                                        meta.span(),
-                                        "#[rustc_legacy_const_generics] functions must \
-                                         only have const generics",
-                                    )
-                                    .span_label(param.span, "non-const generic parameter")
-                                    .emit();
-                                break;
-                            }
-                        }
-                    }
-                    if *val >= arg_count {
-                        let span = meta.span();
-                        self.tcx
-                            .sess
-                            .struct_span_err(span, "index exceeds number of arguments")
-                            .span_label(
-                                span,
-                                format!(
-                                    "there {} only {} argument{}",
-                                    if arg_count != 1 { "are" } else { "is" },
-                                    arg_count,
-                                    pluralize!(arg_count)
-                                ),
-                            )
-                            .emit();
-                        return false;
-                    }
-                } else {
-                    bug!("should be a function item");
+                if *val >= arg_count {
+                    let span = meta.span();
+                    self.tcx
+                        .sess
+                        .struct_span_err(span, "index exceeds number of arguments")
+                        .span_label(
+                            span,
+                            format!(
+                                "there {} only {} argument{}",
+                                if arg_count != 1 { "are" } else { "is" },
+                                arg_count,
+                                pluralize!(arg_count)
+                            ),
+                        )
+                        .emit();
+                    return false;
                 }
             } else {
                 invalid_args.push(meta.span());