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_index::bit_set::DenseBitSet;
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, sess: &rustc_session::Session) -> PassPolicy {
29        PassPolicy::optimization(sess.mir_opt_level() > 1)
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        // Clone dominators because we need them while mutating the body.
36        let dominators = body.basic_blocks.dominators().clone();
37        let mut range_set =
38            RangeSet::new(tcx, typing_env, body, &ssa, &body.local_decls, dominators);
39
40        let reverse_postorder = body.basic_blocks.reverse_postorder().to_vec();
41        for bb in reverse_postorder {
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: Dominators<BasicBlock>,
54    /// Known ranges at each locations.
55    ranges: FxHashMap<Place<'tcx>, Vec<(Location, WrappingRange)>>,
56    /// Determines if the basic block has a single unique predecessor.
57    unique_predecessors: DenseBitSet<BasicBlock>,
58}
59
60impl<'tcx, 'body, 'a> RangeSet<'tcx, 'body, 'a> {
61    fn new(
62        tcx: TyCtxt<'tcx>,
63        typing_env: TypingEnv<'tcx>,
64        body: &Body<'tcx>,
65        ssa: &'a SsaLocals,
66        local_decls: &'body LocalDecls<'tcx>,
67        dominators: Dominators<BasicBlock>,
68    ) -> Self {
69        let predecessors = body.basic_blocks.predecessors();
70        let mut unique_predecessors = DenseBitSet::new_empty(body.basic_blocks.len());
71        for bb in body.basic_blocks.indices() {
72            if predecessors[bb].len() == 1 {
73                unique_predecessors.insert(bb);
74            }
75        }
76        RangeSet {
77            tcx,
78            typing_env,
79            ssa,
80            local_decls,
81            dominators,
82            ranges: FxHashMap::default(),
83            unique_predecessors,
84        }
85    }
86
87    /// Create a new known range at the location.
88    fn insert_range(&mut self, place: Place<'tcx>, location: Location, range: WrappingRange) {
89        assert!(self.is_ssa(place));
90        self.ranges.entry(place).or_default().push((location, range));
91    }
92
93    /// Get the known range at the location.
94    fn get_range(&self, place: &Place<'tcx>, location: Location) -> Option<WrappingRange> {
95        let Some(ranges) = self.ranges.get(place) else {
96            return None;
97        };
98        // FIXME: This should use the intersection of all valid ranges.
99        let (_, range) =
100            ranges.iter().find(|(range_loc, _)| range_loc.dominates(location, &self.dominators))?;
101        Some(*range)
102    }
103
104    fn try_as_constant(
105        &mut self,
106        place: Place<'tcx>,
107        location: Location,
108    ) -> Option<ConstOperand<'tcx>> {
109        if let Some(range) = self.get_range(&place, location)
110            && range.start == range.end
111        {
112            let ty = place.ty(self.local_decls, self.tcx).ty;
113            let layout = self.tcx.layout_of(self.typing_env.as_query_input(ty)).ok()?;
114            let value = ConstValue::Scalar(Scalar::from_uint(range.start, layout.size));
115            let const_ = Const::Val(value, ty);
116            return Some(ConstOperand { span: DUMMY_SP, user_ty: None, const_ });
117        }
118        None
119    }
120
121    fn is_ssa(&self, place: Place<'tcx>) -> bool {
122        self.ssa.is_ssa(place.local) && place.is_stable_offset()
123    }
124}
125
126impl<'tcx> MutVisitor<'tcx> for RangeSet<'tcx, '_, '_> {
127    fn tcx(&self) -> TyCtxt<'tcx> {
128        self.tcx
129    }
130
131    fn visit_operand(&mut self, operand: &mut Operand<'tcx>, location: Location) {
132        // Attempts to simplify an operand to a constant value.
133        if let Some(place) = operand.place()
134            && let Some(const_) = self.try_as_constant(place, location)
135        {
136            *operand = Operand::Constant(Box::new(const_));
137        };
138    }
139
140    fn visit_statement(&mut self, statement: &mut Statement<'tcx>, location: Location) {
141        self.super_statement(statement, location);
142        match &statement.kind {
143            StatementKind::Intrinsic(NonDivergingIntrinsic::Assume(operand))
144                if let Some(place) = operand.place()
145                    && self.is_ssa(place) =>
146            {
147                let successor = location.successor_within_block();
148                let range = WrappingRange { start: 1, end: 1 };
149                self.insert_range(place, successor, range);
150            }
151            _ => {}
152        }
153    }
154
155    fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) {
156        self.super_terminator(terminator, location);
157        match &terminator.kind {
158            TerminatorKind::Assert { cond, expected, target, .. }
159                if let Some(place) = cond.place()
160                    && self.is_ssa(place) =>
161            {
162                let successor = Location { block: *target, statement_index: 0 };
163                if location.strictly_dominates(successor, &self.dominators) {
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 = Location { block: target, statement_index: 0 };
186                    if self.unique_predecessors.contains(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 = Location { block: targets.otherwise(), statement_index: 0 };
196                if place.ty(self.local_decls, self.tcx).ty.is_bool()
197                    && let [val] = targets.all_values()
198                    && self.unique_predecessors.contains(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
