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