Auto merge of #120614 - DianQK:simplify-switch-int, r=cjgillot

Transforms match into an assignment statement

Fixes #106459.

We should be able to do some similar transformations, like `enum` to `enum`.

r? mir-opt

compiler/rustc_middle/src/mir/terminator.rs

@@ -85,6 +85,12 @@ impl SwitchTargets {
         self.values.push(value);
         self.targets.insert(self.targets.len() - 1, bb);
     }
+
+    /// Returns true if all targets (including the fallback target) are distinct.
+    #[inline]
+    pub fn is_distinct(&self) -> bool {
+        self.targets.iter().collect::<FxHashSet<_>>().len() == self.targets.len()
+    }
 }
 
 pub struct SwitchTargetsIter<'a> {

compiler/rustc_mir_transform/src/match_branches.rs

@@ -1,11 +1,125 @@
+use rustc_index::IndexSlice;
+use rustc_middle::mir::patch::MirPatch;
 use rustc_middle::mir::*;
-use rustc_middle::ty::TyCtxt;
+use rustc_middle::ty::{ParamEnv, ScalarInt, Ty, TyCtxt};
+use rustc_target::abi::Size;
 use std::iter;
 
 use super::simplify::simplify_cfg;
 
 pub struct MatchBranchSimplification;
 
+impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
+    fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
+        sess.mir_opt_level() >= 1
+    }
+
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        let def_id = body.source.def_id();
+        let param_env = tcx.param_env_reveal_all_normalized(def_id);
+
+        let mut should_cleanup = false;
+        for i in 0..body.basic_blocks.len() {
+            let bbs = &*body.basic_blocks;
+            let bb_idx = BasicBlock::from_usize(i);
+            if !tcx.consider_optimizing(|| format!("MatchBranchSimplification {def_id:?} ")) {
+                continue;
+            }
+
+            match bbs[bb_idx].terminator().kind {
+                TerminatorKind::SwitchInt {
+                    discr: ref _discr @ (Operand::Copy(_) | Operand::Move(_)),
+                    ref targets,
+                    ..
+                    // We require that the possible target blocks don't contain this block.
+                } if !targets.all_targets().contains(&bb_idx) => {}
+                // Only optimize switch int statements
+                _ => continue,
+            };
+
+            if SimplifyToIf.simplify(tcx, body, bb_idx, param_env).is_some() {
+                should_cleanup = true;
+                continue;
+            }
+            if SimplifyToExp::default().simplify(tcx, body, bb_idx, param_env).is_some() {
+                should_cleanup = true;
+                continue;
+            }
+        }
+
+        if should_cleanup {
+            simplify_cfg(body);
+        }
+    }
+}
+
+trait SimplifyMatch<'tcx> {
+    /// Simplifies a match statement, returning true if the simplification succeeds, false otherwise.
+    /// Generic code is written here, and we generally don't need a custom implementation.
+    fn simplify(
+        &mut self,
+        tcx: TyCtxt<'tcx>,
+        body: &mut Body<'tcx>,
+        switch_bb_idx: BasicBlock,
+        param_env: ParamEnv<'tcx>,
+    ) -> Option<()> {
+        let bbs = &body.basic_blocks;
+        let (discr, targets) = match bbs[switch_bb_idx].terminator().kind {
+            TerminatorKind::SwitchInt { ref discr, ref targets, .. } => (discr, targets),
+            _ => unreachable!(),
+        };
+
+        let discr_ty = discr.ty(body.local_decls(), tcx);
+        self.can_simplify(tcx, targets, param_env, bbs, discr_ty)?;
+
+        let mut patch = MirPatch::new(body);
+
+        // Take ownership of items now that we know we can optimize.
+        let discr = discr.clone();
+
+        // Introduce a temporary for the discriminant value.
+        let source_info = bbs[switch_bb_idx].terminator().source_info;
+        let discr_local = patch.new_temp(discr_ty, source_info.span);
+
+        let (_, first) = targets.iter().next().unwrap();
+        let statement_index = bbs[switch_bb_idx].statements.len();
+        let parent_end = Location { block: switch_bb_idx, statement_index };
+        patch.add_statement(parent_end, StatementKind::StorageLive(discr_local));
+        patch.add_assign(parent_end, Place::from(discr_local), Rvalue::Use(discr));
+        self.new_stmts(tcx, targets, param_env, &mut patch, parent_end, bbs, discr_local, discr_ty);
+        patch.add_statement(parent_end, StatementKind::StorageDead(discr_local));
+        patch.patch_terminator(switch_bb_idx, bbs[first].terminator().kind.clone());
+        patch.apply(body);
+        Some(())
+    }
+
+    /// Check that the BBs to be simplified satisfies all distinct and
+    /// that the terminator are the same.
+    /// There are also conditions for different ways of simplification.
+    fn can_simplify(
+        &mut self,
+        tcx: TyCtxt<'tcx>,
+        targets: &SwitchTargets,
+        param_env: ParamEnv<'tcx>,
+        bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
+        discr_ty: Ty<'tcx>,
+    ) -> Option<()>;
+
+    fn new_stmts(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        targets: &SwitchTargets,
+        param_env: ParamEnv<'tcx>,
+        patch: &mut MirPatch<'tcx>,
+        parent_end: Location,
+        bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
+        discr_local: Local,
+        discr_ty: Ty<'tcx>,
+    );
+}
+
+struct SimplifyToIf;
+
 /// If a source block is found that switches between two blocks that are exactly
 /// the same modulo const bool assignments (e.g., one assigns true another false
 /// to the same place), merge a target block statements into the source block,
@@ -37,144 +151,350 @@ pub struct MatchBranchSimplification;
 ///    goto -> bb3;
 /// }
 /// ```
