Support common enum operations in custom mir

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,10 @@
 use rustc_middle::mir::interpret::{ConstValue, Scalar};
 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 +17,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 +34,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,7 +110,7 @@ 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),
@@ -102,4 +157,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())
+            }),
+        )
+    }
 }