]>
Commit | Line | Data |
---|---|---|
29967ef6 XL |
1 | use std::iter; |
2 | ||
3 | use super::MirPass; | |
1b1a35ee XL |
4 | use rustc_middle::{ |
5 | mir::{ | |
6 | interpret::Scalar, BasicBlock, BinOp, Body, Operand, Place, Rvalue, Statement, | |
29967ef6 | 7 | StatementKind, SwitchTargets, TerminatorKind, |
1b1a35ee XL |
8 | }, |
9 | ty::{Ty, TyCtxt}, | |
10 | }; | |
11 | ||
12 | /// Pass to convert `if` conditions on integrals into switches on the integral. | |
13 | /// For an example, it turns something like | |
14 | /// | |
04454e1e | 15 | /// ```ignore (MIR) |
1b1a35ee XL |
16 | /// _3 = Eq(move _4, const 43i32); |
17 | /// StorageDead(_4); | |
18 | /// switchInt(_3) -> [false: bb2, otherwise: bb3]; | |
19 | /// ``` | |
20 | /// | |
21 | /// into: | |
22 | /// | |
04454e1e | 23 | /// ```ignore (MIR) |
1b1a35ee XL |
24 | /// switchInt(_4) -> [43i32: bb3, otherwise: bb2]; |
25 | /// ``` | |
26 | pub struct SimplifyComparisonIntegral; | |
27 | ||
28 | impl<'tcx> MirPass<'tcx> for SimplifyComparisonIntegral { | |
a2a8927a XL |
29 | fn is_enabled(&self, sess: &rustc_session::Session) -> bool { |
30 | sess.mir_opt_level() > 0 | |
31 | } | |
32 | ||
29967ef6 XL |
33 | fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
34 | trace!("Running SimplifyComparisonIntegral on {:?}", body.source); | |
1b1a35ee XL |
35 | |
36 | let helper = OptimizationFinder { body }; | |
37 | let opts = helper.find_optimizations(); | |
38 | let mut storage_deads_to_insert = vec![]; | |
39 | let mut storage_deads_to_remove: Vec<(usize, BasicBlock)> = vec![]; | |
29967ef6 | 40 | let param_env = tcx.param_env(body.source.def_id()); |
1b1a35ee XL |
41 | for opt in opts { |
42 | trace!("SUCCESS: Applying {:?}", opt); | |
43 | // replace terminator with a switchInt that switches on the integer directly | |
44 | let bbs = &mut body.basic_blocks_mut(); | |
45 | let bb = &mut bbs[opt.bb_idx]; | |
1b1a35ee | 46 | let new_value = match opt.branch_value_scalar { |
29967ef6 XL |
47 | Scalar::Int(int) => { |
48 | let layout = tcx | |
49 | .layout_of(param_env.and(opt.branch_value_ty)) | |
50 | .expect("if we have an evaluated constant we must know the layout"); | |
51 | int.assert_bits(layout.size) | |
52 | } | |
136023e0 | 53 | Scalar::Ptr(..) => continue, |
1b1a35ee XL |
54 | }; |
55 | const FALSE: u128 = 0; | |
29967ef6 XL |
56 | |
57 | let mut new_targets = opt.targets; | |
58 | let first_value = new_targets.iter().next().unwrap().0; | |
59 | let first_is_false_target = first_value == FALSE; | |
1b1a35ee XL |
60 | match opt.op { |
61 | BinOp::Eq => { | |
62 | // if the assignment was Eq we want the true case to be first | |
63 | if first_is_false_target { | |
29967ef6 | 64 | new_targets.all_targets_mut().swap(0, 1); |
1b1a35ee XL |
65 | } |
66 | } | |
67 | BinOp::Ne => { | |
68 | // if the assignment was Ne we want the false case to be first | |
69 | if !first_is_false_target { | |
29967ef6 | 70 | new_targets.all_targets_mut().swap(0, 1); |
1b1a35ee XL |
71 | } |
72 | } | |
73 | _ => unreachable!(), | |
74 | } | |
75 | ||
76 | // delete comparison statement if it the value being switched on was moved, which means it can not be user later on | |
77 | if opt.can_remove_bin_op_stmt { | |
78 | bb.statements[opt.bin_op_stmt_idx].make_nop(); | |
79 | } else { | |
80 | // if the integer being compared to a const integral is being moved into the comparison, | |
81 | // e.g `_2 = Eq(move _3, const 'x');` | |
82 | // we want to avoid making a double move later on in the switchInt on _3. | |
83 | // So to avoid `switchInt(move _3) -> ['x': bb2, otherwise: bb1];`, | |
84 | // we convert the move in the comparison statement to a copy. | |
85 | ||
86 | // unwrap is safe as we know this statement is an assign | |
6a06907d | 87 | let (_, rhs) = bb.statements[opt.bin_op_stmt_idx].kind.as_assign_mut().unwrap(); |
1b1a35ee XL |
88 | |
89 | use Operand::*; | |
90 | match rhs { | |
6a06907d | 91 | Rvalue::BinaryOp(_, box (ref mut left @ Move(_), Constant(_))) => { |
1b1a35ee XL |
92 | *left = Copy(opt.to_switch_on); |
93 | } | |
6a06907d | 94 | Rvalue::BinaryOp(_, box (Constant(_), ref mut right @ Move(_))) => { |
1b1a35ee XL |
95 | *right = Copy(opt.to_switch_on); |
96 | } | |
97 | _ => (), | |
98 | } | |
99 | } | |
100 | ||
101 | let terminator = bb.terminator(); | |
102 | ||
103 | // remove StorageDead (if it exists) being used in the assign of the comparison | |
104 | for (stmt_idx, stmt) in bb.statements.iter().enumerate() { | |
105 | if !matches!(stmt.kind, StatementKind::StorageDead(local) if local == opt.to_switch_on.local) | |
106 | { | |
107 | continue; | |
108 | } | |
109 | storage_deads_to_remove.push((stmt_idx, opt.bb_idx)); | |
110 | // if we have StorageDeads to remove then make sure to insert them at the top of each target | |
29967ef6 | 111 | for bb_idx in new_targets.all_targets() { |
1b1a35ee XL |
112 | storage_deads_to_insert.push(( |
113 | *bb_idx, | |
114 | Statement { | |
115 | source_info: terminator.source_info, | |
116 | kind: StatementKind::StorageDead(opt.to_switch_on.local), | |
117 | }, | |
118 | )); | |
119 | } | |
120 | } | |
121 | ||
29967ef6 XL |
122 | let [bb_cond, bb_otherwise] = match new_targets.all_targets() { |
123 | [a, b] => [*a, *b], | |
124 | e => bug!("expected 2 switch targets, got: {:?}", e), | |
125 | }; | |
1b1a35ee | 126 | |
29967ef6 XL |
127 | let targets = SwitchTargets::new(iter::once((new_value, bb_cond)), bb_otherwise); |
128 | ||
129 | let terminator = bb.terminator_mut(); | |
9c376795 FG |
130 | terminator.kind = |
131 | TerminatorKind::SwitchInt { discr: Operand::Move(opt.to_switch_on), targets }; | |
1b1a35ee XL |
132 | } |
133 | ||
134 | for (idx, bb_idx) in storage_deads_to_remove { | |
135 | body.basic_blocks_mut()[bb_idx].statements[idx].make_nop(); | |
136 | } | |
137 | ||
138 | for (idx, stmt) in storage_deads_to_insert { | |
139 | body.basic_blocks_mut()[idx].statements.insert(0, stmt); | |
140 | } | |
141 | } | |
142 | } | |
143 | ||
144 | struct OptimizationFinder<'a, 'tcx> { | |
145 | body: &'a Body<'tcx>, | |
146 | } | |
147 | ||
a2a8927a | 148 | impl<'tcx> OptimizationFinder<'_, 'tcx> { |
1b1a35ee XL |
149 | fn find_optimizations(&self) -> Vec<OptimizationInfo<'tcx>> { |
150 | self.body | |
f2b60f7d | 151 | .basic_blocks |
1b1a35ee XL |
152 | .iter_enumerated() |
153 | .filter_map(|(bb_idx, bb)| { | |
154 | // find switch | |
29967ef6 XL |
155 | let (place_switched_on, targets, place_switched_on_moved) = |
156 | match &bb.terminator().kind { | |
157 | rustc_middle::mir::TerminatorKind::SwitchInt { discr, targets, .. } => { | |
158 | Some((discr.place()?, targets, discr.is_move())) | |
159 | } | |
160 | _ => None, | |
161 | }?; | |
1b1a35ee XL |
162 | |
163 | // find the statement that assigns the place being switched on | |
164 | bb.statements.iter().enumerate().rev().find_map(|(stmt_idx, stmt)| { | |
165 | match &stmt.kind { | |
166 | rustc_middle::mir::StatementKind::Assign(box (lhs, rhs)) | |
167 | if *lhs == place_switched_on => | |
168 | { | |
169 | match rhs { | |
6a06907d XL |
170 | Rvalue::BinaryOp( |
171 | op @ (BinOp::Eq | BinOp::Ne), | |
172 | box (left, right), | |
173 | ) => { | |
1b1a35ee XL |
174 | let (branch_value_scalar, branch_value_ty, to_switch_on) = |
175 | find_branch_value_info(left, right)?; | |
176 | ||
177 | Some(OptimizationInfo { | |
178 | bin_op_stmt_idx: stmt_idx, | |
179 | bb_idx, | |
180 | can_remove_bin_op_stmt: place_switched_on_moved, | |
181 | to_switch_on, | |
182 | branch_value_scalar, | |
183 | branch_value_ty, | |
184 | op: *op, | |
1b1a35ee XL |
185 | targets: targets.clone(), |
186 | }) | |
187 | } | |
188 | _ => None, | |
189 | } | |
190 | } | |
191 | _ => None, | |
192 | } | |
193 | }) | |
194 | }) | |
195 | .collect() | |
196 | } | |
197 | } | |
198 | ||
199 | fn find_branch_value_info<'tcx>( | |
200 | left: &Operand<'tcx>, | |
201 | right: &Operand<'tcx>, | |
202 | ) -> Option<(Scalar, Ty<'tcx>, Place<'tcx>)> { | |
203 | // check that either left or right is a constant. | |
204 | // if any are, we can use the other to switch on, and the constant as a value in a switch | |
205 | use Operand::*; | |
206 | match (left, right) { | |
207 | (Constant(branch_value), Copy(to_switch_on) | Move(to_switch_on)) | |
208 | | (Copy(to_switch_on) | Move(to_switch_on), Constant(branch_value)) => { | |
6a06907d | 209 | let branch_value_ty = branch_value.literal.ty(); |
1b1a35ee XL |
210 | // we only want to apply this optimization if we are matching on integrals (and chars), as it is not possible to switch on floats |
211 | if !branch_value_ty.is_integral() && !branch_value_ty.is_char() { | |
212 | return None; | |
213 | }; | |
6a06907d | 214 | let branch_value_scalar = branch_value.literal.try_to_scalar()?; |
94222f64 | 215 | Some((branch_value_scalar, branch_value_ty, *to_switch_on)) |
1b1a35ee XL |
216 | } |
217 | _ => None, | |
218 | } | |
219 | } | |
220 | ||
221 | #[derive(Debug)] | |
222 | struct OptimizationInfo<'tcx> { | |
223 | /// Basic block to apply the optimization | |
224 | bb_idx: BasicBlock, | |
225 | /// 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 | |
226 | bin_op_stmt_idx: usize, | |
227 | /// Can remove Eq/Ne assignment | |
228 | can_remove_bin_op_stmt: bool, | |
229 | /// Place that needs to be switched on. This place is of type integral | |
230 | to_switch_on: Place<'tcx>, | |
231 | /// Constant to use in switch target value | |
232 | branch_value_scalar: Scalar, | |
233 | /// Type of the constant value | |
234 | branch_value_ty: Ty<'tcx>, | |
235 | /// Either Eq or Ne | |
236 | op: BinOp, | |
1b1a35ee | 237 | /// Current targets used in the switch |
29967ef6 | 238 | targets: SwitchTargets, |
1b1a35ee | 239 | } |