Custom MIR: Many more improvements

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(),