rustc_mir_transform/
match_branches.rs

1use rustc_abi::Integer;
2use rustc_const_eval::const_eval::mk_eval_cx_for_const_val;
3use rustc_middle::mir::*;
4use rustc_middle::ty::layout::{IntegerExt, TyAndLayout};
5use rustc_middle::ty::util::Discr;
6use rustc_middle::ty::{self, ScalarInt, Ty, TyCtxt};
7
8use super::simplify::simplify_cfg;
9use crate::PassPolicy;
10use crate::patch::MirPatch;
11use crate::unreachable_prop::remove_successors_from_switch;
12
13/// Unifies all targets into one basic block if each statement can have the same statement.
14pub(super) struct MatchBranchSimplification;
15
16impl<'tcx> crate::MirPass<'tcx> for MatchBranchSimplification {
17    fn policy(&self, sess: &rustc_session::Session) -> PassPolicy {
18        // Enable only under -Zmir-opt-level=2 as this can make programs less debuggable.
19        PassPolicy::optimization(sess.mir_opt_level() >= 2)
20    }
21
22    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
23        let typing_env = body.typing_env(tcx);
24        let mut changed = false;
25        for bb in body.basic_blocks.indices() {
26            if !candidate_match(body, bb) {
27                continue;
28            };
29            changed |= simplify_match(tcx, typing_env, body, bb)
30        }
31
32        if changed {
33            simplify_cfg(tcx, body);
34        }
35    }
36}
37
38struct SimplifyMatch<'tcx, 'a> {
39    tcx: TyCtxt<'tcx>,
40    typing_env: ty::TypingEnv<'tcx>,
41    patch: MirPatch<'tcx>,
42    body: &'a Body<'tcx>,
43    switch_bb: BasicBlock,
44    discr: &'a Operand<'tcx>,
45    discr_local: Option<Local>,
46    discr_ty: Ty<'tcx>,
47}
48
49impl<'tcx, 'a> SimplifyMatch<'tcx, 'a> {
50    fn discr_local(&mut self) -> Local {
51        *self.discr_local.get_or_insert_with(|| {
52            // Introduce a temporary for the discriminant value.
53            let source_info = self.body.basic_blocks[self.switch_bb].terminator().source_info;
54            self.patch.new_temp(self.discr_ty, source_info.span)
55        })
56    }
57
58    /// Unifies the assignments if all rvalues are constants and equal.
59    fn unify_if_equal_const(
60        &self,
61        dest: Place<'tcx>,
62        consts: &[(u128, &ConstOperand<'tcx>)],
63        otherwise: Option<&ConstOperand<'tcx>>,
64    ) -> Option<StatementKind<'tcx>> {
65        let (_, first_const, mut others) = split_first_case(consts, otherwise);
66        let first_scalar_int = first_const.const_.try_eval_scalar_int(self.tcx, self.typing_env)?;
67        if others.all(|const_| {
68            const_.const_.try_eval_scalar_int(self.tcx, self.typing_env) == Some(first_scalar_int)
69        }) {
70            Some(StatementKind::Assign(Box::new((
71                dest,
72                // We didn't remember the `WithRetag` of the original assignments, so in case
73                // one of them had "no", we also have to use "no" here.
74                Rvalue::Use(Operand::Constant(Box::new(first_const.clone())), WithRetag::No),
75            ))))
76        } else {
77            None
78        }
79    }
80
81    /// If a source block is found that switches between two blocks that are exactly
82    /// the same modulo const bool assignments (e.g., one assigns true another false
83    /// to the same place), unify a target block statements into the source block,
84    /// using Eq / Ne comparison with switch value where const bools value differ.
85    ///
86    /// For example:
87    ///
88    /// ```ignore (MIR)
89    /// bb0: {
90    ///     switchInt(move _3) -> [42_isize: bb1, otherwise: bb2];
91    /// }
92    ///
93    /// bb1: {
94    ///     _2 = const true;
95    ///     goto -> bb3;
96    /// }
97    ///
98    /// bb2: {
99    ///     _2 = const false;
100    ///     goto -> bb3;
101    /// }
102    /// ```
103    ///
104    /// into:
105    ///
106    /// ```ignore (MIR)
107    /// bb0: {
108    ///    _2 = Eq(move _3, const 42_isize);
109    ///    goto -> bb3;
110    /// }
111    /// ```
112    fn unify_by_eq_op(
113        &mut self,
114        dest: Place<'tcx>,
115        consts: &[(u128, &ConstOperand<'tcx>)],
116        otherwise: Option<&ConstOperand<'tcx>>,
117    ) -> Option<StatementKind<'tcx>> {
118        // FIXME: extend to any case.
119        let (first_case, first_const, mut others) = split_first_case(consts, otherwise);
120        if !first_const.ty().is_bool() {
121            return None;
122        }
123        let first_bool = first_const.const_.try_eval_bool(self.tcx, self.typing_env)?;
124        if others.all(|const_| {
125            const_.const_.try_eval_bool(self.tcx, self.typing_env) == Some(!first_bool)
126        }) {
127            // Make value conditional on switch condition.
128            let size =
129                self.tcx.layout_of(self.typing_env.as_query_input(self.discr_ty)).unwrap().size;
130            let const_cmp = Operand::const_from_scalar(
131                self.tcx,
132                self.discr_ty,
133                rustc_const_eval::interpret::Scalar::from_uint(first_case, size),
134                rustc_span::DUMMY_SP,
135            );
136            let op = if first_bool { BinOp::Eq } else { BinOp::Ne };
137            let rval = Rvalue::BinaryOp(
138                op,
139                Box::new((Operand::Copy(Place::from(self.discr_local())), const_cmp)),
140            );
141            Some(StatementKind::Assign(Box::new((dest, rval))))
142        } else {
143            None
144        }
145    }
146
147    /// Unifies the assignments if all rvalues can be cast from the discriminant value by IntToInt.
148    ///
149    /// For example:
150    ///
151    /// ```ignore (MIR)
152    /// bb0: {
153    ///     switchInt(_1) -> [1: bb2, 2: bb3, 3: bb4, otherwise: bb1];
154    /// }
155    ///
156    /// bb1: {
157    ///     unreachable;
158    /// }
159    ///
160    /// bb2: {
161    ///     _0 = const 1_i16;
162    ///     goto -> bb5;
163    /// }
164    ///
165    /// bb3: {
166    ///     _0 = const 2_i16;
167    ///     goto -> bb5;
168    /// }
169    ///
170    /// bb4: {
171    ///     _0 = const 3_i16;
172    ///     goto -> bb5;
173    /// }
174    /// ```
175    ///
176    /// into:
177    ///
178    /// ```ignore (MIR)
179    /// bb0: {
180    ///    _0 = _1 as i16 (IntToInt);
181    ///    goto -> bb5;
182    /// }
183    /// ```
184    fn unify_by_int_to_int(
185        &mut self,
186        dest: Place<'tcx>,
187        consts: &[(u128, &ConstOperand<'tcx>)],
188    ) -> Option<StatementKind<'tcx>> {
189        let (_, first_const) = consts[0];
190        if !first_const.ty().is_integral() {
191            return None;
192        }
193        let discr_layout =
194            self.tcx.layout_of(self.typing_env.as_query_input(self.discr_ty)).unwrap();
195        if consts.iter().all(|&(case, const_)| {
196            let Some(scalar_int) = const_.const_.try_eval_scalar_int(self.tcx, self.typing_env)
197            else {
198                return false;
199            };
200            can_cast(self.tcx, case, discr_layout, const_.ty(), scalar_int)
201        }) {
202            let operand = Operand::Copy(Place::from(self.discr_local()));
203            let rval = if first_const.ty() == self.discr_ty {
204                Rvalue::Use(operand, WithRetag::No)
205            } else {
206                Rvalue::Cast(CastKind::IntToInt, operand, first_const.ty())
207            };
208            Some(StatementKind::Assign(Box::new((dest, rval))))
209        } else {
210            None
211        }
212    }
213
214    /// This is primarily used to unify these copy statements that simplified the canonical enum clone method by GVN.
215    /// The GVN simplified
216    /// ```ignore (syntax-highlighting-only)
217    /// match a {
218    ///     Foo::A(x) => Foo::A(*x),
219    ///     Foo::B => Foo::B
220    /// }
221    /// ```
222    /// to
223    /// ```ignore (syntax-highlighting-only)
224    /// match a {
225    ///     Foo::A(_x) => a, // copy a
226    ///     Foo::B => Foo::B
227    /// }
228    /// ```
229    /// This will simplify into a copy statement.
230    fn unify_by_copy(
231        &self,
232        dest: Place<'tcx>,
233        rvals: &[(u128, &Rvalue<'tcx>)],
234    ) -> Option<StatementKind<'tcx>> {
235        let bbs = &self.body.basic_blocks;
236        // Check if the copy source matches the following pattern.
237        // _2 = discriminant(*_1); // "*_1" is the expected the copy source.
238        // switchInt(move _2) -> [0: bb3, 1: bb2, otherwise: bb1];
239        let &Statement {
240            kind: StatementKind::Assign((discr_place, Rvalue::Discriminant(copy_src_place))),
241            ..
242        } = bbs[self.switch_bb].statements.last()?
243        else {
244            return None;
245        };
246        if self.discr.place() != Some(discr_place) {
247            return None;
248        }
249        let src_ty = copy_src_place.ty(self.body.local_decls(), self.tcx);
250        if !src_ty.ty.is_enum() || src_ty.variant_index.is_some() {
251            return None;
252        }
253        let dest_ty = dest.ty(self.body.local_decls(), self.tcx);
254        if dest_ty.ty != src_ty.ty || dest_ty.variant_index.is_some() {
255            return None;
256        }
257        let ty::Adt(def, _) = dest_ty.ty.kind() else {
258            return None;
259        };
260
261        for &(case, rvalue) in rvals.iter() {
262            match rvalue {
263                // Check if `_3 = const Foo::B` can be transformed to `_3 = copy *_1`.
264                Rvalue::Use(Operand::Constant(constant), _)
265                    if let Const::Val(const_, ty) = constant.const_ =>
266                {
267                    let (ecx, op) = mk_eval_cx_for_const_val(
268                        self.tcx.at(constant.span),
269                        self.typing_env,
270                        const_,
271                        ty,
272                    )?;
273                    let variant = ecx.read_discriminant(&op).discard_err()?;
274                    if !def.variants()[variant].fields.is_empty() {
275                        return None;
276                    }
277                    let Discr { val, .. } = ty.discriminant_for_variant(self.tcx, variant)?;
278                    if val != case {
279                        return None;
280                    }
281                }
282                Rvalue::Use(Operand::Copy(src_place), _) if *src_place == copy_src_place => {}
283                // Check if `_3 = Foo::B` can be transformed to `_3 = copy *_1`.
284                Rvalue::Aggregate(AggregateKind::Adt(_, variant_index, _, _, None), fields)
285                    if fields.is_empty()
286                        && let Some(Discr { val, .. }) =
287                            src_ty.ty.discriminant_for_variant(self.tcx, *variant_index)
288                        && val == case => {}
289                _ => return None,
290            }
291        }
292        // We didn't remember the `WithRetag` of the original assignments, so in case
293        // one of them had "no", we also have to use "no" here.
294        Some(StatementKind::Assign(Box::new((
295            dest,
296            Rvalue::Use(Operand::Copy(copy_src_place), WithRetag::No),
297        ))))
298    }
299
300    /// Returns a new statement if we can use the statement replace all statements.
301    fn try_unify_stmts(
302        &mut self,
303        index: usize,
304        stmts: &[(u128, &StatementKind<'tcx>)],
305        otherwise: Option<&StatementKind<'tcx>>,
306    ) -> Option<StatementKind<'tcx>> {
307        if let Some(new_stmt) = identical_stmts(stmts, otherwise) {
308            return Some(new_stmt);
309        }
310
311        let (dest, rvals, otherwise) = candidate_assign(stmts, otherwise)?;
312        if let Some((consts, otherwise)) = candidate_const(&rvals, otherwise) {
313            if let Some(new_stmt) = self.unify_if_equal_const(dest, &consts, otherwise) {
314                return Some(new_stmt);
315            }
316            if let Some(new_stmt) = self.unify_by_eq_op(dest, &consts, otherwise) {
317                return Some(new_stmt);
318            }
319            // Requires the otherwise is unreachable.
320            if otherwise.is_none()
321                && let Some(new_stmt) = self.unify_by_int_to_int(dest, &consts)
322            {
323                return Some(new_stmt);
324            }
325        }
326
327        // We only know the first statement is safe to introduce new dereferences.
328        if index == 0
329            // We cannot create overlapping assignments.
330            && dest.is_stable_offset()
331            // Requires the otherwise is unreachable.
332            && otherwise.is_none()
333            && let Some(new_stmt) = self.unify_by_copy(dest, &rvals)
334        {
335            return Some(new_stmt);
336        }
337        None
338    }
339}
340
341/// Returns the first case target if all targets have an equal number of statements and identical destination.
342fn candidate_match<'tcx>(body: &Body<'tcx>, switch_bb: BasicBlock) -> bool {
343    use itertools::Itertools;
344    let targets = match &body.basic_blocks[switch_bb].terminator().kind {
345        TerminatorKind::SwitchInt {
346            discr: Operand::Copy(_) | Operand::Move(_), targets, ..
347        } => targets,
348        // Only optimize switch int statements
349        _ => return false,
350    };
351    // We require that the possible target blocks don't contain this block.
352    if targets.all_targets().contains(&switch_bb) {
353        return false;
354    }
355    // We require that the possible target blocks all be distinct.
356    if !targets.is_distinct() {
357        return false;
358    }
359    // Check that destinations are identical, and if not, then don't optimize this block
360    targets
361        .all_targets()
362        .iter()
363        .map(|&bb| &body.basic_blocks[bb])
364        .filter(|bb| !bb.is_empty_unreachable())
365        .map(|bb| (bb.statements.len(), &bb.terminator().kind))
366        .all_equal()
367}
368
369fn simplify_match<'tcx>(
370    tcx: TyCtxt<'tcx>,
371    typing_env: ty::TypingEnv<'tcx>,
372    body: &mut Body<'tcx>,
373    switch_bb: BasicBlock,
374) -> bool {
375    let (discr, targets) = match &body.basic_blocks[switch_bb].terminator().kind {
376        TerminatorKind::SwitchInt { discr, targets, .. } => (discr, targets),
377        _ => unreachable!(),
378    };
379    let mut simplify_match = SimplifyMatch {
380        tcx,
381        typing_env,
382        patch: MirPatch::new(body),
383        body,
384        switch_bb,
385        discr,
386        discr_local: None,
387        discr_ty: discr.ty(body.local_decls(), tcx),
388    };
389    let reachable_cases: Vec<_> =
390        targets.iter().filter(|&(_, bb)| !body.basic_blocks[bb].is_empty_unreachable()).collect();
391    let mut new_stmts = Vec::new();
392    let otherwise = if body.basic_blocks[targets.otherwise()].is_empty_unreachable() {
393        None
394    } else {
395        Some(targets.otherwise())
396    };
397    // We can patch the terminator to goto because there is a single target.
398    match (reachable_cases.len(), otherwise.is_none()) {
399        (1, true) | (0, false) => {
400            let mut patch = simplify_match.patch;
401            remove_successors_from_switch(tcx, switch_bb, body, &mut patch, |bb| {
402                body.basic_blocks[bb].is_empty_unreachable()
403            });
404            patch.apply(body);
405            return true;
406        }
407        _ => {}
408    }
409    let Some(&(_, first_case_bb)) = reachable_cases.first() else {
410        return false;
411    };
412    let stmt_len = body.basic_blocks[first_case_bb].statements.len();
413    let mut cases = Vec::with_capacity(stmt_len);
414    // Check at each position in the basic blocks whether these statements can be unified.
415    for index in 0..stmt_len {
416        cases.clear();
417        let otherwise = otherwise.map(|bb| &body.basic_blocks[bb].statements[index].kind);
418        for &(case, bb) in &reachable_cases {
419            cases.push((case, &body.basic_blocks[bb].statements[index].kind));
420        }
421        let Some(new_stmt) = simplify_match.try_unify_stmts(index, &cases, otherwise) else {
422            return false;
423        };
424        new_stmts.push(new_stmt);
425    }
426    // Take ownership of items now that we know we can optimize.
427    let discr = discr.clone();
428
429    let statement_index = body.basic_blocks[switch_bb].statements.len();
430    let parent_end = Location { block: switch_bb, statement_index };
431    let mut patch = simplify_match.patch;
432    if let Some(discr_local) = simplify_match.discr_local {
433        patch.add_statement(parent_end, StatementKind::StorageLive(discr_local));
434        patch.add_assign(parent_end, Place::from(discr_local), Rvalue::Use(discr, WithRetag::No));
435    }
436    for new_stmt in new_stmts {
437        patch.add_statement(parent_end, new_stmt);
438    }
439    if let Some(discr_local) = simplify_match.discr_local {
440        patch.add_statement(parent_end, StatementKind::StorageDead(discr_local));
441    }
442    patch.patch_terminator(switch_bb, body.basic_blocks[first_case_bb].terminator().kind.clone());
443    patch.apply(body);
444    true
445}
446
447/// Check if the cast constant using `IntToInt` is equal to the target constant.
448fn can_cast(
449    tcx: TyCtxt<'_>,
450    src_val: impl Into<u128>,
451    src_layout: TyAndLayout<'_>,
452    cast_ty: Ty<'_>,
453    target_scalar: ScalarInt,
454) -> bool {
455    let from_scalar = ScalarInt::try_from_uint(src_val.into(), src_layout.size).unwrap();
456    let v = match src_layout.ty.kind() {
457        ty::Uint(_) => from_scalar.to_uint(src_layout.size),
458        ty::Int(_) => from_scalar.to_int(src_layout.size) as u128,
459        // We can also transform the values of other integer representations (such as char),
460        // although this may not be practical in real-world scenarios.
461        _ => return false,
462    };
463    let size = match *cast_ty.kind() {
464        ty::Int(t) => Integer::from_int_ty(&tcx, t).size(),
465        ty::Uint(t) => Integer::from_uint_ty(&tcx, t).size(),
466        _ => return false,
467    };
468    let v = size.truncate(v);
469    let cast_scalar = ScalarInt::try_from_uint(v, size).unwrap();
470    cast_scalar == target_scalar
471}
472
473fn candidate_assign<'tcx, 'a>(
474    stmts: &'a [(u128, &'a StatementKind<'tcx>)],
475    otherwise: Option<&'a StatementKind<'tcx>>,
476) -> Option<(Place<'tcx>, Vec<(u128, &'a Rvalue<'tcx>)>, Option<&'a Rvalue<'tcx>>)> {
477    let (_, first_stmt) = stmts[0];
478    let (dest, _) = first_stmt.as_assign()?;
479    let otherwise = if let Some(otherwise) = otherwise {
480        let Some((otherwise_dest, rval)) = otherwise.as_assign() else {
481            return None;
482        };
483        if otherwise_dest != dest {
484            return None;
485        }
486        Some(rval)
487    } else {
488        None
489    };
490    let rvals = stmts
491        .into_iter()
492        .map(|&(case, stmt)| {
493            let (other_dest, rval) = stmt.as_assign()?;
494            if other_dest != dest {
495                return None;
496            }
497            Some((case, rval))
498        })
499        .try_collect()?;
500    Some((*dest, rvals, otherwise))
501}
502
503// Returns all ConstOperands if all Rvalues are ConstOperands.
504fn candidate_const<'tcx, 'a>(
505    rvals: &'a [(u128, &'a Rvalue<'tcx>)],
506    otherwise: Option<&'a Rvalue<'tcx>>,
507) -> Option<(Vec<(u128, &'a ConstOperand<'tcx>)>, Option<&'a ConstOperand<'tcx>>)> {
508    // We ignore the retag mode here, which means the `Use` we insert later must be without retag.
509    let otherwise = if let Some(otherwise) = otherwise {
510        let Rvalue::Use(Operand::Constant(const_), _) = otherwise else {
511            return None;
512        };
513        Some(&**const_)
514    } else {
515        None
516    };
517    let consts = rvals
518        .into_iter()
519        .map(|&(case, rval)| {
520            let Rvalue::Use(Operand::Constant(const_), _) = rval else { return None };
521            Some((case, &**const_))
522        })
523        .try_collect()?;
524    Some((consts, otherwise))
525}
526
527// Returns the first case and others (including otherwise if present).
528fn split_first_case<'a, T>(
529    stmts: &'a [(u128, &'a T)],
530    otherwise: Option<&'a T>,
531) -> (u128, &'a T, impl Iterator<Item = &'a T>) {
532    let (first_case, first) = stmts[0];
533    (first_case, first, stmts[1..].into_iter().map(|&(_, val)| val).chain(otherwise))
534}
535
536// If all statements are identical, we can optimize.
537fn identical_stmts<'tcx>(
538    stmts: &[(u128, &StatementKind<'tcx>)],
539    otherwise: Option<&StatementKind<'tcx>>,
540) -> Option<StatementKind<'tcx>> {
541    use itertools::Itertools;
542    let (_, first_stmt, others) = split_first_case(stmts, otherwise);
543    if std::iter::once(first_stmt).chain(others).all_equal() {
544        return Some(first_stmt.clone());
545    }
546    None
547}
rustc_mir_transform/match_branches.rs

rustc_mir_transform/
match_branches.rs
