rustc_mir_transform/
ssa_range_prop.rs

1//! A pass that propagates the known ranges of SSA locals.
2//! We can know the ranges of SSA locals in certain locations for the following code:
3//! ```
4//! fn foo(a: u32) {
5//!   let b = a < 9; // the integer representation of b is within the full range [0, 2).
6//!   if b {
7//!     let c = b; // c is true since b is within the range [1, 2).
8//!     let d = a < 8; // d is true since a is within the range [0, 9).
9//!   }
10//! }
11//! ```
12use rustc_abi::WrappingRange;
13use rustc_const_eval::interpret::Scalar;
14use rustc_data_structures::fx::FxHashMap;
15use rustc_data_structures::graph::dominators::Dominators;
16use rustc_middle::mir::basic_blocks::Predecessors;
17use rustc_middle::mir::visit::MutVisitor;
18use rustc_middle::mir::{BasicBlock, Body, Location, Operand, Place, TerminatorKind, *};
19use rustc_middle::ty::{TyCtxt, TypingEnv};
20use rustc_span::DUMMY_SP;
21
22use crate::PassPolicy;
23use crate::ssa::SsaLocals;
24
25pub(super) struct SsaRangePropagation;
26
27impl<'tcx> crate::MirPass<'tcx> for SsaRangePropagation {
28    fn policy(&self, ctx: &crate::PassCtx<'_>) -> PassPolicy {
29        PassPolicy::optional(ctx.mir_opt_level() >= 2)
30    }
31
32    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
33        let typing_env = body.typing_env(tcx);
34        let ssa = SsaLocals::new(tcx, body, typing_env);
35        let cfg_cache = body.basic_blocks.cfg_cache();
36        let dominators = cfg_cache.dominators(&body.basic_blocks);
37        let predecessors = cfg_cache.predecessors(&body.basic_blocks);
38        let mut range_set =
39            RangeSet::new(tcx, typing_env, &ssa, &body.local_decls, dominators, predecessors);
40
41        for &bb in cfg_cache.reverse_postorder(&body.basic_blocks) {
42            let data = &mut body.basic_blocks.as_mut_preserves_cfg()[bb];
43            range_set.visit_basic_block_data(bb, data);
44        }
45    }
46}
47
48struct RangeSet<'tcx, 'body, 'a> {
49    tcx: TyCtxt<'tcx>,
50    typing_env: TypingEnv<'tcx>,
51    ssa: &'a SsaLocals,
52    local_decls: &'body LocalDecls<'tcx>,
53    dominators: &'a Dominators<BasicBlock>,
54    /// Known ranges at each locations.
55    ranges: FxHashMap<Place<'tcx>, Vec<(Location, WrappingRange)>>,
56    predecessors: &'a Predecessors,
57}
58
59impl<'tcx, 'body, 'a> RangeSet<'tcx, 'body, 'a> {
60    fn new(
61        tcx: TyCtxt<'tcx>,
62        typing_env: TypingEnv<'tcx>,
63        ssa: &'a SsaLocals,
64        local_decls: &'body LocalDecls<'tcx>,
65        dominators: &'a Dominators<BasicBlock>,
66        predecessors: &'a Predecessors,
67    ) -> Self {
68        RangeSet {
69            tcx,
70            typing_env,
71            ssa,
72            local_decls,
73            dominators,
74            ranges: FxHashMap::default(),
75            predecessors,
76        }
77    }
78
79    fn has_unique_predecessor(&self, bb: BasicBlock) -> bool {
80        self.predecessors[bb].len() == 1
81    }
82
83    /// Create a new known range at the location.
84    fn insert_range(&mut self, place: Place<'tcx>, location: Location, range: WrappingRange) {
85        assert!(self.is_ssa(place));
86        self.ranges.entry(place).or_default().push((location, range));
87    }
88
89    /// Get the known range at the location.
90    fn get_range(&self, place: &Place<'tcx>, location: Location) -> Option<WrappingRange> {
91        let Some(ranges) = self.ranges.get(place) else {
92            return None;
93        };
94        // FIXME: This should use the intersection of all valid ranges.
95        let (_, range) =
96            ranges.iter().find(|(range_loc, _)| range_loc.dominates(location, &self.dominators))?;
97        Some(*range)
98    }
99
100    fn try_as_constant(
101        &mut self,
102        place: Place<'tcx>,
103        location: Location,
104    ) -> Option<ConstOperand<'tcx>> {
105        if let Some(range) = self.get_range(&place, location)
106            && range.start == range.end
107        {
108            let ty = place.ty(self.local_decls, self.tcx).ty;
109            let layout = self.tcx.layout_of(self.typing_env.as_query_input(ty)).ok()?;
110            let value = ConstValue::Scalar(Scalar::from_uint(range.start, layout.size));
111            let const_ = Const::Val(value, ty);
112            return Some(ConstOperand { span: DUMMY_SP, user_ty: None, const_ });
113        }
114        None
115    }
116
117    fn is_ssa(&self, place: Place<'tcx>) -> bool {
118        self.ssa.is_ssa(place.local) && place.is_stable_offset()
119    }
120}
121
122impl<'tcx> MutVisitor<'tcx> for RangeSet<'tcx, '_, '_> {
123    fn tcx(&self) -> TyCtxt<'tcx> {
124        self.tcx
125    }
126
127    fn visit_operand(&mut self, operand: &mut Operand<'tcx>, location: Location) {
128        // Attempts to simplify an operand to a constant value.
129        if let Some(place) = operand.place()
130            && let Some(const_) = self.try_as_constant(place, location)
131        {
132            *operand = Operand::Constant(Box::new(const_));
133        };
134    }
135
136    fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
137        self.super_statement(statement, location);
138        match &statement.kind {
139            StatementKind::Intrinsic(NonDivergingIntrinsic::Assume(operand))
140                if let Some(place) = operand.place()
141                    && self.is_ssa(place) =>
142            {
143                let successor = location.successor_within_block();
144                let range = WrappingRange { start: 1, end: 1 };
145                self.insert_range(place, successor, range);
146            }
147            _ => {}
148        }
149    }
150
151    fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) {
152        self.super_terminator(terminator, location);
153        match &terminator.kind {
154            TerminatorKind::Assert { cond, expected, target, msg, .. }
155                if let Some(place) = cond.place()
156                    && self.is_ssa(place) =>
157            {
158                let successor = target.start_location();
159                if location.strictly_dominates(successor, &self.dominators)
160                    // Don't propagate range information from built-in overflow checks,
161                    // this MIR might be lowered without them!
162                    && !msg.is_optional_overflow_check()
163                {
164                    let val = *expected as u128;
165                    let range = WrappingRange { start: val, end: val };
166                    self.insert_range(place, successor, range);
167                }
168            }
169            TerminatorKind::SwitchInt { discr, targets }
170                if let Some(place) = discr.place()
171                    && self.is_ssa(place)
172                    // Reduce the potential compile-time overhead.
173                    && targets.all_targets().len() < 16 =>
174            {
175                let mut distinct_targets: FxHashMap<BasicBlock, u64> = FxHashMap::default();
176                for (_, target) in targets.iter() {
177                    let targets = distinct_targets.entry(target).or_default();
178                    *targets += 1;
179                }
180                for (val, target) in targets.iter() {
181                    if distinct_targets[&target] != 1 {
182                        // FIXME: For multiple targets, the range can be the union of their values.
183                        continue;
184                    }
185                    let successor = target.start_location();
186                    if self.has_unique_predecessor(successor.block) {
187                        assert_ne!(location.block, successor.block);
188                        let range = WrappingRange { start: val, end: val };
189                        self.insert_range(place, successor, range);
190                    }
191                }
192
193                // FIXME: The range for the otherwise target be extend to more types.
194                // For instance, `val` is within the range [4, 1) at the otherwise target of `matches!(val, 1 | 2 | 3)`.
195                let otherwise = targets.otherwise().start_location();
196                if place.ty(self.local_decls, self.tcx).ty.is_bool()
197                    && let [val] = targets.all_values()
198                    && self.has_unique_predecessor(otherwise.block)
199                {
200                    assert_ne!(location.block, otherwise.block);
201                    let range = if val.get() == 0 {
202                        WrappingRange { start: 1, end: 1 }
203                    } else {
204                        WrappingRange { start: 0, end: 0 }
205                    };
206                    self.insert_range(place, otherwise, range);
207                }
208            }
209            _ => {}
210        }
211    }
212}
rustc_mir_transform/ssa_range_prop.rs

rustc_mir_transform/
ssa_range_prop.rs
