Improve proc macro attribute diagnostics

compiler/rustc_error_messages/locales/en-US/passes.ftl

@@ -707,3 +707,24 @@ passes_ignored_derived_impls =
       [one] trait {$trait_list}, but this is
      *[other] traits {$trait_list}, but these are
     } intentionally ignored during dead code analysis
+
+passes_proc_macro_typeerror = mismatched {$kind} signature
+    .label = found {$found}, expected type `proc_macro::TokenStream`
+    .note = {$kind}s must have a signature of `{$expected_signature}`
+
+passes_proc_macro_diff_arg_count = mismatched {$kind} signature
+    .label = found unexpected {$count ->
+      [one] argument
+     *[other] arguments
+    }
+    .note = {$kind}s must have a signature of `{$expected_signature}`
+
+passes_proc_macro_missing_args = mismatched {$kind} signature
+    .label = {$kind} must have {$expected_input_count ->
+      [one] one argument
+     *[other] two arguments
+    } of type `proc_macro::TokenStream`
+
+passes_proc_macro_invalid_abi = proc macro functions may not be `extern "{$abi}"`
+
+passes_proc_macro_unsafe = proc macro functions may not be `unsafe`

compiler/rustc_passes/src/check_attr.rs

@@ -6,11 +6,12 @@
 
 use crate::errors::{
     self, AttrApplication, DebugVisualizerUnreadable, InvalidAttrAtCrateLevel, ObjectLifetimeErr,
-    OnlyHasEffectOn, TransparentIncompatible, UnrecognizedReprHint,
+    OnlyHasEffectOn, ProcMacroDiffArguments, ProcMacroInvalidAbi, ProcMacroMissingArguments,
+    ProcMacroTypeError, ProcMacroUnsafe, TransparentIncompatible, UnrecognizedReprHint,
 };
 use rustc_ast::{ast, AttrStyle, Attribute, LitKind, MetaItemKind, MetaItemLit, NestedMetaItem};
 use rustc_data_structures::fx::FxHashMap;
-use rustc_errors::{fluent, Applicability, MultiSpan};
+use rustc_errors::{fluent, Applicability, IntoDiagnosticArg, MultiSpan};
 use rustc_expand::base::resolve_path;
 use rustc_feature::{AttributeDuplicates, AttributeType, BuiltinAttribute, BUILTIN_ATTRIBUTE_MAP};
 use rustc_hir as hir;
@@ -19,18 +20,20 @@ use rustc_hir::intravisit::{self, Visitor};
 use rustc_hir::{
     self, FnSig, ForeignItem, HirId, Item, ItemKind, TraitItem, CRATE_HIR_ID, CRATE_OWNER_ID,
 };
-use rustc_hir::{MethodKind, Target};
+use rustc_hir::{MethodKind, Target, Unsafety};
 use rustc_middle::hir::nested_filter;
 use rustc_middle::middle::resolve_lifetime::ObjectLifetimeDefault;
+use rustc_middle::ty::fast_reject::{DeepRejectCtxt, TreatParams};
 use rustc_middle::ty::query::Providers;
-use rustc_middle::ty::TyCtxt;
+use rustc_middle::ty::{ParamEnv, TyCtxt};
 use rustc_session::lint::builtin::{
     CONFLICTING_REPR_HINTS, INVALID_DOC_ATTRIBUTES, UNUSED_ATTRIBUTES,
 };
 use rustc_session::parse::feature_err;
 use rustc_span::symbol::{kw, sym, Symbol};
 use rustc_span::{Span, DUMMY_SP};
 use rustc_target::spec::abi::Abi;
+use std::cell::Cell;
 use std::collections::hash_map::Entry;
 
 pub(crate) fn target_from_impl_item<'tcx>(
@@ -62,8 +65,29 @@ enum ItemLike<'tcx> {
     ForeignItem,
 }
 
