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

2 files changed, 37 insertions, 37 deletions

diff --git a/compiler/rustc_mir_build/src/builder/custom/mod.rs b/compiler/rustc_mir_build/src/builder/custom/mod.rs
index 902a6e7f115..792ad6d782c 100644
--- a/compiler/rustc_mir_build/src/builder/custom/mod.rs
+++ b/compiler/rustc_mir_build/src/builder/custom/mod.rs
@@ -19,10 +19,10 @@
 
 use rustc_data_structures::fx::FxHashMap;
 use rustc_hir::def_id::DefId;
-use rustc_hir::{Attribute, HirId};
+use rustc_hir::{HirId, attrs};
 use rustc_index::{IndexSlice, IndexVec};
+use rustc_middle::bug;
 use rustc_middle::mir::*;
-use rustc_middle::span_bug;
 use rustc_middle::thir::*;
 use rustc_middle::ty::{self, Ty, TyCtxt};
 use rustc_span::Span;
@@ -39,7 +39,8 @@ pub(super) fn build_custom_mir<'tcx>(
     return_ty: Ty<'tcx>,
     return_ty_span: Span,
     span: Span,
-    attr: &Attribute,
+    dialect: Option<attrs::MirDialect>,
+    phase: Option<attrs::MirPhase>,
 ) -> Body<'tcx> {
     let mut body = Body {
         basic_blocks: BasicBlocks::new(IndexVec::new()),
@@ -72,7 +73,7 @@ pub(super) fn build_custom_mir<'tcx>(
         inlined_parent_scope: None,
         local_data: ClearCrossCrate::Set(SourceScopeLocalData { lint_root: hir_id }),
     });
-    body.injection_phase = Some(parse_attribute(attr));
+    body.injection_phase = Some(parse_attribute(dialect, phase));
 
     let mut pctxt = ParseCtxt {
         tcx,
@@ -98,40 +99,38 @@ pub(super) fn build_custom_mir<'tcx>(
     body
 }
 
-fn parse_attribute(attr: &Attribute) -> MirPhase {
-    let meta_items = attr.meta_item_list().unwrap();
-    let mut dialect: Option<String> = None;
-    let mut phase: Option<String> = None;
-
-    // Not handling errors properly for this internal attribute; will just abort on errors.
-    for nested in meta_items {
-        let name = nested.name().unwrap();
-        let value = nested.value_str().unwrap().as_str().to_string();
-        match name.as_str() {
-            "dialect" => {
-                assert!(dialect.is_none());
-                dialect = Some(value);
-            }
-            "phase" => {
-                assert!(phase.is_none());
-                phase = Some(value);
-            }
-            other => {
-                span_bug!(
-                    nested.span(),
-                    "Unexpected key while parsing custom_mir attribute: '{}'",
-                    other
-                );
-            }
-        }
-    }
-
+/// Turns the arguments passed to `#[custom_mir(..)]` into a proper
+/// [`MirPhase`]. Panics if this isn't possible for any reason.
+fn parse_attribute(dialect: Option<attrs::MirDialect>, phase: Option<attrs::MirPhase>) -> MirPhase {
     let Some(dialect) = dialect else {
+        // Caught during attribute checking.
         assert!(phase.is_none());
         return MirPhase::Built;
     };
 
-    MirPhase::parse(dialect, phase)
+    match dialect {
+        attrs::MirDialect::Built => {
+            // Caught during attribute checking.
+            assert!(phase.is_none(), "Cannot specify a phase for `Built` MIR");
+            MirPhase::Built
+        }
+        attrs::MirDialect::Analysis => match phase {
+            None | Some(attrs::MirPhase::Initial) => MirPhase::Analysis(AnalysisPhase::Initial),
+
+            Some(attrs::MirPhase::PostCleanup) => MirPhase::Analysis(AnalysisPhase::PostCleanup),
+
+            Some(attrs::MirPhase::Optimized) => {
+                // Caught during attribute checking.
+                bug!("`optimized` dialect is not compatible with the `analysis` dialect")
+            }
+        },
+
+        attrs::MirDialect::Runtime => match phase {
+            None | Some(attrs::MirPhase::Initial) => MirPhase::Runtime(RuntimePhase::Initial),
+            Some(attrs::MirPhase::PostCleanup) => MirPhase::Runtime(RuntimePhase::PostCleanup),
+            Some(attrs::MirPhase::Optimized) => MirPhase::Runtime(RuntimePhase::Optimized),
+        },
+    }
 }
 
 struct ParseCtxt<'a, 'tcx> {
diff --git a/compiler/rustc_mir_build/src/builder/mod.rs b/compiler/rustc_mir_build/src/builder/mod.rs
index 855cd2f3bc0..9570760f943 100644
--- a/compiler/rustc_mir_build/src/builder/mod.rs
+++ b/compiler/rustc_mir_build/src/builder/mod.rs
@@ -11,9 +11,10 @@ use rustc_ast::attr;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sorted_map::SortedIndexMultiMap;
 use rustc_errors::ErrorGuaranteed;
+use rustc_hir::attrs::AttributeKind;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LocalDefId};
-use rustc_hir::{self as hir, BindingMode, ByRef, HirId, ItemLocalId, Node};
+use rustc_hir::{self as hir, BindingMode, ByRef, HirId, ItemLocalId, Node, find_attr};
 use rustc_index::bit_set::GrowableBitSet;
 use rustc_index::{Idx, IndexSlice, IndexVec};
 use rustc_infer::infer::{InferCtxt, TyCtxtInferExt};
@@ -479,8 +480,7 @@ fn construct_fn<'tcx>(
         ty => span_bug!(span_with_body, "unexpected type of body: {ty:?}"),
     };
 
-    if let Some(custom_mir_attr) =
-        tcx.hir_attrs(fn_id).iter().find(|attr| attr.has_name(sym::custom_mir))
+    if let Some((dialect, phase)) = find_attr!(tcx.hir_attrs(fn_id), AttributeKind::CustomMir(dialect, phase, _) => (dialect, phase))
     {
         return custom::build_custom_mir(
             tcx,
@@ -492,7 +492,8 @@ fn construct_fn<'tcx>(
             return_ty,
             return_ty_span,
             span_with_body,
-            custom_mir_attr,
+            dialect.as_ref().map(|(d, _)| *d),
+            phase.as_ref().map(|(p, _)| *p),
         );
     }