Validate `rustc_args_required_const`

4 files changed, 176 insertions, 10 deletions

diff --git a/compiler/rustc_mir/src/transform/promote_consts.rs b/compiler/rustc_mir/src/transform/promote_consts.rs
index cd361e430fa..89f7531b3a7 100644
--- a/compiler/rustc_mir/src/transform/promote_consts.rs
+++ b/compiler/rustc_mir/src/transform/promote_consts.rs
@@ -137,7 +137,7 @@ fn args_required_const(tcx: TyCtxt<'_>, def_id: DefId) -> Option<Vec<usize>> {
             LitKind::Int(a, _) => {
                 ret.push(a as usize);
             }
-            _ => return None,
+            _ => bug!("invalid arg index"),
         }
     }
     Some(ret)
diff --git a/compiler/rustc_passes/src/check_attr.rs b/compiler/rustc_passes/src/check_attr.rs
index efe947daa28..1ec251326a8 100644
--- a/compiler/rustc_passes/src/check_attr.rs
+++ b/compiler/rustc_passes/src/check_attr.rs
@@ -8,12 +8,12 @@ use rustc_middle::hir::map::Map;
 use rustc_middle::ty::query::Providers;
 use rustc_middle::ty::TyCtxt;
 
-use rustc_ast::{Attribute, NestedMetaItem};
-use rustc_errors::struct_span_err;
+use rustc_ast::{Attribute, LitKind, NestedMetaItem};
+use rustc_errors::{pluralize, struct_span_err};
 use rustc_hir as hir;
 use rustc_hir::def_id::LocalDefId;
 use rustc_hir::intravisit::{self, NestedVisitorMap, Visitor};
-use rustc_hir::{self, HirId, Item, ItemKind, TraitItem};
+use rustc_hir::{self, FnSig, ForeignItem, ForeignItemKind, HirId, Item, ItemKind, TraitItem};
 use rustc_hir::{MethodKind, Target};
 use rustc_session::lint::builtin::{CONFLICTING_REPR_HINTS, UNUSED_ATTRIBUTES};
 use rustc_session::parse::feature_err;
@@ -43,6 +43,12 @@ pub(crate) fn target_from_impl_item<'tcx>(
     }
 }
 
+#[derive(Clone, Copy)]
+enum ItemLike<'tcx> {
+    Item(&'tcx Item<'tcx>),
+    ForeignItem(&'tcx ForeignItem<'tcx>),
+}
+
 struct CheckAttrVisitor<'tcx> {
     tcx: TyCtxt<'tcx>,
 }
@@ -55,7 +61,7 @@ impl CheckAttrVisitor<'tcx> {
         attrs: &'hir [Attribute],
         span: &Span,
         target: Target,