+#[derive(Copy, Clone)]
+pub(crate) enum ProcMacroKind {
+    FunctionLike,
+    Derive,
+    Attribute,
+}
+
+impl IntoDiagnosticArg for ProcMacroKind {
+    fn into_diagnostic_arg(self) -> rustc_errors::DiagnosticArgValue<'static> {
+        match self {
+            ProcMacroKind::Attribute => "attribute proc macro",
+            ProcMacroKind::Derive => "derive proc macro",
+            ProcMacroKind::FunctionLike => "function-like proc macro",
+        }
+        .into_diagnostic_arg()
+    }
+}
+
 struct CheckAttrVisitor<'tcx> {
     tcx: TyCtxt<'tcx>,
+
+    // Whether or not this visitor should abort after finding errors
+    abort: Cell<bool>,
 }
 
 impl CheckAttrVisitor<'_> {
@@ -172,7 +196,7 @@ impl CheckAttrVisitor<'_> {
                 sym::path => self.check_generic_attr(hir_id, attr, target, Target::Mod),
                 sym::plugin_registrar => self.check_plugin_registrar(hir_id, attr, target),
                 sym::macro_export => self.check_macro_export(hir_id, attr, target),
-                sym::ignore | sym::should_panic | sym::proc_macro_derive => {
+                sym::ignore | sym::should_panic => {
                     self.check_generic_attr(hir_id, attr, target, Target::Fn)
                 }
                 sym::automatically_derived => {
@@ -182,6 +206,16 @@ impl CheckAttrVisitor<'_> {
                     self.check_generic_attr(hir_id, attr, target, Target::Mod)
                 }
                 sym::rustc_object_lifetime_default => self.check_object_lifetime_default(hir_id),
+                sym::proc_macro => {
+                    self.check_proc_macro(hir_id, target, ProcMacroKind::FunctionLike)
+                }
+                sym::proc_macro_attribute => {
+                    self.check_proc_macro(hir_id, target, ProcMacroKind::Attribute);
+                }
+                sym::proc_macro_derive => {
+                    self.check_generic_attr(hir_id, attr, target, Target::Fn);
+                    self.check_proc_macro(hir_id, target, ProcMacroKind::Derive)
+                }
                 _ => {}
             }
 
@@ -2052,6 +2086,103 @@ impl CheckAttrVisitor<'_> {
             errors::Unused { attr_span: attr.span, note },
         );
     }
+
+    /// A best effort attempt to create an error for a mismatching proc macro signature.
+    ///
+    /// If this best effort goes wrong, it will just emit a worse error later (see #102923)
+    fn check_proc_macro(&self, hir_id: HirId, target: Target, kind: ProcMacroKind) {
+        let expected_input_count = match kind {
+            ProcMacroKind::Attribute => 2,
+            ProcMacroKind::Derive | ProcMacroKind::FunctionLike => 1,
+        };
+
+        let expected_signature = match kind {
+            ProcMacroKind::Attribute => "fn(TokenStream, TokenStream) -> TokenStream",
+            ProcMacroKind::Derive | ProcMacroKind::FunctionLike => "fn(TokenStream) -> TokenStream",
+        };
+
+        let tcx = self.tcx;
+        if target == Target::Fn {
+            let Some(tokenstream) = tcx.get_diagnostic_item(sym::TokenStream) else {return};
+            let tokenstream = tcx.type_of(tokenstream);
+
+            let id = hir_id.expect_owner();
+            let hir_sig = tcx.hir().fn_sig_by_hir_id(hir_id).unwrap();
+
+            let sig = tcx.liberate_late_bound_regions(id.to_def_id(), tcx.fn_sig(id));
+            let sig = tcx.normalize_erasing_regions(ParamEnv::empty(), sig);
+
+            // We don't currently require that the function signature is equal to
+            // `fn(TokenStream) -> TokenStream`, but instead monomorphizes to
+            // `fn(TokenStream) -> TokenStream` after some substitution of generic arguments.
+            //
+            // Properly checking this means pulling in additional `rustc` crates, so we don't.
+            let drcx = DeepRejectCtxt { treat_obligation_params: TreatParams::AsInfer };
+
+            if sig.abi != Abi::Rust {
+                tcx.sess.emit_err(ProcMacroInvalidAbi { span: hir_sig.span, abi: sig.abi.name() });
+                self.abort.set(true);
+            }
+
+            if sig.unsafety == Unsafety::Unsafe {
+                tcx.sess.emit_err(ProcMacroUnsafe { span: hir_sig.span });
+                self.abort.set(true);
+            }
+
+            let output = sig.output();
+
+            // Typecheck the output
+            if !drcx.types_may_unify(output, tokenstream) {
+                tcx.sess.emit_err(ProcMacroTypeError {
+                    span: hir_sig.decl.output.span(),
+                    found: output,
+                    kind,
+                    expected_signature,
+                });
+                self.abort.set(true);
+            }
+
+            if sig.inputs().len() < expected_input_count {
+                tcx.sess.emit_err(ProcMacroMissingArguments {
+                    expected_input_count,
+                    span: hir_sig.span,
+                    kind,
+                    expected_signature,
+                });
+                self.abort.set(true);
+            }
+
+            // Check that the inputs are correct, if there are enough.
+            if sig.inputs().len() >= expected_input_count {
+                for (arg, input) in
+                    sig.inputs().iter().zip(hir_sig.decl.inputs).take(expected_input_count)
+                {
+                    if !drcx.types_may_unify(*arg, tokenstream) {
+                        tcx.sess.emit_err(ProcMacroTypeError {
+                            span: input.span,
+                            found: *arg,
+                            kind,
+                            expected_signature,
+                        });
+                        self.abort.set(true);
+                    }
+                }
+            }
+
+            // Check that there are not too many arguments
+            let body_id = tcx.hir().body_owned_by(id.def_id);
+            let excess = tcx.hir().body(body_id).params.get(expected_input_count..);
+            if let Some(excess @ [begin @ end] | excess @ [begin, .., end]) = excess {
+                tcx.sess.emit_err(ProcMacroDiffArguments {
+                    span: begin.span.to(end.span),
+                    count: excess.len(),
+                    kind,
+                    expected_signature,
+                });
+                self.abort.set(true);
+            }
+        }
+    }
 }
 
 impl<'tcx> Visitor<'tcx> for CheckAttrVisitor<'tcx> {
