Auto merge of #105356 - JakobDegen:more-custom-mir, r=oli-obk

Custom MIR: Many more improvements

Commits are each atomic changes, best reviewed one at a time, with the exception that the last commit includes all the documentation.

### First commit

Unsafetyck was not correctly disabled before for `dialect = "built"` custom MIR. This is fixed and a regression test is added.

### Second commit

Implements `Discriminant`, `SetDiscriminant`, and `SwitchInt`.

### Third commit

Implements indexing, field, and variant projections.

### Fourth commit

Documents the previous commits and everything else.

There is some amount of weirdness here due to having to beat Rust syntax into cooperating with MIR concepts, but it hopefully should not be too much. All of it is documented.

r? `@oli-obk`

compiler/rustc_middle/src/mir/mod.rs

@@ -533,6 +533,11 @@ impl<'tcx> Body<'tcx> {
         };
         injection_phase > self.phase
     }
+
+    #[inline]
+    pub fn is_custom_mir(&self) -> bool {
+        self.injection_phase.is_some()
+    }
 }
 
 #[derive(Copy, Clone, PartialEq, Eq, Debug, TyEncodable, TyDecodable, HashStable)]

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

@@ -25,7 +25,7 @@ use rustc_index::vec::IndexVec;
 use rustc_middle::{
     mir::*,
     thir::*,
-    ty::{Ty, TyCtxt},
+    ty::{ParamEnv, Ty, TyCtxt},
 };
 use rustc_span::Span;
 
