rustc_mir_transform/
simplify_branches.rs

1use rustc_middle::mir::*;
2use rustc_middle::ty::TyCtxt;
3use tracing::trace;
4
5use crate::patch::MirPatch;
6
7pub(super) enum SimplifyConstCondition {
8    AfterInstSimplify,
9    AfterConstProp,
10    Final,
11}
12
13/// A pass that replaces a branch with a goto when its condition is known.
14impl<'tcx> crate::MirPass<'tcx> for SimplifyConstCondition {
15    fn name(&self) -> &'static str {
16        match self {
17            SimplifyConstCondition::AfterInstSimplify => {
18                "SimplifyConstCondition-after-inst-simplify"
19            }
20            SimplifyConstCondition::AfterConstProp => "SimplifyConstCondition-after-const-prop",
21            SimplifyConstCondition::Final => "SimplifyConstCondition-final",
22        }
23    }
24
25    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
26        trace!("Running SimplifyConstCondition on {:?}", body.source);
27        let typing_env = body.typing_env(tcx);
28        let mut patch = MirPatch::new(body);
29
30        fn try_get_const<'tcx, 'a>(
31            operand: &'a Operand<'tcx>,
32            has_place_const: Option<(Place<'tcx>, &'a ConstOperand<'tcx>)>,
33        ) -> Option<&'a ConstOperand<'tcx>> {
34            match operand {
35                Operand::Constant(const_operand) => Some(const_operand),
36                // `has_place_const` must be the LHS of the previous statement.
37                // Soundness: There is nothing can modify the place, as there are no statements between the two statements.
38                Operand::Copy(place) | Operand::Move(place)
39                    if let Some((place_const, const_operand)) = has_place_const
40                        && place_const == *place =>
41                {
42                    Some(const_operand)
43                }
44                Operand::Copy(_) | Operand::Move(_) => None,
45            }
46        }
47
48        'blocks: for (bb, block) in body.basic_blocks.iter_enumerated() {
49            let mut pre_place_const: Option<(Place<'tcx>, &ConstOperand<'tcx>)> = None;
50
51            for (statement_index, stmt) in block.statements.iter().enumerate() {
52                let has_place_const = pre_place_const.take();
53                // Simplify `assume` of a known value: either a NOP or unreachable.
54                if let StatementKind::Intrinsic(box ref intrinsic) = stmt.kind
55                    && let NonDivergingIntrinsic::Assume(discr) = intrinsic
56                    && let Some(c) = try_get_const(discr, has_place_const)
57                    && let Some(constant) = c.const_.try_eval_bool(tcx, typing_env)
58                {
59                    if constant {
60                        patch.nop_statement(Location { block: bb, statement_index });
61                    } else {
62                        patch.patch_terminator(bb, TerminatorKind::Unreachable);
63                        continue 'blocks;
64                    }
65                } else if let StatementKind::Assign(box (lhs, ref rvalue)) = stmt.kind
66                    && let Rvalue::Use(Operand::Constant(c)) = rvalue
67                {
68                    pre_place_const = Some((lhs, c));
69                }
70            }
71
72            let terminator = block.terminator();
73            let terminator = match terminator.kind {
74                TerminatorKind::SwitchInt { ref discr, ref targets, .. }
75                    if let Some(c) = try_get_const(discr, pre_place_const.take())
76                        && let Some(constant) = c.const_.try_eval_bits(tcx, typing_env) =>
77                {
78                    let target = targets.target_for_value(constant);
79                    TerminatorKind::Goto { target }
80                }
81                TerminatorKind::Assert { target, ref cond, expected, .. }
82                    if let Some(c) = try_get_const(&cond, pre_place_const.take())
83                        && let Some(constant) = c.const_.try_eval_bool(tcx, typing_env)
84                        && constant == expected =>
85                {
86                    TerminatorKind::Goto { target }
87                }
88                _ => continue,
89            };
90            patch.patch_terminator(bb, terminator);
91        }
92        patch.apply(body);
93    }
94
95    fn is_required(&self) -> bool {
96        false
97    }
98}