@@ -2214,12 +2345,15 @@ fn check_non_exported_macro_for_invalid_attrs(tcx: TyCtxt<'_>, item: &Item<'_>)
 }
 
 fn check_mod_attrs(tcx: TyCtxt<'_>, module_def_id: LocalDefId) {
-    let check_attr_visitor = &mut CheckAttrVisitor { tcx };
+    let check_attr_visitor = &mut CheckAttrVisitor { tcx, abort: Cell::new(false) };
     tcx.hir().visit_item_likes_in_module(module_def_id, check_attr_visitor);
     if module_def_id.is_top_level_module() {
         check_attr_visitor.check_attributes(CRATE_HIR_ID, DUMMY_SP, Target::Mod, None);
         check_invalid_crate_level_attr(tcx, tcx.hir().krate_attrs());
     }
+    if check_attr_visitor.abort.get() {
+        tcx.sess.abort_if_errors()
+    }
 }
 
 pub(crate) fn provide(providers: &mut Providers) {

compiler/rustc_passes/src/errors.rs

@@ -12,6 +12,7 @@ use rustc_macros::{Diagnostic, LintDiagnostic, Subdiagnostic};
 use rustc_middle::ty::{MainDefinition, Ty};
 use rustc_span::{Span, Symbol, DUMMY_SP};
 
+use crate::check_attr::ProcMacroKind;
 use crate::lang_items::Duplicate;
 
 #[derive(LintDiagnostic)]
@@ -1508,3 +1509,52 @@ pub struct ChangeFieldsToBeOfUnitType {
     #[suggestion_part(code = "()")]
     pub spans: Vec<Span>,
 }
+
+#[derive(Diagnostic)]
+#[diag(passes_proc_macro_typeerror)]
+#[note]
+pub(crate) struct ProcMacroTypeError<'tcx> {
+    #[primary_span]
+    #[label]
+    pub span: Span,
+    pub found: Ty<'tcx>,
+    pub kind: ProcMacroKind,
+    pub expected_signature: &'static str,
+}
+
+#[derive(Diagnostic)]
+#[diag(passes_proc_macro_diff_arg_count)]
+pub(crate) struct ProcMacroDiffArguments {
+    #[primary_span]
+    #[label]
+    pub span: Span,
+    pub count: usize,
+    pub kind: ProcMacroKind,
+    pub expected_signature: &'static str,
+}
+
+#[derive(Diagnostic)]
+#[diag(passes_proc_macro_missing_args)]
+pub(crate) struct ProcMacroMissingArguments {
+    #[primary_span]
+    #[label]
+    pub span: Span,
+    pub expected_input_count: usize,
+    pub kind: ProcMacroKind,
+    pub expected_signature: &'static str,
+}
+
+#[derive(Diagnostic)]
+#[diag(passes_proc_macro_invalid_abi)]
+pub(crate) struct ProcMacroInvalidAbi {
+    #[primary_span]
+    pub span: Span,
+    pub abi: &'static str,
+}
+
+#[derive(Diagnostic)]
+#[diag(passes_proc_macro_unsafe)]
+pub(crate) struct ProcMacroUnsafe {
+    #[primary_span]
+    pub span: Span,
+}

