Auto merge of #79326 - Aaron1011:fix/builtin-macro-stmt, r=petrochenkov

Always invoke statement attributes on the statement itself

This is preparation for PR #78296, which will require us to handle
statement items in addition to normal items.

Author: bors

Diff for: compiler/rustc_ast/src/token.rs

@@ -785,13 +785,20 @@ impl Nonterminal {
     /// See issue #73345 for more details.
     /// FIXME(#73933): Remove this eventually.
     pub fn pretty_printing_compatibility_hack(&self) -> bool {
-        if let NtItem(item) = self {
-            let name = item.ident.name;
-            if name == sym::ProceduralMasqueradeDummyType || name == sym::ProcMacroHack {
-                if let ast::ItemKind::Enum(enum_def, _) = &item.kind {
-                    if let [variant] = &*enum_def.variants {
-                        return variant.ident.name == sym::Input;
-                    }
+        let item = match self {
+            NtItem(item) => item,
+            NtStmt(stmt) => match &stmt.kind {
+                ast::StmtKind::Item(item) => item,
+                _ => return false,
+            },
+            _ => return false,
+        };
+
+        let name = item.ident.name;
+        if name == sym::ProceduralMasqueradeDummyType || name == sym::ProcMacroHack {
+            if let ast::ItemKind::Enum(enum_def, _) = &item.kind {
+                if let [variant] = &*enum_def.variants {
+                    return variant.ident.name == sym::Input;
                 }
             }
         }

Diff for: compiler/rustc_builtin_macros/src/deriving/mod.rs

@@ -54,7 +54,27 @@ impl MultiItemModifier for BuiltinDerive {
         // so we are doing it here in a centralized way.
         let span = ecx.with_def_site_ctxt(span);
         let mut items = Vec::new();
-        (self.0)(ecx, span, meta_item, &item, &mut |a| items.push(a));
+        match item {
+            Annotatable::Stmt(stmt) => {
+                if let ast::StmtKind::Item(item) = stmt.into_inner().kind {
+                    (self.0)(ecx, span, meta_item, &Annotatable::Item(item), &mut |a| {
+                        // Cannot use 'ecx.stmt_item' here, because we need to pass 'ecx'
+                        // to the function
+                        items.push(Annotatable::Stmt(P(ast::Stmt {
+                            id: ast::DUMMY_NODE_ID,
+                            kind: ast::StmtKind::Item(a.expect_item()),
+                            span,
+                            tokens: None,
+                        })));
+                    });
+                } else {
+                    unreachable!("should have already errored on non-item statement")
+                }
+            }
+            _ => {
+                (self.0)(ecx, span, meta_item, &item, &mut |a| items.push(a));
+            }
+        }
         ExpandResult::Ready(items)
     }
 }

Diff for: compiler/rustc_builtin_macros/src/global_allocator.rs

@@ -4,7 +4,7 @@ use rustc_ast::expand::allocator::{
     AllocatorKind, AllocatorMethod, AllocatorTy, ALLOCATOR_METHODS,
 };
 use rustc_ast::ptr::P;
-use rustc_ast::{self as ast, Attribute, Expr, FnHeader, FnSig, Generics, Param};
+use rustc_ast::{self as ast, Attribute, Expr, FnHeader, FnSig, Generics, Param, StmtKind};
 use rustc_ast::{ItemKind, Mutability, Stmt, Ty, TyKind, Unsafe};
 use rustc_expand::base::{Annotatable, ExtCtxt};
 use rustc_span::symbol::{kw, sym, Ident, Symbol};
@@ -14,14 +14,25 @@ pub fn expand(
     ecx: &mut ExtCtxt<'_>,
     _span: Span,
     meta_item: &ast::MetaItem,
-    item: Annotatable,
+    mut item: Annotatable,
 ) -> Vec<Annotatable> {
     check_builtin_macro_attribute(ecx, meta_item, sym::global_allocator);
 
     let not_static = |item: Annotatable| {
         ecx.sess.parse_sess.span_diagnostic.span_err(item.span(), "allocators must be statics");
         vec![item]
     };
+    let orig_item = item.clone();
+    let mut is_stmt = false;
+
+    // Allow using `#[global_allocator]` on an item statement
+    if let Annotatable::Stmt(stmt) = &item {
+        if let StmtKind::Item(item_) = &stmt.kind {
+            item = Annotatable::Item(item_.clone());
+            is_stmt = true;
+        }
+    }
+
     let item = match item {
         Annotatable::Item(item) => match item.kind {
             ItemKind::Static(..) => item,
@@ -41,9 +52,14 @@ pub fn expand(
     let const_ty = ecx.ty(span, TyKind::Tup(Vec::new()));
     let const_body = ecx.expr_block(ecx.block(span, stmts));
     let const_item = ecx.item_const(span, Ident::new(kw::Underscore, span), const_ty, const_body);
+    let const_item = if is_stmt {
+        Annotatable::Stmt(P(ecx.stmt_item(span, const_item)))
+    } else {
+        Annotatable::Item(const_item)
+    };
 
     // Return the original item and the new methods.
-    vec![Annotatable::Item(item), Annotatable::Item(const_item)]
+    vec![orig_item, const_item]
 }
 
 struct AllocFnFactory<'a, 'b> {

Diff for: compiler/rustc_builtin_macros/src/test.rs

@@ -4,6 +4,7 @@ use crate::util::check_builtin_macro_attribute;
 
 use rustc_ast as ast;
 use rustc_ast::attr;
+use rustc_ast::ptr::P;
 use rustc_ast_pretty::pprust;
 use rustc_expand::base::*;
 use rustc_session::Session;
@@ -78,8 +79,16 @@ pub fn expand_test_or_bench(
         return vec![];
     }
 
-    let item = match item {
-        Annotatable::Item(i) => i,
+    let (item, is_stmt) = match item {
+        Annotatable::Item(i) => (i, false),
+        Annotatable::Stmt(stmt) if matches!(stmt.kind, ast::StmtKind::Item(_)) => {
+            // FIXME: Use an 'if let' guard once they are implemented
+            if let ast::StmtKind::Item(i) = stmt.into_inner().kind {
+                (i, true)
+            } else {
+                unreachable!()
+            }
+        }
         other => {
             cx.struct_span_err(
                 other.span(),
@@ -304,14 +313,25 @@ pub fn expand_test_or_bench(
 
     tracing::debug!("synthetic test item:\n{}\n", pprust::item_to_string(&test_const));
 
-    vec![
-        // Access to libtest under a hygienic name
-        Annotatable::Item(test_extern),
-        // The generated test case
-        Annotatable::Item(test_const),
-        // The original item
-        Annotatable::Item(item),
-    ]
+    if is_stmt {
+        vec![
+            // Access to libtest under a hygienic name
+            Annotatable::Stmt(P(cx.stmt_item(sp, test_extern))),
+            // The generated test case
+            Annotatable::Stmt(P(cx.stmt_item(sp, test_const))),
+            // The original item
+            Annotatable::Stmt(P(cx.stmt_item(sp, item))),
+        ]
+    } else {
+        vec![
+            // Access to libtest under a hygienic name
+            Annotatable::Item(test_extern),
+            // The generated test case
+            Annotatable::Item(test_const),
+            // The original item
+            Annotatable::Item(item),
+        ]
+    }
 }
 
 fn item_path(mod_path: &[Ident], item_ident: &Ident) -> String {

Diff for: compiler/rustc_expand/src/base.rs

@@ -234,6 +234,15 @@ impl Annotatable {
 
     pub fn derive_allowed(&self) -> bool {
         match *self {
+            Annotatable::Stmt(ref stmt) => match stmt.kind {
+                ast::StmtKind::Item(ref item) => match item.kind {
+                    ast::ItemKind::Struct(..)
+                    | ast::ItemKind::Enum(..)
+                    | ast::ItemKind::Union(..) => true,
+                    _ => false,
+                },
+                _ => false,
+            },
             Annotatable::Item(ref item) => match item.kind {
                 ast::ItemKind::Struct(..) | ast::ItemKind::Enum(..) | ast::ItemKind::Union(..) => {
                     true

Diff for: compiler/rustc_expand/src/expand.rs

@@ -795,7 +795,14 @@ impl<'a, 'b> MacroExpander<'a, 'b> {
             | Annotatable::TraitItem(_)
             | Annotatable::ImplItem(_)
             | Annotatable::ForeignItem(_) => return,
-            Annotatable::Stmt(_) => "statements",
+            Annotatable::Stmt(stmt) => {
+                // Attributes are stable on item statements,
+                // but unstable on all other kinds of statements
+                if stmt.is_item() {
+                    return;
+                }
+                "statements"
+            }
             Annotatable::Expr(_) => "expressions",
             Annotatable::Arm(..)
             | Annotatable::Field(..)
@@ -1266,9 +1273,19 @@ impl<'a, 'b> MutVisitor for InvocationCollector<'a, 'b> {
 
         // we'll expand attributes on expressions separately
         if !stmt.is_expr() {
-            // FIXME: Handle custom attributes on statements (#15701).
-            let attr =
-                if stmt.is_item() { None } else { self.take_first_attr_no_derive(&mut stmt) };
+            let attr = if stmt.is_item() {
+                // FIXME: Implement proper token collection for statements
+                if let StmtKind::Item(item) = &mut stmt.kind {
+                    stmt.tokens = item.tokens.take()
+                } else {
+                    unreachable!()
+                };
+                self.take_first_attr(&mut stmt)
+            } else {
+                // Ignore derives on non-item statements for backwards compatibility.
+                // This will result in a unused attribute warning
+                self.take_first_attr_no_derive(&mut stmt)
+            };
 
             if let Some(attr) = attr {
                 return self

Diff for: compiler/rustc_expand/src/proc_macro.rs

@@ -1,6 +1,7 @@
 use crate::base::{self, *};
 use crate::proc_macro_server;
 
+use rustc_ast::ptr::P;
 use rustc_ast::token;
 use rustc_ast::tokenstream::{TokenStream, TokenTree};
 use rustc_ast::{self as ast, *};
@@ -74,8 +75,20 @@ impl MultiItemModifier for ProcMacroDerive {
         _meta_item: &ast::MetaItem,
         item: Annotatable,
     ) -> ExpandResult<Vec<Annotatable>, Annotatable> {
+        // We need special handling for statement items
+        // (e.g. `fn foo() { #[derive(Debug)] struct Bar; }`)
+        let mut is_stmt = false;
         let item = match item {
             Annotatable::Item(item) => token::NtItem(item),
+            Annotatable::Stmt(stmt) => {
+                is_stmt = true;
+                assert!(stmt.is_item());
+
+                // A proc macro can't observe the fact that we're passing
+                // them an `NtStmt` - it can only see the underlying tokens
+                // of the wrapped item
+                token::NtStmt(stmt.into_inner())
+            }
             _ => unreachable!(),
         };
         let input = if item.pretty_printing_compatibility_hack() {
@@ -106,7 +119,13 @@ impl MultiItemModifier for ProcMacroDerive {
         loop {
             match parser.parse_item() {
                 Ok(None) => break,
-                Ok(Some(item)) => items.push(Annotatable::Item(item)),
+                Ok(Some(item)) => {
+                    if is_stmt {
+                        items.push(Annotatable::Stmt(P(ecx.stmt_item(span, item))));
+                    } else {
+                        items.push(Annotatable::Item(item));
+                    }
+                }
                 Err(mut err) => {
                     err.emit();
                     break;

Diff for: src/test/ui/proc-macro/allowed-attr-stmt-expr.rs

@@ -0,0 +1,59 @@
+// aux-build:attr-stmt-expr.rs
+// aux-build:test-macros.rs
+// compile-flags: -Z span-debug
+// check-pass
+
+#![feature(proc_macro_hygiene)]
+#![feature(stmt_expr_attributes)]
+#![feature(rustc_attrs)]
+#![allow(dead_code)]
+
+#![no_std] // Don't load unnecessary hygiene information from std
+extern crate std;
+
+extern crate attr_stmt_expr;
+extern crate test_macros;
+use attr_stmt_expr::{expect_let, expect_print_stmt, expect_expr, expect_print_expr};
+use test_macros::print_attr;
+use std::println;
+
+fn print_str(string: &'static str) {
+    // macros are handled a bit differently
+    #[expect_print_expr]
+    println!("{}", string)
+}
+
+macro_rules! make_stmt {
+    ($stmt:stmt) => {
+        $stmt
+    }
+}
+
+macro_rules! second_make_stmt {
+    ($stmt:stmt) => {
+        make_stmt!($stmt);
+    }
+}
+
+
+fn main() {
+    make_stmt!(struct Foo {});
+
+    #[print_attr]
+    #[expect_let]
+    let string = "Hello, world!";
+
+    #[print_attr]
+    #[expect_print_stmt]
+    println!("{}", string);
+
+    #[print_attr]
+    second_make_stmt!(#[allow(dead_code)] struct Bar {});
+
+    #[print_attr]
+    #[rustc_dummy]
+    struct Other {};
+
+    #[expect_expr]
+    print_str("string")
+}