Port `#[custom_mir(..)]` to the new attribute system

2 files changed, 70 insertions, 4 deletions

diff --git a/compiler/rustc_passes/src/check_attr.rs b/compiler/rustc_passes/src/check_attr.rs
index 0e28c51e981..3c329d20700 100644
--- a/compiler/rustc_passes/src/check_attr.rs
+++ b/compiler/rustc_passes/src/check_attr.rs
@@ -18,7 +18,7 @@ use rustc_feature::{
     ACCEPTED_LANG_FEATURES, AttributeDuplicates, AttributeType, BUILTIN_ATTRIBUTE_MAP,
     BuiltinAttribute,
 };
-use rustc_hir::attrs::{AttributeKind, InlineAttr, ReprAttr};
+use rustc_hir::attrs::{AttributeKind, InlineAttr, MirDialect, MirPhase, ReprAttr};
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::LocalModDefId;
 use rustc_hir::intravisit::{self, Visitor};
@@ -197,6 +197,9 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
                 Attribute::Parsed(AttributeKind::MustUse { span, .. }) => {
                     self.check_must_use(hir_id, *span, target)
                 }
+                &Attribute::Parsed(AttributeKind::CustomMir(dialect, phase, attr_span)) => {
+                    self.check_custom_mir(dialect, phase, attr_span)
+                }
                 Attribute::Parsed(
                     AttributeKind::BodyStability { .. }
                     | AttributeKind::ConstStabilityIndirect
@@ -248,7 +251,6 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
                     | AttributeKind::Coroutine(..)
                     | AttributeKind::Linkage(..),
                 ) => { /* do nothing  */ }
-
                 Attribute::Unparsed(attr_item) => {
                     style = Some(attr_item.style);
                     match attr.path().as_slice() {
@@ -331,8 +333,7 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
                             | sym::panic_handler
                             | sym::lang
                             | sym::needs_allocator
-                            | sym::default_lib_allocator
-                            | sym::custom_mir,
+                            | sym::default_lib_allocator,
                             ..
                         ] => {}
                         [name, rest@..] => {
@@ -2113,6 +2114,48 @@ impl<'tcx> CheckAttrVisitor<'tcx> {
             self.dcx().emit_err(errors::ConstContinueAttr { attr_span, node_span });
         };
     }
+
+    fn check_custom_mir(
+        &self,
+        dialect: Option<(MirDialect, Span)>,
+        phase: Option<(MirPhase, Span)>,
+        attr_span: Span,
+    ) {
+        let Some((dialect, dialect_span)) = dialect else {
+            if let Some((_, phase_span)) = phase {
+                self.dcx()
+                    .emit_err(errors::CustomMirPhaseRequiresDialect { attr_span, phase_span });
+            }
+            return;
+        };
+
+        match dialect {
+            MirDialect::Analysis => {
+                if let Some((MirPhase::Optimized, phase_span)) = phase {
+                    self.dcx().emit_err(errors::CustomMirIncompatibleDialectAndPhase {
+                        dialect,
+                        phase: MirPhase::Optimized,
+                        attr_span,
+                        dialect_span,
+                        phase_span,
+                    });
+                }
+            }
+
+            MirDialect::Built => {
+                if let Some((phase, phase_span)) = phase {
+                    self.dcx().emit_err(errors::CustomMirIncompatibleDialectAndPhase {
+                        dialect,
+                        phase,
+                        attr_span,
+                        dialect_span,
+                        phase_span,
+                    });
+                }
+            }
+            MirDialect::Runtime => {}
+        }
+    }
 }
 
 impl<'tcx> Visitor<'tcx> for CheckAttrVisitor<'tcx> {
diff --git a/compiler/rustc_passes/src/errors.rs b/compiler/rustc_passes/src/errors.rs
index c5d5155d0e5..f8ecf10714a 100644
--- a/compiler/rustc_passes/src/errors.rs
+++ b/compiler/rustc_passes/src/errors.rs
@@ -7,6 +7,7 @@ use rustc_errors::{
     MultiSpan, Subdiagnostic,
 };
 use rustc_hir::Target;
+use rustc_hir::attrs::{MirDialect, MirPhase};
 use rustc_macros::{Diagnostic, LintDiagnostic, Subdiagnostic};
 use rustc_middle::ty::{MainDefinition, Ty};
 use rustc_span::{DUMMY_SP, Span, Symbol};
@@ -1570,3 +1571,25 @@ pub(crate) struct ReprAlignShouldBeAlign {
     pub span: Span,
     pub item: &'static str,
 }
+
+#[derive(Diagnostic)]
+#[diag(passes_custom_mir_phase_requires_dialect)]
+pub(crate) struct CustomMirPhaseRequiresDialect {
+    #[primary_span]
+    pub attr_span: Span,
+    #[label]
+    pub phase_span: Span,
+}
+
+#[derive(Diagnostic)]
+#[diag(passes_custom_mir_incompatible_dialect_and_phase)]
+pub(crate) struct CustomMirIncompatibleDialectAndPhase {
+    pub dialect: MirDialect,
+    pub phase: MirPhase,
+    #[primary_span]
+    pub attr_span: Span,
+    #[label]
+    pub dialect_span: Span,
+    #[label]
+    pub phase_span: Span,
+}