rustc_mir_transform/
match_branches.rs

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

rustc_mir_transform/
match_branches.rs
