Port `#[coroutine]` to the new attribute system

Related to https://github.com/rust-lang/rust/issues/131229#issue-2565886367.

8 files changed, 42 insertions, 23 deletions

diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
index 1245d489754..fb42cfea30b 100644
--- a/compiler/rustc_ast_lowering/src/expr.rs
+++ b/compiler/rustc_ast_lowering/src/expr.rs
@@ -98,7 +98,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             }
 
             let expr_hir_id = self.lower_node_id(e.id);
-            self.lower_attrs(expr_hir_id, &e.attrs, e.span);
+            let attrs = self.lower_attrs(expr_hir_id, &e.attrs, e.span);
 
             let kind = match &e.kind {
                 ExprKind::Array(exprs) => hir::ExprKind::Array(self.lower_exprs(exprs)),
@@ -232,10 +232,10 @@ impl<'hir> LoweringContext<'_, 'hir> {
                         *fn_arg_span,
                     ),
                     None => self.lower_expr_closure(
+                        attrs,
                         binder,
                         *capture_clause,
                         e.id,
-                        expr_hir_id,
                         *constness,
                         *movability,
                         fn_decl,
@@ -1052,10 +1052,10 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
     fn lower_expr_closure(
         &mut self,
+        attrs: &[rustc_hir::Attribute],
         binder: &ClosureBinder,
         capture_clause: CaptureBy,
         closure_id: NodeId,
-        closure_hir_id: hir::HirId,
         constness: Const,
         movability: Movability,
         decl: &FnDecl,
@@ -1067,15 +1067,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let (binder_clause, generic_params) = self.lower_closure_binder(binder);
 
         let (body_id, closure_kind) = self.with_new_scopes(fn_decl_span, move |this| {
-            let mut coroutine_kind = if this
-                .attrs
-                .get(&closure_hir_id.local_id)
-                .is_some_and(|attrs| attrs.iter().any(|attr| attr.has_name(sym::coroutine)))
-            {
-                Some(hir::CoroutineKind::Coroutine(Movability::Movable))
-            } else {
-                None
-            };
+
+            let mut coroutine_kind = find_attr!(attrs, AttributeKind::Coroutine(_) => hir::CoroutineKind::Coroutine(Movability::Movable));
+
             // FIXME(contracts): Support contracts on closures?
             let body_id = this.lower_fn_body(decl, None, |this| {
                 this.coroutine_kind = coroutine_kind;
diff --git a/compiler/rustc_attr_parsing/src/attributes/body.rs b/compiler/rustc_attr_parsing/src/attributes/body.rs
new file mode 100644
index 00000000000..ab9330216f6
--- /dev/null
+++ b/compiler/rustc_attr_parsing/src/attributes/body.rs
@@ -0,0 +1,15 @@
+//! Attributes that can be found in function body.
+
+use rustc_hir::attrs::AttributeKind;
+use rustc_span::{Symbol, sym};
+
+use super::{NoArgsAttributeParser, OnDuplicate};
+use crate::context::Stage;
+
+pub(crate) struct CoroutineParser;
+
+impl<S: Stage> NoArgsAttributeParser<S> for CoroutineParser {
+    const PATH: &[Symbol] = &[sym::coroutine];
+    const ON_DUPLICATE: OnDuplicate<S> = OnDuplicate::Error;
+    const CREATE: fn(rustc_span::Span) -> AttributeKind = |span| AttributeKind::Coroutine(span);
+}
diff --git a/compiler/rustc_attr_parsing/src/attributes/mod.rs b/compiler/rustc_attr_parsing/src/attributes/mod.rs
index c574ef78bdf..f7946ade6d2 100644
--- a/compiler/rustc_attr_parsing/src/attributes/mod.rs
+++ b/compiler/rustc_attr_parsing/src/attributes/mod.rs
@@ -26,6 +26,7 @@ use crate::parser::ArgParser;
 use crate::session_diagnostics::UnusedMultiple;
 
 pub(crate) mod allow_unstable;
+pub(crate) mod body;
 pub(crate) mod cfg;
 pub(crate) mod cfg_old;
 pub(crate) mod codegen_attrs;
diff --git a/compiler/rustc_attr_parsing/src/context.rs b/compiler/rustc_attr_parsing/src/context.rs
index b51db9b4b9e..ec74f6fd1c6 100644
--- a/compiler/rustc_attr_parsing/src/context.rs
+++ b/compiler/rustc_attr_parsing/src/context.rs
@@ -16,6 +16,7 @@ use rustc_span::{DUMMY_SP, ErrorGuaranteed, Span, Symbol, sym};
 use crate::attributes::allow_unstable::{
     AllowConstFnUnstableParser, AllowInternalUnstableParser, UnstableFeatureBoundParser,
 };
+use crate::attributes::body::CoroutineParser;
 use crate::attributes::codegen_attrs::{
     ColdParser, CoverageParser, ExportNameParser, NakedParser, NoMangleParser, OptimizeParser,
     TargetFeatureParser, TrackCallerParser, UsedParser,
@@ -184,6 +185,7 @@ attribute_parsers!(
         Single<WithoutArgs<ConstContinueParser>>,
         Single<WithoutArgs<ConstStabilityIndirectParser>>,
         Single<WithoutArgs<ConstTraitParser>>,
+        Single<WithoutArgs<CoroutineParser>>,
         Single<WithoutArgs<DenyExplicitImplParser>>,
         Single<WithoutArgs<DoNotImplementViaObjectParser>>,
         Single<WithoutArgs<ExportStableParser>>,
diff --git a/compiler/rustc_hir/src/attrs/data_structures.rs b/compiler/rustc_hir/src/attrs/data_structures.rs
index 80618422b56..d9d2ec48948 100644
--- a/compiler/rustc_hir/src/attrs/data_structures.rs
+++ b/compiler/rustc_hir/src/attrs/data_structures.rs
@@ -297,6 +297,9 @@ pub enum AttributeKind {
     /// Represents `#[const_trait]`.
     ConstTrait(Span),
 
+    /// Represents `#[coroutine]`.
+    Coroutine(Span),
+
     /// Represents `#[coverage(..)]`.
     Coverage(Span, CoverageAttrKind),
 
diff --git a/compiler/rustc_hir/src/attrs/encode_cross_crate.rs b/compiler/rustc_hir/src/attrs/encode_cross_crate.rs
index 9644a597a31..b66e5bbeabe 100644
--- a/compiler/rustc_hir/src/attrs/encode_cross_crate.rs
+++ b/compiler/rustc_hir/src/attrs/encode_cross_crate.rs
@@ -28,6 +28,7 @@ impl AttributeKind {
             ConstStability { .. } => Yes,
             ConstStabilityIndirect => No,
             ConstTrait(..) => No,
+            Coroutine(..) => No,
             Coverage(..) => No,
             DenyExplicitImpl(..) => No,
             Deprecation { .. } => Yes,
diff --git a/compiler/rustc_passes/src/check_attr.rs b/compiler/rustc_passes/src/check_attr.rs
index 2663d5fe99c..890028d977d 100644
--- a/compiler/rustc_passes/src/check_attr.rs
+++ b/compiler/rustc_passes/src/check_attr.rs
@@ -324,6 +324,9 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
                 &Attribute::Parsed(AttributeKind::Coverage(attr_span, _)) => {
                     self.check_coverage(attr_span, span, target)
                 }
+                &Attribute::Parsed(AttributeKind::Coroutine(attr_span)) => {
+                    self.check_coroutine(attr_span, target)
+                }
                 Attribute::Unparsed(attr_item) => {
                     style = Some(attr_item.style);
                     match attr.path().as_slice() {
@@ -390,9 +393,6 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
                         [sym::autodiff_forward, ..] | [sym::autodiff_reverse, ..] => {
                             self.check_autodiff(hir_id, attr, span, target)
                         }
-                        [sym::coroutine, ..] => {
-                            self.check_coroutine(attr, target);
-                        }
                         [sym::linkage, ..] => self.check_linkage(attr, span, target),
                         [
                             // ok
@@ -2651,11 +2651,11 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
         }
     }
 
-    fn check_coroutine(&self, attr: &Attribute, target: Target) {
+    fn check_coroutine(&self, attr_span: Span, target: Target) {
         match target {
             Target::Closure => return,
             _ => {
-                self.dcx().emit_err(errors::CoroutineOnNonClosure { span: attr.span() });
+                self.dcx().emit_err(errors::CoroutineOnNonClosure { span: attr_span });
             }
         }
     }
diff --git a/tests/ui/attributes/malformed-attrs.stderr b/tests/ui/attributes/malformed-attrs.stderr
index 1b51075b4e8..e8ae4715398 100644
--- a/tests/ui/attributes/malformed-attrs.stderr
+++ b/tests/ui/attributes/malformed-attrs.stderr
@@ -55,12 +55,6 @@ error: malformed `patchable_function_entry` attribute input
 LL | #[patchable_function_entry]
    | ^^^^^^^^^^^^^^^^^^^^^^^^^^^ help: must be of the form: `#[patchable_function_entry(prefix_nops = m, entry_nops = n)]`
 
-error: malformed `coroutine` attribute input
-  --> $DIR/malformed-attrs.rs:108:5
-   |
-LL |     #[coroutine = 63] || {}
-   |     ^^^^^^^^^^^^^^^^^ help: must be of the form: `#[coroutine]`
-
 error: malformed `must_not_suspend` attribute input
   --> $DIR/malformed-attrs.rs:129:1
    |
@@ -436,6 +430,15 @@ LL | #[proc_macro = 18]
    | |            didn't expect any arguments here
    | help: must be of the form: `#[proc_macro]`
 
+error[E0565]: malformed `coroutine` attribute input
+  --> $DIR/malformed-attrs.rs:108:5
+   |
+LL |     #[coroutine = 63] || {}
+   |     ^^^^^^^^^^^^----^
+   |     |           |
+   |     |           didn't expect any arguments here
+   |     help: must be of the form: `#[coroutine]`
+
 error[E0565]: malformed `proc_macro_attribute` attribute input
   --> $DIR/malformed-attrs.rs:113:1
    |