compiler/rustc_span/src/symbol.rs

@@ -287,6 +287,7 @@ symbols! {
         Target,
         ToOwned,
         ToString,
+        TokenStream,
         Try,
         TryCaptureGeneric,
         TryCapturePrintable,

library/proc_macro/src/lib.rs

@@ -74,6 +74,7 @@ pub fn is_available() -> bool {
 ///
 /// This is both the input and output of `#[proc_macro]`, `#[proc_macro_attribute]`
 /// and `#[proc_macro_derive]` definitions.
+#[rustc_diagnostic_item = "TokenStream"]
 #[stable(feature = "proc_macro_lib", since = "1.15.0")]
 #[derive(Clone)]
 pub struct TokenStream(Option<bridge::client::TokenStream>);

tests/ui/proc-macro/allowed-signatures.rs

@@ -0,0 +1,26 @@
+// check-pass
+// force-host
+// no-prefer-dynamic
+
+#![crate_type = "proc-macro"]
+#![allow(private_in_public)]
+extern crate proc_macro;
+use proc_macro::TokenStream;
+
+#[proc_macro]
+pub fn foo<T>(t: T) -> TokenStream {
+  TokenStream::new()
+}
+
+trait Project {
+    type Assoc;
+}
+
+impl Project for () {
+    type Assoc = TokenStream;
+}
+
+#[proc_macro]
+pub fn uwu(_input: <() as Project>::Assoc) -> <() as Project>::Assoc {
+    TokenStream::new()
+}

tests/ui/proc-macro/proc-macro-abi.rs

@@ -0,0 +1,31 @@
+// force-host
+// no-prefer-dynamic
+
+#![crate_type = "proc-macro"]
+#![allow(warnings)]
+
+extern crate proc_macro;
+use proc_macro::TokenStream;
+
+#[proc_macro]
+pub extern "C" fn abi(a: TokenStream) -> TokenStream {
+    //~^ ERROR proc macro functions may not be `extern "C"`
+    a
+}
+
+#[proc_macro]
+pub extern "system" fn abi2(a: TokenStream) -> TokenStream {
+    //~^ ERROR proc macro functions may not be `extern "system"`
+    a
+}
+
+#[proc_macro]
+pub extern fn abi3(a: TokenStream) -> TokenStream {
+    //~^ ERROR proc macro functions may not be `extern "C"`
+    a
+}
+
+#[proc_macro]
+pub extern "Rust" fn abi4(a: TokenStream) -> TokenStream {
+    a
+}

tests/ui/proc-macro/proc-macro-abi.stderr

@@ -0,0 +1,20 @@
+error: proc macro functions may not be `extern "C"`
+  --> $DIR/proc-macro-abi.rs:11:1
+   |
+LL | pub extern "C" fn abi(a: TokenStream) -> TokenStream {
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: proc macro functions may not be `extern "system"`
+  --> $DIR/proc-macro-abi.rs:17:1
+   |
+LL | pub extern "system" fn abi2(a: TokenStream) -> TokenStream {
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: proc macro functions may not be `extern "C"`
+  --> $DIR/proc-macro-abi.rs:23:1
+   |
+LL | pub extern fn abi3(a: TokenStream) -> TokenStream {
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+error: aborting due to 3 previous errors
+