@@ -78,6 +78,7 @@ pub(super) fn build_custom_mir<'tcx>(
 
     let mut pctxt = ParseCtxt {
         tcx,
+        param_env: tcx.param_env(did),
         thir,
         source_scope: OUTERMOST_SOURCE_SCOPE,
         body: &mut body,
@@ -132,6 +133,7 @@ fn parse_attribute(attr: &Attribute) -> MirPhase {
 
 struct ParseCtxt<'tcx, 'body> {
     tcx: TyCtxt<'tcx>,
+    param_env: ParamEnv<'tcx>,
     thir: &'body Thir<'tcx>,
     source_scope: SourceScope,
 

compiler/rustc_mir_build/src/build/custom/parse/instruction.rs

@@ -1,5 +1,11 @@
 use rustc_middle::mir::interpret::{ConstValue, Scalar};
+use rustc_middle::mir::tcx::PlaceTy;
 use rustc_middle::{mir::*, thir::*, ty};
+use rustc_span::Span;
+use rustc_target::abi::VariantIdx;
+
+use crate::build::custom::ParseError;
+use crate::build::expr::as_constant::as_constant_inner;
 
 use super::{parse_by_kind, PResult, ParseCtxt};
 
@@ -12,6 +18,14 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
             @call("mir_retag_raw", args) => {
                 Ok(StatementKind::Retag(RetagKind::Raw, Box::new(self.parse_place(args[0])?)))
             },
+            @call("mir_set_discriminant", args) => {
+                let place = self.parse_place(args[0])?;
+                let var = self.parse_integer_literal(args[1])? as u32;
+                Ok(StatementKind::SetDiscriminant {
+                    place: Box::new(place),
+                    variant_index: VariantIdx::from_u32(var),
+                })
+            },
             ExprKind::Assign { lhs, rhs } => {
                 let lhs = self.parse_place(*lhs)?;
                 let rhs = self.parse_rvalue(*rhs)?;
@@ -21,18 +35,60 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
     }
 
     pub fn parse_terminator(&self, expr_id: ExprId) -> PResult<TerminatorKind<'tcx>> {
-        parse_by_kind!(self, expr_id, _, "terminator",
+        parse_by_kind!(self, expr_id, expr, "terminator",
             @call("mir_return", _args) => {
                 Ok(TerminatorKind::Return)
             },
             @call("mir_goto", args) => {
                 Ok(TerminatorKind::Goto { target: self.parse_block(args[0])? } )
             },
+            ExprKind::Match { scrutinee, arms } => {
+                let discr = self.parse_operand(*scrutinee)?;
+                self.parse_match(arms, expr.span).map(|t| TerminatorKind::SwitchInt { discr, targets: t })
+            },
         )
     }
 
+    fn parse_match(&self, arms: &[ArmId], span: Span) -> PResult<SwitchTargets> {
+        let Some((otherwise, rest)) = arms.split_last() else {
+            return Err(ParseError {
+                span,
+                item_description: format!("no arms"),
+                expected: "at least one arm".to_string(),
+            })
+        };
+
+        let otherwise = &self.thir[*otherwise];
+        let PatKind::Wild = otherwise.pattern.kind else {
+            return Err(ParseError {
+                span: otherwise.span,
+                item_description: format!("{:?}", otherwise.pattern.kind),
+                expected: "wildcard pattern".to_string(),
+            })
+        };
+        let otherwise = self.parse_block(otherwise.body)?;
+
+        let mut values = Vec::new();
+        let mut targets = Vec::new();
+        for arm in rest {
+            let arm = &self.thir[*arm];
+            let PatKind::Constant { value } = arm.pattern.kind else {
+                return Err(ParseError {
+                    span: arm.pattern.span,
+                    item_description: format!("{:?}", arm.pattern.kind),
+                    expected: "constant pattern".to_string(),
+                })
+            };
+            values.push(value.eval_bits(self.tcx, self.param_env, arm.pattern.ty));
+            targets.push(self.parse_block(arm.body)?);
+        }
+
+        Ok(SwitchTargets::new(values.into_iter().zip(targets), otherwise))
+    }
+
     fn parse_rvalue(&self, expr_id: ExprId) -> PResult<Rvalue<'tcx>> {
         parse_by_kind!(self, expr_id, _, "rvalue",
+            @call("mir_discriminant", args) => self.parse_place(args[0]).map(Rvalue::Discriminant),
             ExprKind::Borrow { borrow_kind, arg } => Ok(
                 Rvalue::Ref(self.tcx.lifetimes.re_erased, *borrow_kind, self.parse_place(*arg)?)
             ),
@@ -55,20 +111,50 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
             | ExprKind::ConstParam { .. }
             | ExprKind::ConstBlock { .. } => {
                 Ok(Operand::Constant(Box::new(
-                    crate::build::expr::as_constant::as_constant_inner(expr, |_| None, self.tcx)
+                    as_constant_inner(expr, |_| None, self.tcx)
                 )))
             },
             _ => self.parse_place(expr_id).map(Operand::Copy),
         )
     }
 
     fn parse_place(&self, expr_id: ExprId) -> PResult<Place<'tcx>> {
-        parse_by_kind!(self, expr_id, _, "place",
-            ExprKind::Deref { arg } => Ok(
-                self.parse_place(*arg)?.project_deeper(&[PlaceElem::Deref], self.tcx)
-            ),
-            _ => self.parse_local(expr_id).map(Place::from),
-        )
+        self.parse_place_inner(expr_id).map(|(x, _)| x)
+    }
+
+    fn parse_place_inner(&self, expr_id: ExprId) -> PResult<(Place<'tcx>, PlaceTy<'tcx>)> {
+        let (parent, proj) = parse_by_kind!(self, expr_id, expr, "place",
+            @call("mir_field", args) => {
+                let (parent, ty) = self.parse_place_inner(args[0])?;
+                let field = Field::from_u32(self.parse_integer_literal(args[1])? as u32);
+                let field_ty = ty.field_ty(self.tcx, field);
+                let proj = PlaceElem::Field(field, field_ty);
+                let place = parent.project_deeper(&[proj], self.tcx);
+                return Ok((place, PlaceTy::from_ty(field_ty)));
+            },
+            @call("mir_variant", args) => {
+                (args[0], PlaceElem::Downcast(
+                    None,
+                    VariantIdx::from_u32(self.parse_integer_literal(args[1])? as u32)
+                ))
+            },
+            ExprKind::Deref { arg } => {
+                parse_by_kind!(self, *arg, _, "does not matter",
+                    @call("mir_make_place", args) => return self.parse_place_inner(args[0]),
+                    _ => (*arg, PlaceElem::Deref),
+                )
+            },
+            ExprKind::Index { lhs, index } => (*lhs, PlaceElem::Index(self.parse_local(*index)?)),
+            ExprKind::Field { lhs, name: field, .. } => (*lhs, PlaceElem::Field(*field, expr.ty)),
+            _ => {
+                let place = self.parse_local(expr_id).map(Place::from)?;
+                return Ok((place, PlaceTy::from_ty(expr.ty)))
+            },
+        );
+        let (parent, ty) = self.parse_place_inner(parent)?;
+        let place = parent.project_deeper(&[proj], self.tcx);
+        let ty = ty.projection_ty(self.tcx, proj);
+        Ok((place, ty))
     }
 
     fn parse_local(&self, expr_id: ExprId) -> PResult<Local> {
@@ -102,4 +188,16 @@ impl<'tcx, 'body> ParseCtxt<'tcx, 'body> {
             },
         )
     }
+
+    fn parse_integer_literal(&self, expr_id: ExprId) -> PResult<u128> {
+        parse_by_kind!(self, expr_id, expr, "constant",
+            ExprKind::Literal { .. }
+            | ExprKind::NamedConst { .. }
+            | ExprKind::NonHirLiteral { .. }
+            | ExprKind::ConstBlock { .. } => Ok({
+                let value = as_constant_inner(expr, |_| None, self.tcx);
+                value.literal.eval_bits(self.tcx, self.param_env, value.ty())
+            }),
+        )
+    }
 }

compiler/rustc_mir_transform/src/check_unsafety.rs

@@ -500,7 +500,7 @@ fn unsafety_check_result<'tcx>(
     // `mir_built` force this.
     let body = &tcx.mir_built(def).borrow();
 
-    if body.should_skip() {
+    if body.is_custom_mir() {
         return tcx.arena.alloc(UnsafetyCheckResult {
             violations: Vec::new(),
             used_unsafe_blocks: FxHashSet::default(),