rustc_mir_transform/
simplify_branches.rs

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

rustc_mir_transform/
simplify_branches.rs
