Handle Annotatable::Stmt in some builtin macros

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

Author: Aaron1011

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)
     }
 }

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> {

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 {