Rollup merge of #103464 - JakobDegen:mir-parsing, r=oli-obk

Add support for custom mir

This implements rust-lang/compiler-team#564 . Details about the design, motivation, etc. can be found in there.

r? ```@oli-obk```

compiler/rustc_borrowck/src/lib.rs

@@ -18,6 +18,7 @@ extern crate tracing;
 
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
 use rustc_data_structures::graph::dominators::Dominators;
+use rustc_data_structures::vec_map::VecMap;
 use rustc_errors::{Diagnostic, DiagnosticBuilder, ErrorGuaranteed};
 use rustc_hir as hir;
 use rustc_hir::def_id::LocalDefId;
@@ -129,6 +130,19 @@ fn mir_borrowck<'tcx>(
 ) -> &'tcx BorrowCheckResult<'tcx> {
     let (input_body, promoted) = tcx.mir_promoted(def);
     debug!("run query mir_borrowck: {}", tcx.def_path_str(def.did.to_def_id()));
+
+    if input_body.borrow().should_skip() {
+        debug!("Skipping borrowck because of injected body");
+        // Let's make up a borrowck result! Fun times!
+        let result = BorrowCheckResult {
+            concrete_opaque_types: VecMap::new(),
+            closure_requirements: None,
+            used_mut_upvars: SmallVec::new(),
+            tainted_by_errors: None,
+        };
+        return tcx.arena.alloc(result);
+    }
+
     let hir_owner = tcx.hir().local_def_id_to_hir_id(def.did).owner;
 
     let infcx =

compiler/rustc_feature/src/active.rs

@@ -152,6 +152,8 @@ declare_features! (
     (active, anonymous_lifetime_in_impl_trait, "1.63.0", None, None),
     /// Allows identifying the `compiler_builtins` crate.
     (active, compiler_builtins, "1.13.0", None, None),
+    /// Allows writing custom MIR
+    (active, custom_mir, "1.65.0", None, None),
     /// Outputs useful `assert!` messages
     (active, generic_assert, "1.63.0", None, None),
     /// Allows using the `rust-intrinsic`'s "ABI".

compiler/rustc_feature/src/builtin_attrs.rs

@@ -810,6 +810,10 @@ pub const BUILTIN_ATTRIBUTES: &[BuiltinAttribute] = &[
     rustc_attr!(TEST, rustc_polymorphize_error, Normal, template!(Word), WarnFollowing),
     rustc_attr!(TEST, rustc_def_path, Normal, template!(Word), WarnFollowing),
     rustc_attr!(TEST, rustc_mir, Normal, template!(List: "arg1, arg2, ..."), DuplicatesOk),
+    gated!(
+        custom_mir, Normal, template!(List: r#"dialect = "...", phase = "...""#),
+        ErrorFollowing, "the `#[custom_mir]` attribute is just used for the Rust test suite",
+    ),
     rustc_attr!(TEST, rustc_dump_program_clauses, Normal, template!(Word), WarnFollowing),
     rustc_attr!(TEST, rustc_dump_env_program_clauses, Normal, template!(Word), WarnFollowing),
     rustc_attr!(TEST, rustc_object_lifetime_default, Normal, template!(Word), WarnFollowing),

compiler/rustc_middle/src/mir/mod.rs

@@ -138,6 +138,48 @@ impl MirPhase {
             }
         }
     }
+
+    /// Parses an `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)),
+            _ => panic!("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,
+            _ => panic!("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,
+            _ => panic!("Unknown runtime phase {}", phase),
+        }
+    }
 }
 
 impl Display for MirPhase {
@@ -293,6 +335,13 @@ pub struct Body<'tcx> {
     /// potentially allow things like `[u8; std::mem::size_of::<T>() * 0]` due to this.
     pub is_polymorphic: bool,
 
+    /// The phase at which this MIR should be "injected" into the compilation process.
+    ///
+    /// Everything that comes before this `MirPhase` should be skipped.
+    ///
+    /// This is only `Some` if the function that this body comes from was annotated with `rustc_custom_mir`.
+    pub injection_phase: Option<MirPhase>,
+
     pub tainted_by_errors: Option<ErrorGuaranteed>,
 }
 
@@ -339,6 +388,7 @@ impl<'tcx> Body<'tcx> {
             span,
             required_consts: Vec::new(),
             is_polymorphic: false,
+            injection_phase: None,
             tainted_by_errors,
         };
         body.is_polymorphic = body.has_non_region_param();
@@ -366,6 +416,7 @@ impl<'tcx> Body<'tcx> {
             required_consts: Vec::new(),
             var_debug_info: Vec::new(),
             is_polymorphic: false,
+            injection_phase: None,
             tainted_by_errors: None,
         };
         body.is_polymorphic = body.has_non_region_param();
@@ -508,6 +559,14 @@ impl<'tcx> Body<'tcx> {
     pub fn generator_kind(&self) -> Option<GeneratorKind> {
         self.generator.as_ref().map(|generator| generator.generator_kind)
     }
+
+    #[inline]
+    pub fn should_skip(&self) -> bool {
+        let Some(injection_phase) = self.injection_phase else {
+            return false;
+        };
+        injection_phase > self.phase
+    }
 }
 
 #[derive(Copy, Clone, PartialEq, Eq, Debug, TyEncodable, TyDecodable, HashStable)]

compiler/rustc_mir_build/src/build/custom/mod.rs

@@ -0,0 +1,155 @@
+//! Provides the implementation of the `custom_mir` attribute.
+//!
+//! Up until MIR building, this attribute has absolutely no effect. The `mir!` macro is a normal
+//! decl macro that expands like any other, and the code goes through parsing, name resolution and
+//! type checking like all other code. In MIR building we finally detect whether this attribute is
+//! present, and if so we branch off into this module, which implements the attribute by
+//! implementing a custom lowering from THIR to MIR.
+//!
+//! The result of this lowering is returned "normally" from the `mir_built` query, with the only
+//! notable difference being that the `injected` field in the body is set. Various components of the
+//! MIR pipeline, like borrowck and the pass manager will then consult this field (via
+//! `body.should_skip()`) to skip the parts of the MIR pipeline that precede the MIR phase the user
+//! specified.
+//!
+//! This file defines the general framework for the custom parsing. The parsing for all the
+//! "top-level" constructs can be found in the `parse` submodule, while the parsing for statements,
+//! terminators, and everything below can be found in the `parse::instruction` submodule.
+//!
+
+use rustc_ast::Attribute;
+use rustc_data_structures::fx::FxHashMap;
+use rustc_hir::def_id::DefId;
+use rustc_index::vec::IndexVec;
+use rustc_middle::{
+    mir::*,
+    thir::*,
+    ty::{Ty, TyCtxt},
+};
+use rustc_span::Span;
+
+mod parse;
+
+pub(super) fn build_custom_mir<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    did: DefId,
+    thir: &Thir<'tcx>,
+    expr: ExprId,
+    params: &IndexVec<ParamId, Param<'tcx>>,
+    return_ty: Ty<'tcx>,
+    return_ty_span: Span,
+    span: Span,
+    attr: &Attribute,
+) -> Body<'tcx> {
+    let mut body = Body {
+        basic_blocks: BasicBlocks::new(IndexVec::new()),
+        source: MirSource::item(did),
+        phase: MirPhase::Built,
+        source_scopes: IndexVec::new(),
+        generator: None,
+        local_decls: LocalDecls::new(),
+        user_type_annotations: IndexVec::new(),
+        arg_count: params.len(),
+        spread_arg: None,
+        var_debug_info: Vec::new(),
+        span,
+        required_consts: Vec::new(),
+        is_polymorphic: false,
+        tainted_by_errors: None,
+        injection_phase: None,
+        pass_count: 1,
+    };
+
+    body.local_decls.push(LocalDecl::new(return_ty, return_ty_span));
+    body.basic_blocks_mut().push(BasicBlockData::new(None));
+    body.source_scopes.push(SourceScopeData {
+        span,
+        parent_scope: None,
+        inlined: None,
+        inlined_parent_scope: None,
+        local_data: ClearCrossCrate::Clear,
+    });
+    body.injection_phase = Some(parse_attribute(attr));
+
+    let mut pctxt = ParseCtxt {
+        tcx,
+        thir,
+        source_info: SourceInfo { span, scope: OUTERMOST_SOURCE_SCOPE },
+        body: &mut body,
+        local_map: FxHashMap::default(),
+        block_map: FxHashMap::default(),
+    };
+
+    let res = (|| {
+        pctxt.parse_args(&params)?;
+        pctxt.parse_body(expr)
+    })();
+    if let Err(err) = res {
+        tcx.sess.diagnostic().span_fatal(
+            err.span,
+            format!("Could not parse {}, found: {:?}", err.expected, err.item_description),
+        )
+    }
+
+    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;
+
+    for nested in meta_items {
+        let name = nested.name_or_empty();
+        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 => {
+                panic!("Unexpected key {}", other);
+            }
+        }
+    }
+
+    let Some(dialect) = dialect else {
+        assert!(phase.is_none());
+        return MirPhase::Built;
+    };
+
+    MirPhase::parse(dialect, phase)
+}
+
+struct ParseCtxt<'tcx, 'body> {
+    tcx: TyCtxt<'tcx>,
+    thir: &'body Thir<'tcx>,
+    source_info: SourceInfo,
+
+    body: &'body mut Body<'tcx>,
+    local_map: FxHashMap<LocalVarId, Local>,
+    block_map: FxHashMap<LocalVarId, BasicBlock>,
+}
+
+struct ParseError {
+    span: Span,
+    item_description: String,
+    expected: String,
+}
+
+impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
+    fn expr_error(&self, expr: ExprId, expected: &'static str) -> ParseError {
+        let expr = &self.thir[expr];
+        ParseError {
+            span: expr.span,
+            item_description: format!("{:?}", expr.kind),
+            expected: expected.to_string(),
+        }
+    }
+}
+
+type PResult<T> = Result<T, ParseError>;