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

compiler/rustc_attr_parsing/src/attributes/mod.rs

@@ -43,6 +43,7 @@ pub(crate) mod no_implicit_prelude;
 pub(crate) mod non_exhaustive;
 pub(crate) mod path;
 pub(crate) mod proc_macro_attrs;
+pub(crate) mod prototype;
 pub(crate) mod repr;
 pub(crate) mod rustc_internal;
 pub(crate) mod semantics;

compiler/rustc_attr_parsing/src/attributes/prototype.rs

@@ -0,0 +1,136 @@
+//! Attributes that are only used on function prototypes.
+
+use rustc_feature::{AttributeTemplate, template};
+use rustc_hir::attrs::{AttributeKind, MirDialect, MirPhase};
+use rustc_span::{Span, Symbol, sym};
+
+use super::{AttributeOrder, OnDuplicate};
+use crate::attributes::SingleAttributeParser;
+use crate::context::{AcceptContext, Stage};
+use crate::parser::ArgParser;
+
+pub(crate) struct CustomMirParser;
+
+impl<S: Stage> SingleAttributeParser<S> for CustomMirParser {
+    const PATH: &[rustc_span::Symbol] = &[sym::custom_mir];
+
+    const ATTRIBUTE_ORDER: AttributeOrder = AttributeOrder::KeepOutermost;
+
+    const ON_DUPLICATE: OnDuplicate<S> = OnDuplicate::Error;
+
+    const TEMPLATE: AttributeTemplate = template!(List: &[r#"dialect = "...", phase = "...""#]);
+
+    fn convert(cx: &mut AcceptContext<'_, '_, S>, args: &ArgParser<'_>) -> Option<AttributeKind> {
+        let Some(list) = args.list() else {
+            cx.expected_list(cx.attr_span);
+            return None;
+        };
+
+        let mut dialect = None;
+        let mut phase = None;
+        let mut failed = false;
+
+        for item in list.mixed() {
+            let Some(meta_item) = item.meta_item() else {
+                cx.expected_name_value(item.span(), None);
+                failed = true;
+                break;
+            };
+
+            if let Some(arg) = meta_item.word_is(sym::dialect) {
+                extract_value(cx, sym::dialect, arg, meta_item.span(), &mut dialect, &mut failed);
+            } else if let Some(arg) = meta_item.word_is(sym::phase) {
+                extract_value(cx, sym::phase, arg, meta_item.span(), &mut phase, &mut failed);
+            } else if let Some(word) = meta_item.path().word() {
+                let word = word.to_string();
+                cx.unknown_key(meta_item.span(), word, &["dialect", "phase"]);
+                failed = true;
+            } else {
+                cx.expected_name_value(meta_item.span(), None);
+                failed = true;
+            };
+        }
+
+        let dialect = parse_dialect(cx, dialect, &mut failed);
+        let phase = parse_phase(cx, phase, &mut failed);
+
+        if failed {
+            return None;
+        }
+
+        Some(AttributeKind::CustomMir(dialect, phase, cx.attr_span))
+    }
+}
+
+fn extract_value<S: Stage>(
+    cx: &mut AcceptContext<'_, '_, S>,
+    key: Symbol,
+    arg: &ArgParser<'_>,
+    span: Span,
+    out_val: &mut Option<(Symbol, Span)>,
+    failed: &mut bool,
+) {
+    if out_val.is_some() {
+        cx.duplicate_key(span, key);
+        *failed = true;
+        return;
+    }
+
+    let Some(val) = arg.name_value() else {
+        cx.expected_single_argument(arg.span().unwrap_or(span));
+        *failed = true;
+        return;
+    };
+
+    let Some(value_sym) = val.value_as_str() else {
+        cx.expected_string_literal(val.value_span, Some(val.value_as_lit()));
+        *failed = true;
+        return;
+    };
+
+    *out_val = Some((value_sym, val.value_span));
+}
+
+fn parse_dialect<S: Stage>(
+    cx: &mut AcceptContext<'_, '_, S>,
+    dialect: Option<(Symbol, Span)>,
+    failed: &mut bool,
+) -> Option<(MirDialect, Span)> {
+    let (dialect, span) = dialect?;
+
+    let dialect = match dialect {
+        sym::analysis => MirDialect::Analysis,
+        sym::built => MirDialect::Built,
+        sym::runtime => MirDialect::Runtime,
+
+        _ => {
+            cx.expected_specific_argument(span, vec!["analysis", "built", "runtime"]);
+            *failed = true;
+            return None;
+        }
+    };
+
+    Some((dialect, span))
+}
+
+fn parse_phase<S: Stage>(
+    cx: &mut AcceptContext<'_, '_, S>,
+    phase: Option<(Symbol, Span)>,
+    failed: &mut bool,
+) -> Option<(MirPhase, Span)> {
+    let (phase, span) = phase?;
+
+    let phase = match phase {
+        sym::initial => MirPhase::Initial,
+        sym::post_cleanup => MirPhase::PostCleanup,
+        sym::optimized => MirPhase::Optimized,
+
+        _ => {
+            cx.expected_specific_argument(span, vec!["initial", "post-cleanup", "optimized"]);
+            *failed = true;
+            return None;
+        }
+    };
+
+    Some((phase, span))
+}

compiler/rustc_attr_parsing/src/context.rs

@@ -43,6 +43,7 @@ use crate::attributes::path::PathParser as PathAttributeParser;
 use crate::attributes::proc_macro_attrs::{
     ProcMacroAttributeParser, ProcMacroDeriveParser, ProcMacroParser, RustcBuiltinMacroParser,
 };
+use crate::attributes::prototype::CustomMirParser;
 use crate::attributes::repr::{AlignParser, ReprParser};
 use crate::attributes::rustc_internal::{
     RustcLayoutScalarValidRangeEnd, RustcLayoutScalarValidRangeStart,
@@ -159,6 +160,7 @@ attribute_parsers!(
 
         // tidy-alphabetical-start
         Single<CoverageParser>,
+        Single<CustomMirParser>,
         Single<DeprecationParser>,
         Single<DummyParser>,
         Single<ExportNameParser>,

compiler/rustc_errors/src/diagnostic_impls.rs

@@ -9,6 +9,7 @@ use rustc_abi::TargetDataLayoutErrors;
 use rustc_ast::util::parser::ExprPrecedence;
 use rustc_ast_pretty::pprust;
 use rustc_hir::RustcVersion;
+use rustc_hir::attrs::{MirDialect, MirPhase};
 use rustc_macros::Subdiagnostic;
 use rustc_span::edition::Edition;
 use rustc_span::{Ident, MacroRulesNormalizedIdent, Span, Symbol};
@@ -312,6 +313,28 @@ impl IntoDiagArg for ExprPrecedence {
     }
 }
 
+impl IntoDiagArg for MirDialect {
+    fn into_diag_arg(self, _path: &mut Option<PathBuf>) -> DiagArgValue {
+        let arg = match self {
+            MirDialect::Analysis => "analysis",
+            MirDialect::Built => "built",
+            MirDialect::Runtime => "runtime",
+        };
+        DiagArgValue::Str(Cow::Borrowed(arg))
+    }
+}
+
+impl IntoDiagArg for MirPhase {
+    fn into_diag_arg(self, _path: &mut Option<PathBuf>) -> DiagArgValue {
+        let arg = match self {
+            MirPhase::Initial => "initial",
+            MirPhase::PostCleanup => "post-cleanup",
+            MirPhase::Optimized => "optimized",
+        };
+        DiagArgValue::Str(Cow::Borrowed(arg))
+    }
+}
+
 #[derive(Clone)]
 pub struct DiagSymbolList<S = Symbol>(Vec<S>);
 

compiler/rustc_hir/src/attrs/data_structures.rs

@@ -187,6 +187,22 @@ pub enum CfgEntry {
     Version(Option<RustcVersion>, Span),
 }
 
+#[derive(Clone, Copy, Decodable, Debug, Encodable, PartialEq)]
+#[derive(HashStable_Generic, PrintAttribute)]
+pub enum MirDialect {
+    Analysis,
+    Built,
+    Runtime,
+}
+
+#[derive(Clone, Copy, Decodable, Debug, Encodable, PartialEq)]
+#[derive(HashStable_Generic, PrintAttribute)]
+pub enum MirPhase {
+    Initial,
+    PostCleanup,
+    Optimized,
+}
+
 /// Represents parsed *built-in* inert attributes.
 ///
 /// ## Overview
@@ -306,6 +322,9 @@ pub enum AttributeKind {
     /// Represents `#[coverage(..)]`.
     Coverage(Span, CoverageAttrKind),
 
+    /// Represents `#[custom_mir]`.
+    CustomMir(Option<(MirDialect, Span)>, Option<(MirPhase, Span)>, Span),
+
     ///Represents `#[rustc_deny_explicit_impl]`.
     DenyExplicitImpl(Span),
 

compiler/rustc_hir/src/attrs/encode_cross_crate.rs

@@ -31,6 +31,7 @@ impl AttributeKind {
             ConstTrait(..) => No,
             Coroutine(..) => No,
             Coverage(..) => No,
+            CustomMir(_, _, _) => Yes,
             DenyExplicitImpl(..) => No,
             Deprecation { .. } => Yes,
             DoNotImplementViaObject(..) => No,

compiler/rustc_middle/src/mir/mod.rs

@@ -115,48 +115,6 @@ impl MirPhase {
             MirPhase::Runtime(runtime_phase) => (3, 1 + runtime_phase as usize),
         }
     }
-
-    /// Parses a `MirPhase` from a pair of strings. Panics if this isn't possible for any reason.
-    pub fn parse(dialect: String, phase: Option<String>) -> Self {
-        match &*dialect.to_ascii_lowercase() {
-            "built" => {
-                assert!(phase.is_none(), "Cannot specify a phase for `Built` MIR");
-                MirPhase::Built
-            }
-            "analysis" => Self::Analysis(AnalysisPhase::parse(phase)),
-            "runtime" => Self::Runtime(RuntimePhase::parse(phase)),
-            _ => bug!("Unknown MIR dialect: '{}'", dialect),
-        }
-    }
-}
-
-impl AnalysisPhase {
-    pub fn parse(phase: Option<String>) -> Self {
-        let Some(phase) = phase else {
-            return Self::Initial;
-        };
-
-        match &*phase.to_ascii_lowercase() {
-            "initial" => Self::Initial,
-            "post_cleanup" | "post-cleanup" | "postcleanup" => Self::PostCleanup,
-            _ => bug!("Unknown analysis phase: '{}'", phase),
-        }
-    }
-}
-
-impl RuntimePhase {
-    pub fn parse(phase: Option<String>) -> Self {
-        let Some(phase) = phase else {
-            return Self::Initial;
-        };
-
-        match &*phase.to_ascii_lowercase() {
-            "initial" => Self::Initial,
-            "post_cleanup" | "post-cleanup" | "postcleanup" => Self::PostCleanup,
-            "optimized" => Self::Optimized,
-            _ => bug!("Unknown runtime phase: '{}'", phase),
-        }
-    }
 }
 
 /// Where a specific `mir::Body` comes from.

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