Refactor MatchBranchSimplification

compiler/rustc_mir_transform/src/match_branches.rs

@@ -1,11 +1,116 @@
+use rustc_index::IndexVec;
 use rustc_middle::mir::*;
-use rustc_middle::ty::TyCtxt;
+use rustc_middle::ty::{ParamEnv, Ty, TyCtxt};
 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 bbs = body.basic_blocks.as_mut();
+        let mut should_cleanup = false;
+        for bb_idx in bbs.indices() {
+            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, &mut body.local_decls, bbs, bb_idx, param_env) {
+                should_cleanup = true;
+                continue;
+            }
+        }
+
+        if should_cleanup {
+            simplify_cfg(body);
+        }
+    }
+}
+
+trait SimplifyMatch<'tcx> {
+    fn simplify(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>,
+        bbs: &mut IndexVec<BasicBlock, BasicBlockData<'tcx>>,
+        switch_bb_idx: BasicBlock,
+        param_env: ParamEnv<'tcx>,
+    ) -> bool {
+        let (discr, targets) = match bbs[switch_bb_idx].terminator().kind {
+            TerminatorKind::SwitchInt { ref discr, ref targets, .. } => (discr, targets),
+            _ => unreachable!(),
+        };
+
+        if !self.can_simplify(tcx, targets, param_env, bbs) {
+            return false;
+        }
+
+        // Take ownership of items now that we know we can optimize.
+        let discr = discr.clone();
+        let discr_ty = discr.ty(local_decls, tcx);
+
+        // Introduce a temporary for the discriminant value.
+        let source_info = bbs[switch_bb_idx].terminator().source_info;
+        let discr_local = 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 new_stmts = self.new_stmts(tcx, targets, param_env, bbs, discr_local.clone(), discr_ty);
+        let (_, first) = targets.iter().next().unwrap();
+        let (from, first) = bbs.pick2_mut(switch_bb_idx, first);
+        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();
+        true
+    }
+
+    fn can_simplify(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        targets: &SwitchTargets,
+        param_env: ParamEnv<'tcx>,
+        bbs: &IndexVec<BasicBlock, BasicBlockData<'tcx>>,
+    ) -> bool;
+
+    fn new_stmts(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        targets: &SwitchTargets,
+        param_env: ParamEnv<'tcx>,
+        bbs: &IndexVec<BasicBlock, BasicBlockData<'tcx>>,
+        discr_local: Local,
+        discr_ty: Ty<'tcx>,
+    ) -> Vec<Statement<'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,57 +142,36 @@ 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 {
+        &self,
-        sess.mir_opt_level() >= 1
+        tcx: TyCtxt<'tcx>,
+        targets: &SwitchTargets,
+        param_env: ParamEnv<'tcx>,
+        bbs: &IndexVec<BasicBlock, BasicBlockData<'tcx>>,
+    ) -> bool {
+        if targets.iter().len() != 1 {
+            return false;
         }
-
+        // We require that the possible target blocks all be distinct.
-    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        let (_, first) = targets.iter().next().unwrap();
-        let def_id = body.source.def_id();
+        let second = targets.otherwise();
-        let param_env = tcx.param_env_reveal_all_normalized(def_id);
+        if first == second {
-
+            return false;
-        let bbs = body.basic_blocks.as_mut();
-        let mut should_cleanup = false;
-        'outer: for bb_idx in bbs.indices() {
-            if !tcx.consider_optimizing(|| format!("MatchBranchSimplification {def_id:?} ")) {
-                continue;
         }
-
-            let (discr, val, first, second) = match bbs[bb_idx].terminator().kind {
-                TerminatorKind::SwitchInt {
-                    discr: ref discr @ (Operand::Copy(_) | Operand::Move(_)),
-                    ref targets,
-                    ..
-                } 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;
+            return false;
         }
 
         // 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;
+        let second_stmts = &bbs[second].statements;
-            if first_stmts.len() != scnd_stmts.len() {
+        if first_stmts.len() != second_stmts.len() {
-                continue;
+            return false;
         }
-            for (f, s) in iter::zip(first_stmts, scnd_stmts) {
+        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 => {}
@@ -103,20 +187,27 @@ impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
                     && s_c.const_.try_eval_bool(tcx, param_env).is_some() => {}
 
                 // Otherwise we cannot optimize. Try another block.
-                    _ => continue 'outer,
+                _ => return false,
             }
         }
-            // Take ownership of items now that we know we can optimize.
+        true
-            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));
 
+    fn new_stmts(
+        &self,
+        tcx: TyCtxt<'tcx>,
+        targets: &SwitchTargets,
+        param_env: ParamEnv<'tcx>,
+        bbs: &IndexVec<BasicBlock, BasicBlockData<'tcx>>,
+        discr_local: Local,
+        discr_ty: Ty<'tcx>,
+    ) -> Vec<Statement<'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 (from, first, second) = bbs.pick3_mut(bb_idx, first, second);
+        let first = &bbs[first];
+        let second = &bbs[second];
 
         let new_stmts = iter::zip(&first.statements, &second.statements).map(|(f, s)| {
             match (&f.kind, &s.kind) {
@@ -156,25 +247,6 @@ impl<'tcx> MirPass<'tcx> for MatchBranchSimplification {
                 _ => unreachable!(),
             }
         });
-
+        new_stmts.collect()
-            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 {
-            simplify_cfg(body);
-        }
     }
 }