-        item: Option<&Item<'_>>,
+        item: Option<ItemLike<'_>>,
     ) {
         let mut is_valid = true;
         for attr in attrs {
@@ -75,6 +81,8 @@ impl CheckAttrVisitor<'tcx> {
                 self.check_no_link(&attr, span, target)
             } else if self.tcx.sess.check_name(attr, sym::export_name) {
                 self.check_export_name(&attr, span, target)
+            } else if self.tcx.sess.check_name(attr, sym::rustc_args_required_const) {
+                self.check_rustc_args_required_const(&attr, span, target, item)
             } else {
                 // lint-only checks
                 if self.tcx.sess.check_name(attr, sym::cold) {
@@ -400,6 +408,71 @@ impl CheckAttrVisitor<'tcx> {
         }
     }
 
+    /// Checks if `#[rustc_args_required_const]` is applied to a function and has a valid argument.
+    fn check_rustc_args_required_const(
+        &self,
+        attr: &Attribute,
+        span: &Span,
+        target: Target,
+        item: Option<ItemLike<'_>>,
+    ) -> bool {
+        if let Target::Fn | Target::Method(..) | Target::ForeignFn = target {
+            let mut invalid_args = vec![];
+            for meta in attr.meta_item_list().expect("no meta item list") {
+                if let Some(LitKind::Int(val, _)) = meta.literal().map(|lit| &lit.kind) {
+                    if let Some(ItemLike::Item(Item {
+                        kind: ItemKind::Fn(FnSig { decl, .. }, ..),
+                        ..
+                    }))
+                    | Some(ItemLike::ForeignItem(ForeignItem {
+                        kind: ForeignItemKind::Fn(decl, ..),
+                        ..
+                    })) = item
+                    {
+                        let arg_count = decl.inputs.len() as u128;
+                        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");
+                    }
+                } else {
+                    invalid_args.push(meta.span());
+                }
+            }
+            if !invalid_args.is_empty() {
+                self.tcx
+                    .sess
+                    .struct_span_err(invalid_args, "arguments should be non-negative integers")
+                    .emit();
+                false
+            } else {
+                true
+            }
+        } else {
+            self.tcx
+                .sess
+                .struct_span_err(attr.span, "attribute should be applied to a function")
+                .span_label(*span, "not a function")
+                .emit();
+            false
+        }
+    }
+
     /// Checks if `#[link_section]` is applied to a function or static.
     fn check_link_section(&self, hir_id: HirId, attr: &Attribute, span: &Span, target: Target) {
         match target {
@@ -448,7 +521,7 @@ impl CheckAttrVisitor<'tcx> {
         attrs: &'hir [Attribute],
         span: &Span,
         target: Target,
-        item: Option<&Item<'_>>,
+        item: Option<ItemLike<'_>>,
         hir_id: HirId,
     ) {
         // Extract the names of all repr hints, e.g., [foo, bar, align] for:
@@ -564,7 +637,14 @@ impl CheckAttrVisitor<'tcx> {
         // Warn on repr(u8, u16), repr(C, simd), and c-like-enum-repr(C, u8)
         if (int_reprs > 1)
             || (is_simd && is_c)
-            || (int_reprs == 1 && is_c && item.map_or(false, |item| is_c_like_enum(item)))
+            || (int_reprs == 1
+                && is_c
+                && item.map_or(false, |item| {
+                    if let ItemLike::Item(item) = item {
+                        return is_c_like_enum(item);
+                    }
+                    return false;
+                }))
         {
             self.tcx.struct_span_lint_hir(
                 CONFLICTING_REPR_HINTS,
@@ -649,7 +729,13 @@ impl Visitor<'tcx> for CheckAttrVisitor<'tcx> {
 
     fn visit_item(&mut self, item: &'tcx Item<'tcx>) {
         let target = Target::from_item(item);
-        self.check_attributes(item.hir_id, item.attrs, &item.span, target, Some(item));
+        self.check_attributes(
+            item.hir_id,
+            item.attrs,
+            &item.span,
+            target,
+            Some(ItemLike::Item(item)),
+        );
         intravisit::walk_item(self, item)
     }
 
@@ -659,9 +745,15 @@ impl Visitor<'tcx> for CheckAttrVisitor<'tcx> {
         intravisit::walk_trait_item(self, trait_item)
     }
 
-    fn visit_foreign_item(&mut self, f_item: &'tcx hir::ForeignItem<'tcx>) {
+    fn visit_foreign_item(&mut self, f_item: &'tcx ForeignItem<'tcx>) {
         let target = Target::from_foreign_item(f_item);
-        self.check_attributes(f_item.hir_id, &f_item.attrs, &f_item.span, target, None);
+        self.check_attributes(
+            f_item.hir_id,
+            &f_item.attrs,
+            &f_item.span,
+            target,
+            Some(ItemLike::ForeignItem(f_item)),
+        );
         intravisit::walk_foreign_item(self, f_item)
     }
 
diff --git a/src/test/ui/invalid-rustc_args_required_const-arguments.rs b/src/test/ui/invalid-rustc_args_required_const-arguments.rs
new file mode 100644
index 00000000000..76c01c21301
--- /dev/null
+++ b/src/test/ui/invalid-rustc_args_required_const-arguments.rs
@@ -0,0 +1,26 @@
+#![feature(rustc_attrs)]
+
+#[rustc_args_required_const(0)] //~ ERROR index exceeds number of arguments
+fn foo1() {}
+
+#[rustc_args_required_const(1)] //~ ERROR index exceeds number of arguments
+fn foo2(_: u8) {}
+
+#[rustc_args_required_const(a)] //~ ERROR arguments should be non-negative integers
+fn foo4() {}
+
+#[rustc_args_required_const(1, a, 2, b)] //~ ERROR arguments should be non-negative integers
+fn foo5(_: u8, _: u8, _: u8) {}
+
+#[rustc_args_required_const(0)] //~ ERROR attribute should be applied to a function
+struct S;
+
+#[rustc_args_required_const(0usize)] //~ ERROR suffixed literals are not allowed in attributes
+fn foo6(_: u8) {}
+
+extern {
+    #[rustc_args_required_const(1)] //~ ERROR index exceeds number of arguments
+    fn foo7(_: u8);
+}
+
+fn main() {}
diff --git a/src/test/ui/invalid-rustc_args_required_const-arguments.stderr b/src/test/ui/invalid-rustc_args_required_const-arguments.stderr
new file mode 100644
index 00000000000..39d04626168
--- /dev/null
+++ b/src/test/ui/invalid-rustc_args_required_const-arguments.stderr
@@ -0,0 +1,48 @@
+error: suffixed literals are not allowed in attributes
+  --> $DIR/invalid-rustc_args_required_const-arguments.rs:18:29
+   |
+LL | #[rustc_args_required_const(0usize)]
+   |                             ^^^^^^
+   |
+   = help: instead of using a suffixed literal (`1u8`, `1.0f32`, etc.), use an unsuffixed version (`1`, `1.0`, etc.)
+
+error: index exceeds number of arguments
+  --> $DIR/invalid-rustc_args_required_const-arguments.rs:3:29
+   |
+LL | #[rustc_args_required_const(0)]
+   |                             ^ there are only 0 arguments
+
+error: index exceeds number of arguments
+  --> $DIR/invalid-rustc_args_required_const-arguments.rs:6:29
+   |
+LL | #[rustc_args_required_const(1)]
+   |                             ^ there is only 1 argument
+
+error: arguments should be non-negative integers
+  --> $DIR/invalid-rustc_args_required_const-arguments.rs:9:29
+   |
+LL | #[rustc_args_required_const(a)]
+   |                             ^
+
+error: arguments should be non-negative integers
+  --> $DIR/invalid-rustc_args_required_const-arguments.rs:12:32
+   |
+LL | #[rustc_args_required_const(1, a, 2, b)]
+   |                                ^     ^
+
+error: attribute should be applied to a function
+  --> $DIR/invalid-rustc_args_required_const-arguments.rs:15:1
+   |
+LL | #[rustc_args_required_const(0)]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+LL | struct S;
+   | --------- not a function
+
+error: index exceeds number of arguments
+  --> $DIR/invalid-rustc_args_required_const-arguments.rs:22:33
+   |
+LL |     #[rustc_args_required_const(1)]
+   |                                 ^ there is only 1 argument
+
+error: aborting due to 7 previous errors
+