]>
Commit | Line | Data |
---|---|---|
3dfed10e XL |
1 | use crate::transform::{MirPass, MirSource}; |
2 | use rustc_middle::mir::*; | |
3 | use rustc_middle::ty::TyCtxt; | |
4 | ||
5 | pub struct MatchBranchSimplification; | |
6 | ||
7 | /// If a source block is found that switches between two blocks that are exactly | |
8 | /// the same modulo const bool assignments (e.g., one assigns true another false | |
9 | /// to the same place), merge a target block statements into the source block, | |
10 | /// using Eq / Ne comparison with switch value where const bools value differ. | |
11 | /// | |
12 | /// For example: | |
13 | /// | |
14 | /// ```rust | |
15 | /// bb0: { | |
16 | /// switchInt(move _3) -> [42_isize: bb1, otherwise: bb2]; | |
17 | /// } | |
18 | /// | |
19 | /// bb1: { | |
20 | /// _2 = const true; | |
21 | /// goto -> bb3; | |
22 | /// } | |
23 | /// | |
24 | /// bb2: { | |
25 | /// _2 = const false; | |
26 | /// goto -> bb3; | |
27 | /// } | |
28 | /// ``` | |
29 | /// | |
30 | /// into: | |
31 | /// | |
32 | /// ```rust | |
33 | /// bb0: { | |
34 | /// _2 = Eq(move _3, const 42_isize); | |
35 | /// goto -> bb3; | |
36 | /// } | |
37 | /// ``` | |
38 | ||
39 | impl<'tcx> MirPass<'tcx> for MatchBranchSimplification { | |
40 | fn run_pass(&self, tcx: TyCtxt<'tcx>, src: MirSource<'tcx>, body: &mut Body<'tcx>) { | |
41 | let param_env = tcx.param_env(src.def_id()); | |
42 | let bbs = body.basic_blocks_mut(); | |
43 | 'outer: for bb_idx in bbs.indices() { | |
44 | let (discr, val, switch_ty, first, second) = match bbs[bb_idx].terminator().kind { | |
45 | TerminatorKind::SwitchInt { | |
46 | discr: Operand::Copy(ref place) | Operand::Move(ref place), | |
47 | switch_ty, | |
48 | ref targets, | |
49 | ref values, | |
50 | .. | |
51 | } if targets.len() == 2 && values.len() == 1 && targets[0] != targets[1] => { | |
52 | (place, values[0], switch_ty, targets[0], targets[1]) | |
53 | } | |
54 | // Only optimize switch int statements | |
55 | _ => continue, | |
56 | }; | |
57 | ||
58 | // Check that destinations are identical, and if not, then don't optimize this block | |
59 | if &bbs[first].terminator().kind != &bbs[second].terminator().kind { | |
60 | continue; | |
61 | } | |
62 | ||
63 | // Check that blocks are assignments of consts to the same place or same statement, | |
64 | // and match up 1-1, if not don't optimize this block. | |
65 | let first_stmts = &bbs[first].statements; | |
66 | let scnd_stmts = &bbs[second].statements; | |
67 | if first_stmts.len() != scnd_stmts.len() { | |
68 | continue; | |
69 | } | |
70 | for (f, s) in first_stmts.iter().zip(scnd_stmts.iter()) { | |
71 | match (&f.kind, &s.kind) { | |
72 | // If two statements are exactly the same, we can optimize. | |
73 | (f_s, s_s) if f_s == s_s => {} | |
74 | ||
75 | // If two statements are const bool assignments to the same place, we can optimize. | |
76 | ( | |
77 | StatementKind::Assign(box (lhs_f, Rvalue::Use(Operand::Constant(f_c)))), | |
78 | StatementKind::Assign(box (lhs_s, Rvalue::Use(Operand::Constant(s_c)))), | |
79 | ) if lhs_f == lhs_s | |
80 | && f_c.literal.ty.is_bool() | |
81 | && s_c.literal.ty.is_bool() | |
82 | && f_c.literal.try_eval_bool(tcx, param_env).is_some() | |
83 | && s_c.literal.try_eval_bool(tcx, param_env).is_some() => {} | |
84 | ||
85 | // Otherwise we cannot optimize. Try another block. | |
86 | _ => continue 'outer, | |
87 | } | |
88 | } | |
89 | // Take ownership of items now that we know we can optimize. | |
90 | let discr = discr.clone(); | |
91 | ||
92 | // We already checked that first and second are different blocks, | |
93 | // and bb_idx has a different terminator from both of them. | |
94 | let (from, first, second) = bbs.pick3_mut(bb_idx, first, second); | |
95 | ||
96 | let new_stmts = first.statements.iter().zip(second.statements.iter()).map(|(f, s)| { | |
97 | match (&f.kind, &s.kind) { | |
98 | (f_s, s_s) if f_s == s_s => (*f).clone(), | |
99 | ||
100 | ( | |
101 | StatementKind::Assign(box (lhs, Rvalue::Use(Operand::Constant(f_c)))), | |
102 | StatementKind::Assign(box (_, Rvalue::Use(Operand::Constant(s_c)))), | |
103 | ) => { | |
104 | // From earlier loop we know that we are dealing with bool constants only: | |
105 | let f_b = f_c.literal.try_eval_bool(tcx, param_env).unwrap(); | |
106 | let s_b = s_c.literal.try_eval_bool(tcx, param_env).unwrap(); | |
107 | if f_b == s_b { | |
108 | // Same value in both blocks. Use statement as is. | |
109 | (*f).clone() | |
110 | } else { | |
111 | // Different value between blocks. Make value conditional on switch condition. | |
112 | let size = tcx.layout_of(param_env.and(switch_ty)).unwrap().size; | |
113 | let const_cmp = Operand::const_from_scalar( | |
114 | tcx, | |
115 | switch_ty, | |
116 | crate::interpret::Scalar::from_uint(val, size), | |
117 | rustc_span::DUMMY_SP, | |
118 | ); | |
119 | let op = if f_b { BinOp::Eq } else { BinOp::Ne }; | |
120 | let rhs = Rvalue::BinaryOp(op, Operand::Copy(discr.clone()), const_cmp); | |
121 | Statement { | |
122 | source_info: f.source_info, | |
123 | kind: StatementKind::Assign(box (*lhs, rhs)), | |
124 | } | |
125 | } | |
126 | } | |
127 | ||
128 | _ => unreachable!(), | |
129 | } | |
130 | }); | |
131 | from.statements.extend(new_stmts); | |
132 | from.terminator_mut().kind = first.terminator().kind.clone(); | |
133 | } | |
134 | } | |
135 | } |