|
| 1 | +//! A pass that propagates the known ranges of SSA locals. |
| 2 | +//! We can know the ranges of SSA locals at some locations for the following code: |
| 3 | +//! ``` |
| 4 | +//! fn foo(a: u32) { |
| 5 | +//! let b = a < 9; |
| 6 | +//! if b { |
| 7 | +//! let c = b; // c is true since b is whitin the range [1, 2) |
| 8 | +//! let d = a < 8; // d is true since b whitin the range [0, 9) |
| 9 | +//! } |
| 10 | +//! } |
| 11 | +//! ``` |
| 12 | +use rustc_abi::WrappingRange; |
| 13 | +use rustc_const_eval::interpret::Scalar; |
| 14 | +use rustc_data_structures::fx::FxHashMap; |
| 15 | +use rustc_data_structures::graph::dominators::Dominators; |
| 16 | +use rustc_index::bit_set::DenseBitSet; |
| 17 | +use rustc_middle::mir::visit::MutVisitor; |
| 18 | +use rustc_middle::mir::{BasicBlock, Body, Location, Operand, Place, TerminatorKind, *}; |
| 19 | +use rustc_middle::ty::{TyCtxt, TypingEnv}; |
| 20 | +use rustc_span::DUMMY_SP; |
| 21 | + |
| 22 | +use crate::ssa::SsaLocals; |
| 23 | + |
| 24 | +pub(super) struct SsaRangePropagation; |
| 25 | + |
| 26 | +impl<'tcx> crate::MirPass<'tcx> for SsaRangePropagation { |
| 27 | + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { |
| 28 | + sess.mir_opt_level() > 1 |
| 29 | + } |
| 30 | + |
| 31 | + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
| 32 | + let typing_env = body.typing_env(tcx); |
| 33 | + let ssa = SsaLocals::new(tcx, body, typing_env); |
| 34 | + // Clone dominators because we need them while mutating the body. |
| 35 | + let dominators = body.basic_blocks.dominators().clone(); |
| 36 | + let mut range_set = |
| 37 | + RangeSet::new(tcx, typing_env, body, &ssa, &body.local_decls, dominators); |
| 38 | + |
| 39 | + let reverse_postorder = body.basic_blocks.reverse_postorder().to_vec(); |
| 40 | + for bb in reverse_postorder { |
| 41 | + let data = &mut body.basic_blocks.as_mut_preserves_cfg()[bb]; |
| 42 | + range_set.visit_basic_block_data(bb, data); |
| 43 | + } |
| 44 | + } |
| 45 | + |
| 46 | + fn is_required(&self) -> bool { |
| 47 | + false |
| 48 | + } |
| 49 | +} |
| 50 | + |
| 51 | +struct RangeSet<'tcx, 'body, 'a> { |
| 52 | + tcx: TyCtxt<'tcx>, |
| 53 | + typing_env: TypingEnv<'tcx>, |
| 54 | + ssa: &'a SsaLocals, |
| 55 | + local_decls: &'body LocalDecls<'tcx>, |
| 56 | + dominators: Dominators<BasicBlock>, |
| 57 | + /// Known ranges at each locations. |
| 58 | + ranges: FxHashMap<Place<'tcx>, Vec<(Location, WrappingRange)>>, |
| 59 | + /// Determines if the basic block has a single unique predecessor. |
| 60 | + unique_predecessors: DenseBitSet<BasicBlock>, |
| 61 | +} |
| 62 | + |
| 63 | +impl<'tcx, 'body, 'a> RangeSet<'tcx, 'body, 'a> { |
| 64 | + fn new( |
| 65 | + tcx: TyCtxt<'tcx>, |
| 66 | + typing_env: TypingEnv<'tcx>, |
| 67 | + body: &Body<'tcx>, |
| 68 | + ssa: &'a SsaLocals, |
| 69 | + local_decls: &'body LocalDecls<'tcx>, |
| 70 | + dominators: Dominators<BasicBlock>, |
| 71 | + ) -> Self { |
| 72 | + let predecessors = body.basic_blocks.predecessors(); |
| 73 | + let mut unique_predecessors = DenseBitSet::new_empty(body.basic_blocks.len()); |
| 74 | + for bb in body.basic_blocks.indices() { |
| 75 | + if predecessors[bb].len() == 1 { |
| 76 | + unique_predecessors.insert(bb); |
| 77 | + } |
| 78 | + } |
| 79 | + RangeSet { |
| 80 | + tcx, |
| 81 | + typing_env, |
| 82 | + ssa, |
| 83 | + local_decls, |
| 84 | + dominators, |
| 85 | + ranges: FxHashMap::default(), |
| 86 | + unique_predecessors, |
| 87 | + } |
| 88 | + } |
| 89 | + |
| 90 | + /// Create a new known range at the location. |
| 91 | + fn insert_range(&mut self, place: Place<'tcx>, location: Location, range: WrappingRange) { |
| 92 | + self.ranges.entry(place).or_default().push((location, range)); |
| 93 | + } |
| 94 | + |
| 95 | + /// Get the known range at the location. |
| 96 | + fn get_range(&self, place: &Place<'tcx>, location: Location) -> Option<WrappingRange> { |
| 97 | + let Some(ranges) = self.ranges.get(place) else { |
| 98 | + return None; |
| 99 | + }; |
| 100 | + // FIXME: This should use the intersection of all valid ranges. |
| 101 | + let (_, range) = |
| 102 | + ranges.iter().find(|(range_loc, _)| range_loc.dominates(location, &self.dominators))?; |
| 103 | + Some(*range) |
| 104 | + } |
| 105 | + |
| 106 | + fn try_as_constant( |
| 107 | + &mut self, |
| 108 | + place: Place<'tcx>, |
| 109 | + location: Location, |
| 110 | + ) -> Option<ConstOperand<'tcx>> { |
| 111 | + if let Some(range) = self.get_range(&place, location) |
| 112 | + && range.start == range.end |
| 113 | + { |
| 114 | + let ty = place.ty(self.local_decls, self.tcx).ty; |
| 115 | + let layout = self.tcx.layout_of(self.typing_env.as_query_input(ty)).ok()?; |
| 116 | + let value = ConstValue::Scalar(Scalar::from_uint(range.start, layout.size)); |
| 117 | + let const_ = Const::Val(value, ty); |
| 118 | + return Some(ConstOperand { span: DUMMY_SP, user_ty: None, const_ }); |
| 119 | + } |
| 120 | + None |
| 121 | + } |
| 122 | + |
| 123 | + /// Attempts to simplify an operand to a constant value. |
| 124 | + /// |
| 125 | + /// Returns |
| 126 | + /// - `Ok(())` if the operand is or can be simplified to a constant. |
| 127 | + /// - `Err(Some(place))` if simplification fails for an SSA local. |
| 128 | + /// - `Err(None)` if simplification fails with no further optimization possible. |
| 129 | + fn simplify_operand( |
| 130 | + &mut self, |
| 131 | + operand: &mut Operand<'tcx>, |
| 132 | + location: Location, |
| 133 | + ) -> Result<(), Option<Place<'tcx>>> { |
| 134 | + let Some(place) = operand.place() else { |
| 135 | + return Ok(()); |
| 136 | + }; |
| 137 | + let Some(const_) = self.try_as_constant(place, location) else { |
| 138 | + if self.is_ssa(place) { |
| 139 | + return Err(Some(place)); |
| 140 | + } else { |
| 141 | + return Err(None); |
| 142 | + } |
| 143 | + }; |
| 144 | + *operand = Operand::Constant(Box::new(const_)); |
| 145 | + Ok(()) |
| 146 | + } |
| 147 | + |
| 148 | + fn is_ssa(&self, place: Place<'tcx>) -> bool { |
| 149 | + self.ssa.is_ssa(place.local) && place.is_stable_offset() |
| 150 | + } |
| 151 | +} |
| 152 | + |
| 153 | +impl<'tcx> MutVisitor<'tcx> for RangeSet<'tcx, '_, '_> { |
| 154 | + fn tcx(&self) -> TyCtxt<'tcx> { |
| 155 | + self.tcx |
| 156 | + } |
| 157 | + |
| 158 | + fn visit_operand(&mut self, operand: &mut Operand<'tcx>, location: Location) { |
| 159 | + let _ = self.simplify_operand(operand, location); |
| 160 | + } |
| 161 | + |
| 162 | + fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) { |
| 163 | + match &mut terminator.kind { |
| 164 | + TerminatorKind::Assert { cond, expected, target, .. } => { |
| 165 | + if let Err(Some(place)) = self.simplify_operand(cond, location) { |
| 166 | + let successor = Location { block: *target, statement_index: 0 }; |
| 167 | + if location.block != successor.block |
| 168 | + && self.unique_predecessors.contains(successor.block) |
| 169 | + { |
| 170 | + let val = *expected as u128; |
| 171 | + let range = WrappingRange { start: val, end: val }; |
| 172 | + self.insert_range(place, successor, range); |
| 173 | + } |
| 174 | + } |
| 175 | + } |
| 176 | + TerminatorKind::SwitchInt { discr, targets } => { |
| 177 | + if let Err(Some(place)) = self.simplify_operand(discr, location) |
| 178 | + && targets.all_targets().len() < 8 |
| 179 | + { |
| 180 | + let mut distinct_targets: FxHashMap<BasicBlock, u8> = FxHashMap::default(); |
| 181 | + for (_, target) in targets.iter() { |
| 182 | + let targets = distinct_targets.entry(target).or_default(); |
| 183 | + if *targets == 0 { |
| 184 | + *targets = 1; |
| 185 | + } else { |
| 186 | + *targets = 2; |
| 187 | + } |
| 188 | + } |
| 189 | + for (val, target) in targets.iter() { |
| 190 | + if distinct_targets[&target] != 1 { |
| 191 | + continue; |
| 192 | + } |
| 193 | + let successor = Location { block: target, statement_index: 0 }; |
| 194 | + if location.block != successor.block |
| 195 | + && self.unique_predecessors.contains(successor.block) |
| 196 | + { |
| 197 | + let range = WrappingRange { start: val, end: val }; |
| 198 | + self.insert_range(place, successor, range); |
| 199 | + } |
| 200 | + } |
| 201 | + |
| 202 | + let otherwise = Location { block: targets.otherwise(), statement_index: 0 }; |
| 203 | + if place.ty(self.local_decls, self.tcx).ty.is_bool() |
| 204 | + && let [val] = targets.all_values() |
| 205 | + && location.block != otherwise.block |
| 206 | + && self.unique_predecessors.contains(otherwise.block) |
| 207 | + { |
| 208 | + let range = if val.get() == 0 { |
| 209 | + WrappingRange { start: 1, end: 1 } |
| 210 | + } else { |
| 211 | + WrappingRange { start: 0, end: 0 } |
| 212 | + }; |
| 213 | + self.insert_range(place, otherwise, range); |
| 214 | + } |
| 215 | + } |
| 216 | + } |
| 217 | + _ => {} |
| 218 | + } |
| 219 | + } |
| 220 | +} |
0 commit comments