4 interpret::Scalar, BasicBlock, BinOp, Body, Operand, Place, Rvalue, Statement,
5 StatementKind, TerminatorKind,
10 /// Pass to convert `if` conditions on integrals into switches on the integral.
11 /// For an example, it turns something like
14 /// _3 = Eq(move _4, const 43i32);
16 /// switchInt(_3) -> [false: bb2, otherwise: bb3];
22 /// switchInt(_4) -> [43i32: bb3, otherwise: bb2];
24 pub struct SimplifyComparisonIntegral;
26 impl<'tcx> MirPass<'tcx> for SimplifyComparisonIntegral {
27 fn run_pass(&self, _: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
28 trace!("Running SimplifyComparisonIntegral on {:?}", body.source);
30 let helper = OptimizationFinder { body };
31 let opts = helper.find_optimizations();
32 let mut storage_deads_to_insert = vec![];
33 let mut storage_deads_to_remove: Vec<(usize, BasicBlock)> = vec![];
35 trace!("SUCCESS: Applying {:?}", opt);
36 // replace terminator with a switchInt that switches on the integer directly
37 let bbs = &mut body.basic_blocks_mut();
38 let bb = &mut bbs[opt.bb_idx];
39 // We only use the bits for the untyped, not length checked `values` field. Thus we are
40 // not using any of the convenience wrappers here and directly access the bits.
41 let new_value = match opt.branch_value_scalar {
42 Scalar::Raw { data, .. } => data,
43 Scalar::Ptr(_) => continue,
45 const FALSE: u128 = 0;
46 let mut new_targets = opt.targets.clone();
47 let first_is_false_target = opt.values[0] == FALSE;
50 // if the assignment was Eq we want the true case to be first
51 if first_is_false_target {
52 new_targets.swap(0, 1);
56 // if the assignment was Ne we want the false case to be first
57 if !first_is_false_target {
58 new_targets.swap(0, 1);
64 // delete comparison statement if it the value being switched on was moved, which means it can not be user later on
65 if opt.can_remove_bin_op_stmt {
66 bb.statements[opt.bin_op_stmt_idx].make_nop();
68 // if the integer being compared to a const integral is being moved into the comparison,
69 // e.g `_2 = Eq(move _3, const 'x');`
70 // we want to avoid making a double move later on in the switchInt on _3.
71 // So to avoid `switchInt(move _3) -> ['x': bb2, otherwise: bb1];`,
72 // we convert the move in the comparison statement to a copy.
74 // unwrap is safe as we know this statement is an assign
75 let box (_, rhs) = bb.statements[opt.bin_op_stmt_idx].kind.as_assign_mut().unwrap();
79 Rvalue::BinaryOp(_, ref mut left @ Move(_), Constant(_)) => {
80 *left = Copy(opt.to_switch_on);
82 Rvalue::BinaryOp(_, Constant(_), ref mut right @ Move(_)) => {
83 *right = Copy(opt.to_switch_on);
89 let terminator = bb.terminator();
91 // remove StorageDead (if it exists) being used in the assign of the comparison
92 for (stmt_idx, stmt) in bb.statements.iter().enumerate() {
93 if !matches!(stmt.kind, StatementKind::StorageDead(local) if local == opt.to_switch_on.local)
97 storage_deads_to_remove.push((stmt_idx, opt.bb_idx));
98 // if we have StorageDeads to remove then make sure to insert them at the top of each target
99 for bb_idx in new_targets.iter() {
100 storage_deads_to_insert.push((
103 source_info: terminator.source_info,
104 kind: StatementKind::StorageDead(opt.to_switch_on.local),
110 let terminator = bb.terminator_mut();
112 terminator.kind = TerminatorKind::SwitchInt {
113 discr: Operand::Move(opt.to_switch_on),
114 switch_ty: opt.branch_value_ty,
115 values: vec![new_value].into(),
116 targets: new_targets,
120 for (idx, bb_idx) in storage_deads_to_remove {
121 body.basic_blocks_mut()[bb_idx].statements[idx].make_nop();
124 for (idx, stmt) in storage_deads_to_insert {
125 body.basic_blocks_mut()[idx].statements.insert(0, stmt);
130 struct OptimizationFinder<'a, 'tcx> {
131 body: &'a Body<'tcx>,
134 impl<'a, 'tcx> OptimizationFinder<'a, 'tcx> {
135 fn find_optimizations(&self) -> Vec<OptimizationInfo<'tcx>> {
139 .filter_map(|(bb_idx, bb)| {
141 let (place_switched_on, values, targets, place_switched_on_moved) = match &bb
145 rustc_middle::mir::TerminatorKind::SwitchInt {
146 discr, values, targets, ..
147 } => Some((discr.place()?, values, targets, discr.is_move())),
151 // find the statement that assigns the place being switched on
152 bb.statements.iter().enumerate().rev().find_map(|(stmt_idx, stmt)| {
154 rustc_middle::mir::StatementKind::Assign(box (lhs, rhs))
155 if *lhs == place_switched_on =>
158 Rvalue::BinaryOp(op @ (BinOp::Eq | BinOp::Ne), left, right) => {
159 let (branch_value_scalar, branch_value_ty, to_switch_on) =
160 find_branch_value_info(left, right)?;
162 Some(OptimizationInfo {
163 bin_op_stmt_idx: stmt_idx,
165 can_remove_bin_op_stmt: place_switched_on_moved,
170 values: values.clone().into_owned(),
171 targets: targets.clone(),
185 fn find_branch_value_info<'tcx>(
186 left: &Operand<'tcx>,
187 right: &Operand<'tcx>,
188 ) -> Option<(Scalar, Ty<'tcx>, Place<'tcx>)> {
189 // check that either left or right is a constant.
190 // if any are, we can use the other to switch on, and the constant as a value in a switch
192 match (left, right) {
193 (Constant(branch_value), Copy(to_switch_on) | Move(to_switch_on))
194 | (Copy(to_switch_on) | Move(to_switch_on), Constant(branch_value)) => {
195 let branch_value_ty = branch_value.literal.ty;
196 // we only want to apply this optimization if we are matching on integrals (and chars), as it is not possible to switch on floats
197 if !branch_value_ty.is_integral() && !branch_value_ty.is_char() {
200 let branch_value_scalar = branch_value.literal.val.try_to_scalar()?;
201 Some((branch_value_scalar, branch_value_ty, *to_switch_on))
208 struct OptimizationInfo<'tcx> {
209 /// Basic block to apply the optimization
211 /// Statement index of Eq/Ne assignment that can be removed. None if the assignment can not be removed - i.e the statement is used later on
212 bin_op_stmt_idx: usize,
213 /// Can remove Eq/Ne assignment
214 can_remove_bin_op_stmt: bool,
215 /// Place that needs to be switched on. This place is of type integral
216 to_switch_on: Place<'tcx>,
217 /// Constant to use in switch target value
218 branch_value_scalar: Scalar,
219 /// Type of the constant value
220 branch_value_ty: Ty<'tcx>,
223 /// Current values used in the switch target. This needs to be replaced with the branch_value
225 /// Current targets used in the switch
226 targets: Vec<BasicBlock>,