diff --git a/compiler/rustc_mir_build/src/build/custom/mod.rs b/compiler/rustc_mir_build/src/build/custom/mod.rs index 1fd2e40c187e..34fefb99e09c 100644 --- 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 index 03206af33bfb..bc822113ce1f 100644 --- 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> { - 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 { + 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> { 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 { + 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 index 8ba1c122884c..4759f128bbbb 100644 --- a/library/core/src/intrinsics/mir.rs +++ b/library/core/src/intrinsics/mir.rs @@ -82,6 +82,12 @@ define!("mir_retag_raw", fn RetagRaw(place: T)); define!("mir_move", fn Move(place: T) -> T); define!("mir_static", fn Static(s: T) -> &'static T); define!("mir_static_mut", fn StaticMut(s: T) -> *mut T); +define!( + "mir_discriminant", + /// Gets the discriminant of a place. + fn Discriminant(place: T) -> ::Discriminant +); +define!("mir_set_discriminant", fn SetDiscriminant(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 new file mode 100644 index 000000000000..e5cd45637784 --- /dev/null +++ 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 new file mode 100644 index 000000000000..7de9ed0983fe --- /dev/null +++ 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 new file mode 100644 index 000000000000..6fdc3d0f4d4e --- /dev/null +++ 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 new file mode 100644 index 000000000000..95c57d2dca33 --- /dev/null +++ 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 new file mode 100644 index 000000000000..a659ba7c114d --- /dev/null +++ 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 new file mode 100644 index 000000000000..d60e4b1b712f --- /dev/null +++ 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 + } +}