MIR opt: separate constant predecessors of a switch

compiler/rustc_mir/src/transform/mod.rs

@@ -48,6 +48,7 @@ pub mod remove_unneeded_drops;
 pub mod remove_zsts;
 pub mod required_consts;
 pub mod rustc_peek;
+pub mod separate_const_switch;
 pub mod simplify;
 pub mod simplify_branches;
 pub mod simplify_comparison_integral;
@@ -501,6 +502,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         // inst combine is after MatchBranchSimplification to clean up Ne(_1, false)
         &multiple_return_terminators::MultipleReturnTerminators,
         &instcombine::InstCombine,
+        &separate_const_switch::SeparateConstSwitch,
         &const_prop::ConstProp,
         &simplify_branches::SimplifyBranches::new("after-const-prop"),
         &early_otherwise_branch::EarlyOtherwiseBranch,

compiler/rustc_mir/src/transform/separate_const_switch.rs

@@ -0,0 +1,343 @@
+//! A pass that duplicates switch-terminated blocks
+//! into a new copy for each predecessor, provided
+//! the predecessor sets the value being switched
+//! over to a constant.
+//!
+//! The purpose of this pass is to help constant
+//! propagation passes to simplify the switch terminator
+//! of the copied blocks into gotos when some predecessors
+//! statically determine the output of switches.
+//!
+//! ```text
+//!     x = 12 ---              ---> something
+//!               \            / 12
+//!                --> switch x
+//!               /            \ otherwise
+//!     x = y  ---              ---> something else
+//! ```
+//! becomes
+//! ```text
+//!     x = 12 ---> switch x ------> something
+//!                          \ / 12
+//!                           X
+//!                          / \ otherwise
+//!     x = y  ---> switch x ------> something else
+//! ```
+//! so it can hopefully later be turned by another pass into
+//! ```text
+//!     x = 12 --------------------> something
+//!                            / 12
+//!                           /
+//!                          /   otherwise
+//!     x = y  ---- switch x ------> something else
+//! ```
+//!
+//! This optimization is meant to cover simple cases
+//! like `?` desugaring. For now, it thus focuses on
+//! simplicity rather than completeness (it notably
+//! sometimes duplicates abusively).
+
+use crate::transform::MirPass;
+use rustc_middle::mir::*;
+use rustc_middle::ty::TyCtxt;
+use smallvec::SmallVec;
+
+pub struct SeparateConstSwitch;
+
+impl<'tcx> MirPass<'tcx> for SeparateConstSwitch {
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        if tcx.sess.mir_opt_level() < 4 {
+            return;
+        }
+
+        // If execution did something, applying a simplification layer
+        // helps later passes optimize the copy away.
+        if separate_const_switch(body) > 0 {
+            super::simplify::simplify_cfg(tcx, body);
+        }
+    }
+}
+
+/// Returns the amount of blocks that were duplicated
+pub fn separate_const_switch<'tcx>(body: &mut Body<'tcx>) -> usize {
+    let mut new_blocks: SmallVec<[(BasicBlock, BasicBlock); 6]> = SmallVec::new();
+    let predecessors = body.predecessors();
+    'block_iter: for (block_id, block) in body.basic_blocks().iter_enumerated() {
+        if let TerminatorKind::SwitchInt {
+            discr: Operand::Copy(switch_place) | Operand::Move(switch_place),
+            ..
+        } = block.terminator().kind
+        {
+            // If the block is on an unwind path, do not
+            // apply the optimization as unwind paths
+            // rely on a unique parent invariant
+            if block.is_cleanup {
+                continue 'block_iter;
+            }
+
+            // If the block has fewer than 2 predecessors, ignore it
+            // we could maybe chain blocks that have exactly one
+            // predecessor, but for now we ignore that
+            if predecessors[block_id].len() < 2 {
+                continue 'block_iter;
+            }
+
+            // First, let's find a non-const place
+            // that determines the result of the switch
+            if let Some(switch_place) = find_determining_place(switch_place, block) {
+                // We now have an input place for which it would
+                // be interesting if predecessors assigned it from a const
+
+                let mut predecessors_left = predecessors[block_id].len();
+                'predec_iter: for predecessor_id in predecessors[block_id].iter().copied() {
+                    let predecessor = &body.basic_blocks()[predecessor_id];
+
+                    // First we make sure the predecessor jumps
+                    // in a reasonable way
+                    match &predecessor.terminator().kind {
+                        // The following terminators are
+                        // unconditionally valid
+                        TerminatorKind::Goto { .. } | TerminatorKind::SwitchInt { .. } => {}
+
+                        TerminatorKind::FalseEdge { real_target, .. } => {
+                            if *real_target != block_id {
+                                continue 'predec_iter;
+                            }
+                        }
+
+                        // The following terminators are not allowed
+                        TerminatorKind::Resume
+                        | TerminatorKind::Drop { .. }
+                        | TerminatorKind::DropAndReplace { .. }
+                        | TerminatorKind::Call { .. }
+                        | TerminatorKind::Assert { .. }
+                        | TerminatorKind::FalseUnwind { .. }
+                        | TerminatorKind::Yield { .. }
+                        | TerminatorKind::Abort
+                        | TerminatorKind::Return
+                        | TerminatorKind::Unreachable
+                        | TerminatorKind::InlineAsm { .. }
+                        | TerminatorKind::GeneratorDrop => {
+                            continue 'predec_iter;
+                        }
+                    }
+
+                    if is_likely_const(switch_place, predecessor) {
+                        new_blocks.push((predecessor_id, block_id));
+                        predecessors_left -= 1;
+                        if predecessors_left < 2 {
+                            // If the original block only has one predecessor left,
+                            // we have nothing left to do
+                            break 'predec_iter;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    // Once the analysis is done, perform the duplication
+    let body_span = body.span;
+    let copied_blocks = new_blocks.len();
+    let blocks = body.basic_blocks_mut();
+    for (pred_id, target_id) in new_blocks {
+        let new_block = blocks[target_id].clone();
+        let new_block_id = blocks.push(new_block);
+        let terminator = blocks[pred_id].terminator_mut();
+
+        match terminator.kind {
+            TerminatorKind::Goto { ref mut target } => {
+                *target = new_block_id;
+            }
+
+            TerminatorKind::FalseEdge { ref mut real_target, .. } => {
+                if *real_target == target_id {
+                    *real_target = new_block_id;
+                }
+            }
+
+            TerminatorKind::SwitchInt { ref mut targets, .. } => {
+                targets.all_targets_mut().iter_mut().for_each(|x| {
+                    if *x == target_id {
+                        *x = new_block_id;
+                    }
+                });
+            }
+
+            TerminatorKind::Resume
+            | TerminatorKind::Abort
+            | TerminatorKind::Return
+            | TerminatorKind::Unreachable
+            | TerminatorKind::GeneratorDrop
+            | TerminatorKind::Assert { .. }
+            | TerminatorKind::DropAndReplace { .. }
+            | TerminatorKind::FalseUnwind { .. }
+            | TerminatorKind::Drop { .. }
+            | TerminatorKind::Call { .. }
+            | TerminatorKind::InlineAsm { .. }
+            | TerminatorKind::Yield { .. } => {
+                span_bug!(
+                    body_span,
+                    "basic block terminator had unexpected kind {:?}",
+                    &terminator.kind
+                )
+            }
+        }
+    }
+
+    copied_blocks
+}
+
+/// This function describes a rough heuristic guessing
+/// whether a place is last set with a const within the block.
+/// Notably, it will be overly pessimistic in cases that are already
+/// not handled by `separate_const_switch`.
+fn is_likely_const<'tcx>(mut tracked_place: Place<'tcx>, block: &BasicBlockData<'tcx>) -> bool {
+    for statement in block.statements.iter().rev() {
+        match &statement.kind {
+            StatementKind::Assign(assign) => {
+                if assign.0 == tracked_place {
+                    match assign.1 {
+                        // These rvalues are definitely constant
+                        Rvalue::Use(Operand::Constant(_))
+                        | Rvalue::Ref(_, _, _)
+                        | Rvalue::AddressOf(_, _)
+                        | Rvalue::Cast(_, Operand::Constant(_), _)
+                        | Rvalue::NullaryOp(_, _)
+                        | Rvalue::UnaryOp(_, Operand::Constant(_)) => return true,
+
+                        // These rvalues make things ambiguous
+                        Rvalue::Repeat(_, _)
+                        | Rvalue::ThreadLocalRef(_)
+                        | Rvalue::Len(_)
+                        | Rvalue::BinaryOp(_, _)
+                        | Rvalue::CheckedBinaryOp(_, _)
+                        | Rvalue::Aggregate(_, _) => return false,
+
+                        // These rvalues move the place to track
+                        Rvalue::Cast(_, Operand::Copy(place) | Operand::Move(place), _)
+                        | Rvalue::Use(Operand::Copy(place) | Operand::Move(place))
+                        | Rvalue::UnaryOp(_, Operand::Copy(place) | Operand::Move(place))
+                        | Rvalue::Discriminant(place) => tracked_place = place,
+                    }
+                }
+            }
+
+            // If the discriminant is set, it is always set
+            // as a constant, so the job is done.
+            // As we are **ignoring projections**, if the place
+            // we are tracking sees its discriminant be set,
+            // that means we had to be tracking the discriminant
+            // specifically (as it is impossible to switch over
+            // an enum directly, and if we were switching over
+            // its content, we would have had to at least cast it to
+            // some variant first)
+            StatementKind::SetDiscriminant { place, .. } => {
+                if **place == tracked_place {
+                    return true;
+                }
+            }
+
+            // If inline assembly is found, we probably should
+            // not try to analyze the code
+            StatementKind::LlvmInlineAsm(_) => return false,
+
+            // These statements have no influence on the place
+            // we are interested in
+            StatementKind::FakeRead(_)
+            | StatementKind::StorageLive(_)
+            | StatementKind::Retag(_, _)
+            | StatementKind::AscribeUserType(_, _)
+            | StatementKind::Coverage(_)
+            | StatementKind::StorageDead(_)
+            | StatementKind::CopyNonOverlapping(_)
+            | StatementKind::Nop => {}
+        }
+    }
+
+    // If no good reason for the place to be const is found,
+    // give up. We could maybe go up predecessors, but in
+    // most cases giving up now should be sufficient.
+    false
+}
+
+/// Finds a unique place that entirely determines the value
+/// of `switch_place`, if it exists. This is only a heuristic.
+/// Ideally we would like to track multiple determining places
+/// for some edge cases, but one is enough for a lot of situations.
+fn find_determining_place<'tcx>(
+    mut switch_place: Place<'tcx>,
+    block: &BasicBlockData<'tcx>,
+) -> Option<Place<'tcx>> {
+    for statement in block.statements.iter().rev() {
+        match &statement.kind {
+            StatementKind::Assign(op) => {
+                if op.0 != switch_place {
+                    continue;
+                }
+
+                match op.1 {
+                    // The following rvalues move the place
+                    // that may be const in the predecessor
+                    Rvalue::Use(Operand::Move(new) | Operand::Copy(new))
+                    | Rvalue::UnaryOp(_, Operand::Copy(new) | Operand::Move(new))
+                    | Rvalue::Cast(_, Operand::Move(new) | Operand::Copy(new), _)
+                    | Rvalue::Repeat(Operand::Move(new) | Operand::Copy(new), _)
+                    | Rvalue::Discriminant(new)
+                    => switch_place = new,
+
+                    // The following rvalues might still make the block
+                    // be valid but for now we reject them
+                    Rvalue::Len(_)
+                    | Rvalue::Ref(_, _, _)
+                    | Rvalue::BinaryOp(_, _)
+                    | Rvalue::CheckedBinaryOp(_, _)
+                    | Rvalue::Aggregate(_, _)
+
+                    // The following rvalues definitely mean we cannot
+                    // or should not apply this optimization
+                    | Rvalue::Use(Operand::Constant(_))
+                    | Rvalue::Repeat(Operand::Constant(_), _)
+                    | Rvalue::ThreadLocalRef(_)
+                    | Rvalue::AddressOf(_, _)
+                    | Rvalue::NullaryOp(_, _)
+                    | Rvalue::UnaryOp(_, Operand::Constant(_))
+                    | Rvalue::Cast(_, Operand::Constant(_), _)
+                    => return None,
+                }
+            }
+
+            // These statements have no influence on the place
+            // we are interested in
+            StatementKind::FakeRead(_)
+            | StatementKind::StorageLive(_)
+            | StatementKind::StorageDead(_)
+            | StatementKind::Retag(_, _)
+            | StatementKind::AscribeUserType(_, _)
+            | StatementKind::Coverage(_)
+            | StatementKind::CopyNonOverlapping(_)
+            | StatementKind::Nop => {}
+
+            // If inline assembly is found, we probably should
+            // not try to analyze the code
+            StatementKind::LlvmInlineAsm(_) => return None,
+
+            // If the discriminant is set, it is always set
+            // as a constant, so the job is already done.
+            // As we are **ignoring projections**, if the place
+            // we are tracking sees its discriminant be set,
+            // that means we had to be tracking the discriminant
+            // specifically (as it is impossible to switch over
+            // an enum directly, and if we were switching over
+            // its content, we would have had to at least cast it to
+            // some variant first)
+            StatementKind::SetDiscriminant { place, .. } => {
+                if **place == switch_place {
+                    return None;
+                }
+            }
+        }
+    }
+
+    Some(switch_place)
+}