Support common enum operations in custom mir

JakobDegen · JakobDegen · commit 409f4d2adb6c · 2022-12-14T01:10:16.000-08:00
diff --git a/compiler/rustc_mir_build/src/build/custom/mod.rs b/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,
 
diff --git a/compiler/rustc_mir_build/src/build/custom/parse/instruction.rs b/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())
+            }),
+        )
+    }
 }
diff --git a/library/core/src/intrinsics/mir.rs b/library/core/src/intrinsics/mir.rs
@@ -82,6 +82,12 @@ define!("mir_retag_raw", fn RetagRaw<T>(place: T));
 define!("mir_move", fn Move<T>(place: T) -> T);
 define!("mir_static", fn Static<T>(s: T) -> &'static T);
 define!("mir_static_mut", fn StaticMut<T>(s: T) -> *mut T);
+define!(
+    "mir_discriminant",
+    /// Gets the discriminant of a place.
+    fn Discriminant<T>(place: T) -> <T as ::core::marker::DiscriminantKind>::Discriminant
+);
+define!("mir_set_discriminant", fn SetDiscriminant<T>(place: T, index: u32));
 
 /// Convenience macro for generating custom MIR.
 ///
diff --git a/src/test/mir-opt/building/custom/enums.rs b/src/test/mir-opt/building/custom/enums.rs
@@ -0,0 +1,120 @@
+#![feature(custom_mir, core_intrinsics)]
+
+extern crate core;
+use core::intrinsics::mir::*;
+
+// EMIT_MIR enums.switch_bool.built.after.mir
+#[custom_mir(dialect = "built")]
+pub fn switch_bool(b: bool) -> u32 {
+    mir!(
+        {
+            match b {
+                true => t,
+                false => f,
+                _ => f,
+            }
+        }
+
+        t = {
+            RET = 5;
+            Return()
+        }
+
+        f = {
+            RET = 10;
+            Return()
+        }
+    )
+}
+
+// EMIT_MIR enums.switch_option.built.after.mir
+#[custom_mir(dialect = "built")]
+pub fn switch_option(option: Option<()>) -> bool {
+    mir!(
+        {
+            let discr = Discriminant(option);
+            match discr {
+                0 => n,
+                1 => s,
+                _ => s,
+            }
+        }
+
+        n = {
+            RET = false;
+            Return()
+        }
+
+        s = {
+            RET = true;
+            Return()
+        }
+    )
+}
+
+#[repr(u8)]
+enum Bool {
+    False = 0,
+    True = 1,
+}
+
+// EMIT_MIR enums.switch_option_repr.built.after.mir
+#[custom_mir(dialect = "built")]
+fn switch_option_repr(option: Bool) -> bool {
+    mir!(
+        {
+            let discr = Discriminant(option);
+            match discr {
+                0 => f,
+                _ => t,
+            }
+        }
+
+        t = {
+            RET = true;
+            Return()
+        }
+
+        f = {
+            RET = false;
+            Return()
+        }
+    )
+}
+
+// EMIT_MIR enums.set_discr.built.after.mir
+#[custom_mir(dialect = "runtime", phase = "initial")]
+fn set_discr(option: &mut Option<()>) {
+    mir!({
+        SetDiscriminant(*option, 0);
+        Return()
+    })
+}
+
+// EMIT_MIR enums.set_discr_repr.built.after.mir
+#[custom_mir(dialect = "runtime", phase = "initial")]
+fn set_discr_repr(b: &mut Bool) {
+    mir!({
+        SetDiscriminant(*b, 0);
+        Return()
+    })
+}
+
+fn main() {
+    assert_eq!(switch_bool(true), 5);
+    assert_eq!(switch_bool(false), 10);
+
+    assert_eq!(switch_option(Some(())), true);
+    assert_eq!(switch_option(None), false);
+
+    assert_eq!(switch_option_repr(Bool::True), true);
+    assert_eq!(switch_option_repr(Bool::False), false);
+
+    let mut opt = Some(());
+    set_discr(&mut opt);
+    assert_eq!(opt, None);
+
+    let mut b = Bool::True;
+    set_discr_repr(&mut b);
+    assert!(matches!(b, Bool::False));
+}
diff --git a/src/test/mir-opt/building/custom/enums.set_discr.built.after.mir b/src/test/mir-opt/building/custom/enums.set_discr.built.after.mir
@@ -0,0 +1,10 @@
+// MIR for `set_discr` after built
+
+fn set_discr(_1: &mut Option<()>) -> () {
+    let mut _0: ();                      // return place in scope 0 at $DIR/enums.rs:+0:39: +0:39
+
+    bb0: {
+        discriminant((*_1)) = 0;         // scope 0 at $DIR/enums.rs:+2:9: +2:36
+        return;                          // scope 0 at $DIR/enums.rs:+3:9: +3:17
+    }
+}
diff --git a/src/test/mir-opt/building/custom/enums.set_discr_repr.built.after.mir b/src/test/mir-opt/building/custom/enums.set_discr_repr.built.after.mir
@@ -0,0 +1,10 @@
+// MIR for `set_discr_repr` after built
+
+fn set_discr_repr(_1: &mut Bool) -> () {
+    let mut _0: ();                      // return place in scope 0 at $DIR/enums.rs:+0:33: +0:33
+
+    bb0: {
+        discriminant((*_1)) = 0;         // scope 0 at $DIR/enums.rs:+2:9: +2:31
+        return;                          // scope 0 at $DIR/enums.rs:+3:9: +3:17
+    }
+}
diff --git a/src/test/mir-opt/building/custom/enums.switch_bool.built.after.mir b/src/test/mir-opt/building/custom/enums.switch_bool.built.after.mir
@@ -0,0 +1,19 @@
+// MIR for `switch_bool` after built
+
+fn switch_bool(_1: bool) -> u32 {
+    let mut _0: u32;                     // return place in scope 0 at $DIR/enums.rs:+0:32: +0:35
+
+    bb0: {
+        switchInt(_1) -> [1: bb1, 0: bb2, otherwise: bb2]; // scope 0 at $DIR/enums.rs:+3:13: +7:14
+    }
+
+    bb1: {
+        _0 = const 5_u32;                // scope 0 at $DIR/enums.rs:+11:13: +11:20
+        return;                          // scope 0 at $DIR/enums.rs:+12:13: +12:21
+    }
+
+    bb2: {
+        _0 = const 10_u32;               // scope 0 at $DIR/enums.rs:+16:13: +16:21
+        return;                          // scope 0 at $DIR/enums.rs:+17:13: +17:21
+    }
+}
diff --git a/src/test/mir-opt/building/custom/enums.switch_option.built.after.mir b/src/test/mir-opt/building/custom/enums.switch_option.built.after.mir
@@ -0,0 +1,21 @@
+// MIR for `switch_option` after built
+
+fn switch_option(_1: Option<()>) -> bool {
+    let mut _0: bool;                    // return place in scope 0 at $DIR/enums.rs:+0:45: +0:49
+    let mut _2: isize;                   // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+
+    bb0: {
+        _2 = discriminant(_1);           // scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+        switchInt(_2) -> [0: bb1, 1: bb2, otherwise: bb2]; // scope 0 at $DIR/enums.rs:+4:13: +8:14
+    }
+
+    bb1: {
+        _0 = const false;                // scope 0 at $DIR/enums.rs:+12:13: +12:24
+        return;                          // scope 0 at $DIR/enums.rs:+13:13: +13:21
+    }
+
+    bb2: {
+        _0 = const true;                 // scope 0 at $DIR/enums.rs:+17:13: +17:23
+        return;                          // scope 0 at $DIR/enums.rs:+18:13: +18:21
+    }
+}
diff --git a/src/test/mir-opt/building/custom/enums.switch_option_repr.built.after.mir b/src/test/mir-opt/building/custom/enums.switch_option_repr.built.after.mir
@@ -0,0 +1,21 @@
+// MIR for `switch_option_repr` after built
+
+fn switch_option_repr(_1: Bool) -> bool {
+    let mut _0: bool;                    // return place in scope 0 at $DIR/enums.rs:+0:40: +0:44
+    let mut _2: u8;                      // in scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+
+    bb0: {
+        _2 = discriminant(_1);           // scope 0 at $SRC_DIR/core/src/intrinsics/mir.rs:LL:COL
+        switchInt(_2) -> [0: bb2, otherwise: bb1]; // scope 0 at $DIR/enums.rs:+4:13: +7:14
+    }
+
+    bb1: {
+        _0 = const true;                 // scope 0 at $DIR/enums.rs:+11:13: +11:23
+        return;                          // scope 0 at $DIR/enums.rs:+12:13: +12:21
+    }
+
+    bb2: {
+        _0 = const false;                // scope 0 at $DIR/enums.rs:+16:13: +16:24
+        return;                          // scope 0 at $DIR/enums.rs:+17:13: +17:21
+    }
+}
