rustc_mir_transform/
simplify_comparison_integral.rs

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

rustc_mir_transform/
simplify_comparison_integral.rs