-
+impl<'tcx> SimplifyMatch<'tcx> for SimplifyToIf {
-impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
+    fn can_simplify(
-    fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
+        &mut self,
-        sess.mir_opt_level() >= 1
+        tcx: TyCtxt<'tcx>,
-    }
+        targets: &SwitchTargets,
-
+        param_env: ParamEnv<'tcx>,
-    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
-        let def_id = body.source.def_id();
+        _discr_ty: Ty<'tcx>,
-        let param_env = tcx.param_env_reveal_all_normalized(def_id);
+    ) -> Option<()> {
-
+        if targets.iter().len() != 1 {
-        let bbs = body.basic_blocks.as_mut();
+            return None;
-        let mut should_cleanup = false;
+        }
-        'outer: for bb_idx in bbs.indices() {
+        // We require that the possible target blocks all be distinct.
-            if !tcx.consider_optimizing(|| format!("MatchBranchSimplification {def_id:?} ")) {
+        let (_, first) = targets.iter().next().unwrap();
-                continue;
+        let second = targets.otherwise();
-            }
+        if first == second {
-
+            return None;
-            let (discr, val, first, second) = match bbs[bb_idx].terminator().kind {
+        }
-                TerminatorKind::SwitchInt {
+        // Check that destinations are identical, and if not, then don't optimize this block
-                    discr: ref discr @ (Operand::Copy(_) | Operand::Move(_)),
+        if bbs[first].terminator().kind != bbs[second].terminator().kind {
-                    ref targets,
+            return None;
-                    ..
-                } if targets.iter().len() == 1 => {
-                    let (value, target) = targets.iter().next().unwrap();
-                    // We require that this block and the two possible target blocks all be
-                    // distinct.
-                    if target == targets.otherwise()
-                        || bb_idx == target
-                        || bb_idx == targets.otherwise()
-                    {
-                        continue;
-                    }
-                    (discr, value, target, targets.otherwise())
-                }
-                // Only optimize switch int statements
-                _ => continue,
-            };
-
-            // Check that destinations are identical, and if not, then don't optimize this block
-            if bbs[first].terminator().kind != bbs[second].terminator().kind {
-                continue;
-            }
-
-            // Check that blocks are assignments of consts to the same place or same statement,
-            // and match up 1-1, if not don't optimize this block.
-            let first_stmts = &bbs[first].statements;
-            let scnd_stmts = &bbs[second].statements;
-            if first_stmts.len() != scnd_stmts.len() {
-                continue;
-            }
-            for (f, s) in iter::zip(first_stmts, scnd_stmts) {
-                match (&f.kind, &s.kind) {
-                    // If two statements are exactly the same, we can optimize.
-                    (f_s, s_s) if f_s == s_s => {}
-
-                    // If two statements are const bool assignments to the same place, we can optimize.
-                    (
-                        StatementKind::Assign(box (lhs_f, Rvalue::Use(Operand::Constant(f_c)))),
-                        StatementKind::Assign(box (lhs_s, Rvalue::Use(Operand::Constant(s_c)))),
-                    ) if lhs_f == lhs_s
-                        && f_c.const_.ty().is_bool()
-                        && s_c.const_.ty().is_bool()
-                        && f_c.const_.try_eval_bool(tcx, param_env).is_some()
-                        && s_c.const_.try_eval_bool(tcx, param_env).is_some() => {}
-
-                    // Otherwise we cannot optimize. Try another block.
-                    _ => continue 'outer,
-                }
-            }
-            // Take ownership of items now that we know we can optimize.
-            let discr = discr.clone();
-            let discr_ty = discr.ty(&body.local_decls, tcx);
-
-            // Introduce a temporary for the discriminant value.
-            let source_info = bbs[bb_idx].terminator().source_info;
-            let discr_local = body.local_decls.push(LocalDecl::new(discr_ty, source_info.span));
-
-            // We already checked that first and second are different blocks,
-            // and bb_idx has a different terminator from both of them.
-            let (from, first, second) = bbs.pick3_mut(bb_idx, first, second);
-
-            let new_stmts = iter::zip(&first.statements, &second.statements).map(|(f, s)| {
-                match (&f.kind, &s.kind) {
-                    (f_s, s_s) if f_s == s_s => (*f).clone(),
-
-                    (
-                        StatementKind::Assign(box (lhs, Rvalue::Use(Operand::Constant(f_c)))),
-                        StatementKind::Assign(box (_, Rvalue::Use(Operand::Constant(s_c)))),
-                    ) => {
-                        // From earlier loop we know that we are dealing with bool constants only:
-                        let f_b = f_c.const_.try_eval_bool(tcx, param_env).unwrap();
-                        let s_b = s_c.const_.try_eval_bool(tcx, param_env).unwrap();
-                        if f_b == s_b {
-                            // Same value in both blocks. Use statement as is.
-                            (*f).clone()
-                        } else {
-                            // Different value between blocks. Make value conditional on switch condition.
-                            let size = tcx.layout_of(param_env.and(discr_ty)).unwrap().size;
-                            let const_cmp = Operand::const_from_scalar(
-                                tcx,
-                                discr_ty,
-                                rustc_const_eval::interpret::Scalar::from_uint(val, size),
-                                rustc_span::DUMMY_SP,
-                            );
-                            let op = if f_b { BinOp::Eq } else { BinOp::Ne };
-                            let rhs = Rvalue::BinaryOp(
-                                op,
-                                Box::new((Operand::Copy(Place::from(discr_local)), const_cmp)),
-                            );
-                            Statement {
-                                source_info: f.source_info,
-                                kind: StatementKind::Assign(Box::new((*lhs, rhs))),
-                            }
-                        }
-                    }
-
-                    _ => unreachable!(),
-                }
-            });
-
-            from.statements
-                .push(Statement { source_info, kind: StatementKind::StorageLive(discr_local) });
-            from.statements.push(Statement {
-                source_info,
-                kind: StatementKind::Assign(Box::new((
-                    Place::from(discr_local),
-                    Rvalue::Use(discr),
-                ))),
-            });
-            from.statements.extend(new_stmts);
-            from.statements
-                .push(Statement { source_info, kind: StatementKind::StorageDead(discr_local) });
-            from.terminator_mut().kind = first.terminator().kind.clone();
-            should_cleanup = true;
         }
 
-        if should_cleanup {
+        // Check that blocks are assignments of consts to the same place or same statement,
-            simplify_cfg(body);
+        // and match up 1-1, if not don't optimize this block.
+        let first_stmts = &bbs[first].statements;
+        let second_stmts = &bbs[second].statements;
+        if first_stmts.len() != second_stmts.len() {
+            return None;
+        }
+        for (f, s) in iter::zip(first_stmts, second_stmts) {
+            match (&f.kind, &s.kind) {
+                // If two statements are exactly the same, we can optimize.
+                (f_s, s_s) if f_s == s_s => {}
+
+                // If two statements are const bool assignments to the same place, we can optimize.
+                (
+                    StatementKind::Assign(box (lhs_f, Rvalue::Use(Operand::Constant(f_c)))),
+                    StatementKind::Assign(box (lhs_s, Rvalue::Use(Operand::Constant(s_c)))),
+                ) if lhs_f == lhs_s
+                    && f_c.const_.ty().is_bool()
+                    && s_c.const_.ty().is_bool()
+                    && f_c.const_.try_eval_bool(tcx, param_env).is_some()
+                    && s_c.const_.try_eval_bool(tcx, param_env).is_some() => {}
+
+                // Otherwise we cannot optimize. Try another block.
+                _ => return None,
+            }
+        }
+        Some(())
+    }
+
+    fn new_stmts(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        targets: &SwitchTargets,
+        param_env: ParamEnv<'tcx>,
+        patch: &mut MirPatch<'tcx>,
+        parent_end: Location,
+        bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
+        discr_local: Local,
+        discr_ty: Ty<'tcx>,
+    ) {
+        let (val, first) = targets.iter().next().unwrap();
+        let second = targets.otherwise();
+        // We already checked that first and second are different blocks,
+        // and bb_idx has a different terminator from both of them.
+        let first = &bbs[first];
+        let second = &bbs[second];
+        for (f, s) in iter::zip(&first.statements, &second.statements) {
+            match (&f.kind, &s.kind) {
+                (f_s, s_s) if f_s == s_s => {
+                    patch.add_statement(parent_end, f.kind.clone());
+                }
+
+                (
+                    StatementKind::Assign(box (lhs, Rvalue::Use(Operand::Constant(f_c)))),
+                    StatementKind::Assign(box (_, Rvalue::Use(Operand::Constant(s_c)))),
+                ) => {
+                    // From earlier loop we know that we are dealing with bool constants only:
+                    let f_b = f_c.const_.try_eval_bool(tcx, param_env).unwrap();
+                    let s_b = s_c.const_.try_eval_bool(tcx, param_env).unwrap();
+                    if f_b == s_b {
+                        // Same value in both blocks. Use statement as is.
+                        patch.add_statement(parent_end, f.kind.clone());
+                    } else {
+                        // Different value between blocks. Make value conditional on switch condition.
+                        let size = tcx.layout_of(param_env.and(discr_ty)).unwrap().size;
+                        let const_cmp = Operand::const_from_scalar(
+                            tcx,
+                            discr_ty,
+                            rustc_const_eval::interpret::Scalar::from_uint(val, size),
+                            rustc_span::DUMMY_SP,
+                        );
+                        let op = if f_b { BinOp::Eq } else { BinOp::Ne };
+                        let rhs = Rvalue::BinaryOp(
+                            op,
+                            Box::new((Operand::Copy(Place::from(discr_local)), const_cmp)),
+                        );
+                        patch.add_assign(parent_end, *lhs, rhs);
+                    }
+                }
+
+                _ => unreachable!(),
+            }
+        }
+    }
+}
+
+#[derive(Default)]
+struct SimplifyToExp {
+    transfrom_types: Vec<TransfromType>,
+}
+
+#[derive(Clone, Copy)]
+enum CompareType<'tcx, 'a> {
+    /// Identical statements.
+    Same(&'a StatementKind<'tcx>),
+    /// Assignment statements have the same value.
+    Eq(&'a Place<'tcx>, Ty<'tcx>, ScalarInt),
+    /// Enum variant comparison type.
+    Discr { place: &'a Place<'tcx>, ty: Ty<'tcx>, is_signed: bool },
+}
+
+enum TransfromType {
+    Same,
+    Eq,
+    Discr,
+}
+
+impl From<CompareType<'_, '_>> for TransfromType {
+    fn from(compare_type: CompareType<'_, '_>) -> Self {
+        match compare_type {
+            CompareType::Same(_) => TransfromType::Same,
+            CompareType::Eq(_, _, _) => TransfromType::Eq,
+            CompareType::Discr { .. } => TransfromType::Discr,
+        }
+    }
+}
+
+/// If we find that the value of match is the same as the assignment,
+/// merge a target block statements into the source block,
+/// using cast to transform different integer types.
+///
+/// For example:
+///
+/// ```ignore (MIR)
+/// bb0: {
+///     switchInt(_1) -> [1: bb2, 2: bb3, 3: bb4, otherwise: bb1];
+/// }
+///
+/// bb1: {
+///     unreachable;
+/// }
+///
+/// bb2: {
+///     _0 = const 1_i16;
+///     goto -> bb5;
+/// }
+///
+/// bb3: {
+///     _0 = const 2_i16;
+///     goto -> bb5;
+/// }
+///
+/// bb4: {
+///     _0 = const 3_i16;
+///     goto -> bb5;
+/// }
+/// ```
+///
+/// into:
+///
+/// ```ignore (MIR)
+/// bb0: {
+///    _0 = _3 as i16 (IntToInt);
+///    goto -> bb5;
+/// }
+/// ```
+impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
+    fn can_simplify(
+        &mut self,
+        tcx: TyCtxt<'tcx>,
+        targets: &SwitchTargets,
+        param_env: ParamEnv<'tcx>,
+        bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
+        discr_ty: Ty<'tcx>,
+    ) -> Option<()> {
+        if targets.iter().len() < 2 || targets.iter().len() > 64 {
+            return None;
+        }
+        // We require that the possible target blocks all be distinct.
+        if !targets.is_distinct() {
+            return None;
+        }
+        if !bbs[targets.otherwise()].is_empty_unreachable() {
+            return None;
+        }
+        let mut target_iter = targets.iter();
+        let (first_val, first_target) = target_iter.next().unwrap();
+        let first_terminator_kind = &bbs[first_target].terminator().kind;
+        // Check that destinations are identical, and if not, then don't optimize this block
+        if !targets
+            .iter()
+            .all(|(_, other_target)| first_terminator_kind == &bbs[other_target].terminator().kind)
+        {
+            return None;
+        }
+
+        let discr_size = tcx.layout_of(param_env.and(discr_ty)).unwrap().size;
+        let first_stmts = &bbs[first_target].statements;
+        let (second_val, second_target) = target_iter.next().unwrap();
+        let second_stmts = &bbs[second_target].statements;
+        if first_stmts.len() != second_stmts.len() {
+            return None;
+        }
+
+        fn int_equal(l: ScalarInt, r: impl Into<u128>, size: Size) -> bool {
+            l.try_to_int(l.size()).unwrap()
+                == ScalarInt::try_from_uint(r, size).unwrap().try_to_int(size).unwrap()
+        }
+
+        // We first compare the two branches, and then the other branches need to fulfill the same conditions.
+        let mut compare_types = Vec::new();
+        for (f, s) in iter::zip(first_stmts, second_stmts) {
+            let compare_type = match (&f.kind, &s.kind) {
+                // If two statements are exactly the same, we can optimize.
+                (f_s, s_s) if f_s == s_s => CompareType::Same(f_s),
+
+                // If two statements are assignments with the match values to the same place, we can optimize.
+                (
+                    StatementKind::Assign(box (lhs_f, Rvalue::Use(Operand::Constant(f_c)))),
+                    StatementKind::Assign(box (lhs_s, Rvalue::Use(Operand::Constant(s_c)))),
+                ) if lhs_f == lhs_s
+                    && f_c.const_.ty() == s_c.const_.ty()
+                    && f_c.const_.ty().is_integral() =>
+                {
+                    match (
+                        f_c.const_.try_eval_scalar_int(tcx, param_env),
+                        s_c.const_.try_eval_scalar_int(tcx, param_env),
+                    ) {
+                        (Some(f), Some(s)) if f == s => CompareType::Eq(lhs_f, f_c.const_.ty(), f),
+                        // Enum variants can also be simplified to an assignment statement if their values are equal.
+                        // We need to consider both unsigned and signed scenarios here.
+                        (Some(f), Some(s))
+                            if ((f_c.const_.ty().is_signed() || discr_ty.is_signed())
+                                && int_equal(f, first_val, discr_size)
+                                && int_equal(s, second_val, discr_size))
+                                || (Some(f) == ScalarInt::try_from_uint(first_val, f.size())
+                                    && Some(s)
+                                        == ScalarInt::try_from_uint(second_val, s.size())) =>
+                        {
+                            CompareType::Discr {
+                                place: lhs_f,
+                                ty: f_c.const_.ty(),
+                                is_signed: f_c.const_.ty().is_signed() || discr_ty.is_signed(),
+                            }
+                        }
+                        _ => {
+                            return None;
+                        }
+                    }
+                }
+
+                // Otherwise we cannot optimize. Try another block.
+                _ => return None,
+            };
+            compare_types.push(compare_type);
+        }
+
+        // All remaining BBs need to fulfill the same pattern as the two BBs from the previous step.
+        for (other_val, other_target) in target_iter {
+            let other_stmts = &bbs[other_target].statements;
+            if compare_types.len() != other_stmts.len() {
+                return None;
+            }
+            for (f, s) in iter::zip(&compare_types, other_stmts) {
+                match (*f, &s.kind) {
+                    (CompareType::Same(f_s), s_s) if f_s == s_s => {}
+                    (
+                        CompareType::Eq(lhs_f, f_ty, val),
+                        StatementKind::Assign(box (lhs_s, Rvalue::Use(Operand::Constant(s_c)))),
+                    ) if lhs_f == lhs_s
+                        && s_c.const_.ty() == f_ty
+                        && s_c.const_.try_eval_scalar_int(tcx, param_env) == Some(val) => {}
+                    (
+                        CompareType::Discr { place: lhs_f, ty: f_ty, is_signed },
+                        StatementKind::Assign(box (lhs_s, Rvalue::Use(Operand::Constant(s_c)))),
+                    ) if lhs_f == lhs_s && s_c.const_.ty() == f_ty => {
+                        let Some(f) = s_c.const_.try_eval_scalar_int(tcx, param_env) else {
+                            return None;
+                        };
+                        if is_signed
+                            && s_c.const_.ty().is_signed()
+                            && int_equal(f, other_val, discr_size)
+                        {
+                            continue;
+                        }
+                        if Some(f) == ScalarInt::try_from_uint(other_val, f.size()) {
+                            continue;
+                        }
+                        return None;
+                    }
+                    _ => return None,
+                }
+            }
+        }
+        self.transfrom_types = compare_types.into_iter().map(|c| c.into()).collect();
+        Some(())
+    }
+
+    fn new_stmts(
+        &self,
+        _tcx: TyCtxt<'tcx>,
+        targets: &SwitchTargets,
+        _param_env: ParamEnv<'tcx>,
+        patch: &mut MirPatch<'tcx>,
+        parent_end: Location,
+        bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
+        discr_local: Local,
+        discr_ty: Ty<'tcx>,
+    ) {
+        let (_, first) = targets.iter().next().unwrap();
+        let first = &bbs[first];
+
+        for (t, s) in iter::zip(&self.transfrom_types, &first.statements) {
+            match (t, &s.kind) {
+                (TransfromType::Same, _) | (TransfromType::Eq, _) => {
+                    patch.add_statement(parent_end, s.kind.clone());
+                }
+                (
+                    TransfromType::Discr,
+                    StatementKind::Assign(box (lhs, Rvalue::Use(Operand::Constant(f_c)))),
+                ) => {
+                    let operand = Operand::Copy(Place::from(discr_local));
+                    let r_val = if f_c.const_.ty() == discr_ty {
+                        Rvalue::Use(operand)
+                    } else {
+                        Rvalue::Cast(CastKind::IntToInt, operand, f_c.const_.ty())
+                    };
+                    patch.add_assign(parent_end, *lhs, r_val);
+                }
+                _ => unreachable!(),
+            }
         }
     }
 }

tests/codegen/match-optimized.rs

@@ -26,12 +26,12 @@ pub fn exhaustive_match(e: E) -> u8 {
 // CHECK-NEXT: store i8 1, ptr %_0, align 1
 // CHECK-NEXT: br label %[[EXIT]]
 // CHECK: [[C]]:
-// CHECK-NEXT: store i8 2, ptr %_0, align 1
+// CHECK-NEXT: store i8 3, ptr %_0, align 1
 // CHECK-NEXT: br label %[[EXIT]]
     match e {
         E::A => 0,
         E::B => 1,
-        E::C => 2,
+        E::C => 3,
     }
 }
 

tests/mir-opt/matches_reduce_branches.match_i128_u128.MatchBranchSimplification.diff

@@ -0,0 +1,47 @@
+- // MIR for `match_i128_u128` before MatchBranchSimplification
++ // MIR for `match_i128_u128` after MatchBranchSimplification
+  
+  fn match_i128_u128(_1: EnumAi128) -> u128 {
+      debug i => _1;
+      let mut _0: u128;
+      let mut _2: i128;
++     let mut _3: i128;
+  
+      bb0: {
+          _2 = discriminant(_1);
+-         switchInt(move _2) -> [1: bb3, 2: bb4, 3: bb5, 340282366920938463463374607431768211455: bb2, otherwise: bb1];
+-     }
+- 
+-     bb1: {
+-         unreachable;
+-     }
+- 
+-     bb2: {
+-         _0 = const core::num::<impl u128>::MAX;
+-         goto -> bb6;
+-     }
+- 
+-     bb3: {
+-         _0 = const 1_u128;
+-         goto -> bb6;
+-     }
+- 
+-     bb4: {
+-         _0 = const 2_u128;
+-         goto -> bb6;
+-     }
+- 
+-     bb5: {
+-         _0 = const 3_u128;
+-         goto -> bb6;
+-     }
+- 
+-     bb6: {
++         StorageLive(_3);
++         _3 = move _2;
++         _0 = _3 as u128 (IntToInt);
++         StorageDead(_3);
+          return;
+      }
+  }
+  

tests/mir-opt/matches_reduce_branches.match_i16_i8.MatchBranchSimplification.diff

@@ -0,0 +1,42 @@
+- // MIR for `match_i16_i8` before MatchBranchSimplification
++ // MIR for `match_i16_i8` after MatchBranchSimplification
+  
+  fn match_i16_i8(_1: EnumAi16) -> i8 {
+      debug i => _1;
+      let mut _0: i8;
+      let mut _2: i16;
++     let mut _3: i16;
+  
+      bb0: {
+          _2 = discriminant(_1);
+-         switchInt(move _2) -> [65535: bb3, 2: bb4, 65533: bb2, otherwise: bb1];
+-     }
+- 
+-     bb1: {
+-         unreachable;
+-     }
+- 
+-     bb2: {
+-         _0 = const -3_i8;
+-         goto -> bb5;
+-     }
+- 
+-     bb3: {
+-         _0 = const -1_i8;
+-         goto -> bb5;
+-     }
+- 
+-     bb4: {
+-         _0 = const 2_i8;
+-         goto -> bb5;
+-     }
+- 
+-     bb5: {
++         StorageLive(_3);
++         _3 = move _2;
++         _0 = _3 as i8 (IntToInt);
++         StorageDead(_3);
+          return;
+      }
+  }
+  

tests/mir-opt/matches_reduce_branches.match_i8_i16.MatchBranchSimplification.diff

@@ -0,0 +1,42 @@
+- // MIR for `match_i8_i16` before MatchBranchSimplification
++ // MIR for `match_i8_i16` after MatchBranchSimplification
+  
+  fn match_i8_i16(_1: EnumAi8) -> i16 {
+      debug i => _1;
+      let mut _0: i16;
+      let mut _2: i8;
++     let mut _3: i8;
+  
+      bb0: {
+          _2 = discriminant(_1);
+-         switchInt(move _2) -> [255: bb3, 2: bb4, 253: bb2, otherwise: bb1];
+-     }
+- 
+-     bb1: {
+-         unreachable;
+-     }
+- 
+-     bb2: {
+-         _0 = const -3_i16;
+-         goto -> bb5;
+-     }
+- 
+-     bb3: {
+-         _0 = const -1_i16;
+-         goto -> bb5;
+-     }
+- 
+-     bb4: {
+-         _0 = const 2_i16;
+-         goto -> bb5;
+-     }
+- 
+-     bb5: {
++         StorageLive(_3);
++         _3 = move _2;
++         _0 = _3 as i16 (IntToInt);
++         StorageDead(_3);
+          return;
+      }
+  }
+  

tests/mir-opt/matches_reduce_branches.match_i8_i16_failed.MatchBranchSimplification.diff

@@ -0,0 +1,37 @@
+- // MIR for `match_i8_i16_failed` before MatchBranchSimplification
++ // MIR for `match_i8_i16_failed` after MatchBranchSimplification
+  
+  fn match_i8_i16_failed(_1: EnumAi8) -> i16 {
+      debug i => _1;
+      let mut _0: i16;
+      let mut _2: i8;
+  
+      bb0: {
+          _2 = discriminant(_1);
+          switchInt(move _2) -> [255: bb3, 2: bb4, 253: bb2, otherwise: bb1];
+      }
+  
+      bb1: {
+          unreachable;
+      }
+  
+      bb2: {
+          _0 = const 3_i16;
+          goto -> bb5;
+      }
+  
+      bb3: {
+          _0 = const -1_i16;
+          goto -> bb5;
+      }
+  
+      bb4: {
+          _0 = const 2_i16;
+          goto -> bb5;
+      }
+  
+      bb5: {
+          return;
+      }
+  }
+  

tests/mir-opt/matches_reduce_branches.match_u8_i16.MatchBranchSimplification.diff

@@ -0,0 +1,37 @@
+- // MIR for `match_u8_i16` before MatchBranchSimplification
++ // MIR for `match_u8_i16` after MatchBranchSimplification
+  
+  fn match_u8_i16(_1: EnumAu8) -> i16 {
+      debug i => _1;
+      let mut _0: i16;
+      let mut _2: u8;
++     let mut _3: u8;
+  
+      bb0: {
+          _2 = discriminant(_1);
+-         switchInt(move _2) -> [1: bb3, 2: bb2, otherwise: bb1];
+-     }
+- 
+-     bb1: {
+-         unreachable;
+-     }
+- 
+-     bb2: {
+-         _0 = const 2_i16;
+-         goto -> bb4;
+-     }
+- 
+-     bb3: {
+-         _0 = const 1_i16;
+-         goto -> bb4;
+-     }
+- 
+-     bb4: {
++         StorageLive(_3);
++         _3 = move _2;
++         _0 = _3 as i16 (IntToInt);
++         StorageDead(_3);
+          return;
+      }
+  }
+  

tests/mir-opt/matches_reduce_branches.match_u8_i16_2.MatchBranchSimplification.diff

@@ -0,0 +1,26 @@
+- // MIR for `match_u8_i16_2` before MatchBranchSimplification
++ // MIR for `match_u8_i16_2` after MatchBranchSimplification
+  
+  fn match_u8_i16_2(_1: EnumAu8) -> i16 {
+      let mut _0: i16;
+      let mut _2: u8;
+  
+      bb0: {
+          _2 = discriminant(_1);
+          switchInt(_2) -> [1: bb3, 2: bb1, otherwise: bb2];
+      }
+  
+      bb1: {
+          _0 = const 2_i16;
+          goto -> bb3;
+      }
+  
+      bb2: {
+          unreachable;
+      }
+  
+      bb3: {
+          return;
+      }
+  }
+  

tests/mir-opt/matches_reduce_branches.match_u8_i16_failed.MatchBranchSimplification.diff

@@ -0,0 +1,32 @@
+- // MIR for `match_u8_i16_failed` before MatchBranchSimplification
++ // MIR for `match_u8_i16_failed` after MatchBranchSimplification
+  
+  fn match_u8_i16_failed(_1: EnumAu8) -> i16 {
+      debug i => _1;
+      let mut _0: i16;
+      let mut _2: u8;
+  
+      bb0: {
+          _2 = discriminant(_1);
+          switchInt(move _2) -> [1: bb3, 2: bb2, otherwise: bb1];
+      }
+  
+      bb1: {
+          unreachable;
+      }
+  
+      bb2: {
+          _0 = const 3_i16;
+          goto -> bb4;
+      }
+  
+      bb3: {
+          _0 = const 1_i16;
+          goto -> bb4;
+      }
+  
+      bb4: {
+          return;
+      }
+  }
+  

tests/mir-opt/matches_reduce_branches.match_u8_i16_fallback.MatchBranchSimplification.diff

@@ -0,0 +1,31 @@
+- // MIR for `match_u8_i16_fallback` before MatchBranchSimplification
++ // MIR for `match_u8_i16_fallback` after MatchBranchSimplification
+  
+  fn match_u8_i16_fallback(_1: u8) -> i16 {
+      debug i => _1;
+      let mut _0: i16;
+  
+      bb0: {
+          switchInt(_1) -> [1: bb2, 2: bb3, otherwise: bb1];
+      }
+  
+      bb1: {
+          _0 = const 3_i16;
+          goto -> bb4;
+      }
+  
+      bb2: {
+          _0 = const 1_i16;
+          goto -> bb4;
+      }
+  
+      bb3: {
+          _0 = const 2_i16;
+          goto -> bb4;
+      }
+  
+      bb4: {
+          return;
+      }
+  }
+  

tests/mir-opt/matches_reduce_branches.match_u8_u16.MatchBranchSimplification.diff

@@ -0,0 +1,42 @@
+- // MIR for `match_u8_u16` before MatchBranchSimplification
++ // MIR for `match_u8_u16` after MatchBranchSimplification
+  
+  fn match_u8_u16(_1: EnumBu8) -> u16 {
+      debug i => _1;
+      let mut _0: u16;
+      let mut _2: u8;
++     let mut _3: u8;
+  
+      bb0: {
+          _2 = discriminant(_1);
+-         switchInt(move _2) -> [1: bb3, 2: bb4, 5: bb2, otherwise: bb1];
+-     }
+- 
+-     bb1: {
+-         unreachable;
+-     }
+- 
+-     bb2: {
+-         _0 = const 5_u16;
+-         goto -> bb5;
+-     }
+- 
+-     bb3: {
+-         _0 = const 1_u16;
+-         goto -> bb5;
+-     }
+- 
+-     bb4: {
+-         _0 = const 2_u16;
+-         goto -> bb5;
+-     }
+- 
+-     bb5: {
++         StorageLive(_3);
++         _3 = move _2;
++         _0 = _3 as u16 (IntToInt);
++         StorageDead(_3);
+          return;
+      }
+  }
+  

tests/mir-opt/matches_reduce_branches.match_u8_u16_2.MatchBranchSimplification.diff

@@ -0,0 +1,37 @@
+- // MIR for `match_u8_u16_2` before MatchBranchSimplification
++ // MIR for `match_u8_u16_2` after MatchBranchSimplification
+  
+  fn match_u8_u16_2(_1: EnumBu8) -> i16 {
+      let mut _0: i16;
+      let mut _2: u8;
+  
+      bb0: {
+          _2 = discriminant(_1);
+          switchInt(_2) -> [1: bb1, 2: bb2, 5: bb3, otherwise: bb4];
+      }
+  
+      bb1: {
+          _0 = const 1_i16;
+          goto -> bb5;
+      }
+  
+      bb2: {
+          _0 = const 2_i16;
+          goto -> bb5;
+      }
+  
+      bb3: {
+          _0 = const 5_i16;
+          _0 = const 5_i16;
+          goto -> bb5;
+      }
+  
+      bb4: {
+          unreachable;
+      }
+  
+      bb5: {
+          return;
+      }
+  }
+  

tests/mir-opt/matches_reduce_branches.rs

@@ -1,18 +1,28 @@
-// skip-filecheck
 //@ unit-test: MatchBranchSimplification
 
+#![feature(repr128)]
+#![feature(core_intrinsics)]
+#![feature(custom_mir)]
+
+use std::intrinsics::mir::*;
 
 // EMIT_MIR matches_reduce_branches.foo.MatchBranchSimplification.diff
-// EMIT_MIR matches_reduce_branches.bar.MatchBranchSimplification.diff
-// EMIT_MIR matches_reduce_branches.match_nested_if.MatchBranchSimplification.diff
-
 fn foo(bar: Option<()>) {
+    // CHECK-LABEL: fn foo(
+    // CHECK: = Eq(
+    // CHECK: switchInt
+    // CHECK-NOT: switchInt
     if matches!(bar, None) {
         ()
     }
 }
 
+// EMIT_MIR matches_reduce_branches.bar.MatchBranchSimplification.diff
 fn bar(i: i32) -> (bool, bool, bool, bool) {
+    // CHECK-LABEL: fn bar(
+    // CHECK: = Ne(
+    // CHECK: = Eq(
+    // CHECK-NOT: switchInt
     let a;
     let b;
     let c;
@@ -38,7 +48,10 @@ fn bar(i: i32) -> (bool, bool, bool, bool) {
     (a, b, c, d)
 }
 
+// EMIT_MIR matches_reduce_branches.match_nested_if.MatchBranchSimplification.diff
 fn match_nested_if() -> bool {
+    // CHECK-LABEL: fn match_nested_if(
+    // CHECK-NOT: switchInt
     let val = match () {
         () if if if if true { true } else { false } { true } else { false } {
             true
@@ -53,9 +66,221 @@ fn match_nested_if() -> bool {
     val
 }
 
+#[repr(u8)]
+enum EnumAu8 {
+    A = 1,
+    B = 2,
+}
+
+// EMIT_MIR matches_reduce_branches.match_u8_i16.MatchBranchSimplification.diff
+fn match_u8_i16(i: EnumAu8) -> i16 {
+    // CHECK-LABEL: fn match_u8_i16(
+    // CHECK-NOT: switchInt
+    // CHECK: _0 = _3 as i16 (IntToInt);
+    // CHECH: return
+    match i {
+        EnumAu8::A => 1,
+        EnumAu8::B => 2,
+    }
+}
+
+// EMIT_MIR matches_reduce_branches.match_u8_i16_2.MatchBranchSimplification.diff
+// Check for different instruction lengths
+#[custom_mir(dialect = "built")]
+fn match_u8_i16_2(i: EnumAu8) -> i16 {
+    // CHECK-LABEL: fn match_u8_i16_2(
+    // CHECK: switchInt
+    mir!(
+        {
+            let a = Discriminant(i);
+            match a {
+                1 => bb1,
+                2 => bb2,
+                _ => unreachable_bb,
+            }
+        }
+        bb1 = {
+            Goto(ret)
+        }
+        bb2 = {
+            RET = 2;
+            Goto(ret)
+        }
+        unreachable_bb = {
+            Unreachable()
+        }
+        ret = {
+            Return()
+        }
+    )
+}
+
+// EMIT_MIR matches_reduce_branches.match_u8_i16_failed.MatchBranchSimplification.diff
+fn match_u8_i16_failed(i: EnumAu8) -> i16 {
+    // CHECK-LABEL: fn match_u8_i16_failed(
+    // CHECK: switchInt
+    match i {
+        EnumAu8::A => 1,
+        EnumAu8::B => 3,
+    }
+}
+
+// EMIT_MIR matches_reduce_branches.match_u8_i16_fallback.MatchBranchSimplification.diff
+fn match_u8_i16_fallback(i: u8) -> i16 {
+    // CHECK-LABEL: fn match_u8_i16_fallback(
+    // CHECK: switchInt
+    match i {
+        1 => 1,
+        2 => 2,
+        _ => 3,
+    }
+}
+
+#[repr(u8)]
+enum EnumBu8 {
+    A = 1,
+    B = 2,
+    C = 5,
+}
+
+// EMIT_MIR matches_reduce_branches.match_u8_u16.MatchBranchSimplification.diff
+fn match_u8_u16(i: EnumBu8) -> u16 {
+    // CHECK-LABEL: fn match_u8_u16(
+    // CHECK-NOT: switchInt
+    // CHECK: _0 = _3 as u16 (IntToInt);
+    // CHECH: return
+    match i {
+        EnumBu8::A => 1,
+        EnumBu8::B => 2,
+        EnumBu8::C => 5,
+    }
+}
+
+// EMIT_MIR matches_reduce_branches.match_u8_u16_2.MatchBranchSimplification.diff
+// Check for different instruction lengths
+#[custom_mir(dialect = "built")]
+fn match_u8_u16_2(i: EnumBu8) -> i16 {
+    // CHECK-LABEL: fn match_u8_u16_2(
+    // CHECK: switchInt
+    mir!(
+        {
+            let a = Discriminant(i);
+            match a {
+                1 => bb1,
+                2 => bb2,
+                5 => bb5,
+                _ => unreachable_bb,
+            }
+        }
+        bb1 = {
+            RET = 1;
+            Goto(ret)
+        }
+        bb2 = {
+            RET = 2;
+            Goto(ret)
+        }
+        bb5 = {
+            RET = 5;
+            RET = 5;
+            Goto(ret)
+        }
+        unreachable_bb = {
+            Unreachable()
+        }
+        ret = {
+            Return()
+        }
+    )
+}
+
+#[repr(i8)]
+enum EnumAi8 {
+    A = -1,
+    B = 2,
+    C = -3,
+}
+
+// EMIT_MIR matches_reduce_branches.match_i8_i16.MatchBranchSimplification.diff
+fn match_i8_i16(i: EnumAi8) -> i16 {
+    // CHECK-LABEL: fn match_i8_i16(
+    // CHECK-NOT: switchInt
+    // CHECK: _0 = _3 as i16 (IntToInt);
+    // CHECH: return
+    match i {
+        EnumAi8::A => -1,
+        EnumAi8::B => 2,
+        EnumAi8::C => -3,
+    }
+}
+
+// EMIT_MIR matches_reduce_branches.match_i8_i16_failed.MatchBranchSimplification.diff
+fn match_i8_i16_failed(i: EnumAi8) -> i16 {
+    // CHECK-LABEL: fn match_i8_i16_failed(
+    // CHECK: switchInt
+    match i {
+        EnumAi8::A => -1,
+        EnumAi8::B => 2,
+        EnumAi8::C => 3,
+    }
+}
+
+#[repr(i16)]
+enum EnumAi16 {
+    A = -1,
+    B = 2,
+    C = -3,
+}
+
+// EMIT_MIR matches_reduce_branches.match_i16_i8.MatchBranchSimplification.diff
+fn match_i16_i8(i: EnumAi16) -> i8 {
+    // CHECK-LABEL: fn match_i16_i8(
+    // CHECK-NOT: switchInt
+    // CHECK: _0 = _3 as i8 (IntToInt);
+    // CHECH: return
+    match i {
+        EnumAi16::A => -1,
+        EnumAi16::B => 2,
+        EnumAi16::C => -3,
+    }
+}
+
+#[repr(i128)]
+enum EnumAi128 {
+    A = 1,
+    B = 2,
+    C = 3,
+    D = -1,
+}
+
+// EMIT_MIR matches_reduce_branches.match_i128_u128.MatchBranchSimplification.diff
+fn match_i128_u128(i: EnumAi128) -> u128 {
+    // CHECK-LABEL: fn match_i128_u128(
+    // CHECK-NOT: switchInt
+    // CHECK: _0 = _3 as u128 (IntToInt);
+    // CHECH: return
+    match i {
+        EnumAi128::A => 1,
+        EnumAi128::B => 2,
+        EnumAi128::C => 3,
+        EnumAi128::D => u128::MAX,
+    }
+}
+
 fn main() {
     let _ = foo(None);
     let _ = foo(Some(()));
     let _ = bar(0);
     let _ = match_nested_if();
+    let _ = match_u8_i16(EnumAu8::A);
+    let _ = match_u8_i16_2(EnumAu8::A);
+    let _ = match_u8_i16_failed(EnumAu8::A);
+    let _ = match_u8_i16_fallback(1);
+    let _ = match_u8_u16(EnumBu8::A);
+    let _ = match_u8_u16_2(EnumBu8::A);
+    let _ = match_i8_i16(EnumAi8::A);
+    let _ = match_i8_i16_failed(EnumAi8::A);
+    let _ = match_i8_i16(EnumAi8::A);
+    let _ = match_i16_i8(EnumAi16::A);
+    let _ = match_i128_u128(EnumAi128::A);
 }

tests/mir-opt/matches_u8.exhaustive_match.MatchBranchSimplification.diff

@@ -5,27 +5,32 @@
       debug e => _1;
       let mut _0: u8;
       let mut _2: isize;
++     let mut _3: isize;
   
       bb0: {
           _2 = discriminant(_1);
-          switchInt(move _2) -> [0: bb3, 1: bb2, otherwise: bb1];
+-         switchInt(move _2) -> [0: bb3, 1: bb2, otherwise: bb1];
-      }
+-     }
-  
+- 
-      bb1: {
+-     bb1: {
-          unreachable;
+-         unreachable;
-      }
+-     }
-  
+- 
-      bb2: {
+-     bb2: {
-          _0 = const 1_u8;
+-         _0 = const 1_u8;
-          goto -> bb4;
+-         goto -> bb4;
-      }
+-     }
-  
+- 
-      bb3: {
+-     bb3: {
-          _0 = const 0_u8;
+-         _0 = const 0_u8;
-          goto -> bb4;
+-         goto -> bb4;
-      }
+-     }
-  
+- 
-      bb4: {
+-     bb4: {
++         StorageLive(_3);
++         _3 = move _2;
++         _0 = _3 as u8 (IntToInt);
++         StorageDead(_3);
           return;
       }
   }

tests/mir-opt/matches_u8.exhaustive_match_i8.MatchBranchSimplification.diff

@@ -5,27 +5,32 @@
       debug e => _1;
       let mut _0: i8;
       let mut _2: isize;
++     let mut _3: isize;
   
       bb0: {
           _2 = discriminant(_1);
-          switchInt(move _2) -> [0: bb3, 1: bb2, otherwise: bb1];
+-         switchInt(move _2) -> [0: bb3, 1: bb2, otherwise: bb1];
-      }
+-     }
-  
+- 
-      bb1: {
+-     bb1: {
-          unreachable;
+-         unreachable;
-      }
+-     }
-  
+- 
-      bb2: {
+-     bb2: {
-          _0 = const 1_i8;
+-         _0 = const 1_i8;
-          goto -> bb4;
+-         goto -> bb4;
-      }
+-     }
-  
+- 
-      bb3: {
+-     bb3: {
-          _0 = const 0_i8;
+-         _0 = const 0_i8;
-          goto -> bb4;
+-         goto -> bb4;
-      }
+-     }
-  
+- 
-      bb4: {
+-     bb4: {
++         StorageLive(_3);
++         _3 = move _2;
++         _0 = _3 as i8 (IntToInt);
++         StorageDead(_3);
           return;
       }
   }