|
| 1 | +//! This pass eliminates casting of arrays into slices when their length |
| 2 | +//! is taken using `.len()` method. Handy to preserve information in MIR for const prop |
| 3 | +
|
| 4 | +use crate::MirPass; |
| 5 | +use rustc_data_structures::fx::FxIndexMap; |
| 6 | +use rustc_index::bit_set::BitSet; |
| 7 | +use rustc_index::vec::IndexVec; |
| 8 | +use rustc_middle::mir::*; |
| 9 | +use rustc_middle::ty::{self, TyCtxt}; |
| 10 | + |
| 11 | +const MAX_NUM_BLOCKS: usize = 800; |
| 12 | +const MAX_NUM_LOCALS: usize = 3000; |
| 13 | + |
| 14 | +pub struct NormalizeArrayLen; |
| 15 | + |
| 16 | +impl<'tcx> MirPass<'tcx> for NormalizeArrayLen { |
| 17 | + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
| 18 | + if tcx.sess.mir_opt_level() < 4 { |
| 19 | + return; |
| 20 | + } |
| 21 | + |
| 22 | + // early returns for edge cases of highly unrolled functions |
| 23 | + if body.basic_blocks().len() > MAX_NUM_BLOCKS { |
| 24 | + return; |
| 25 | + } |
| 26 | + if body.local_decls().len() > MAX_NUM_LOCALS { |
| 27 | + return; |
| 28 | + } |
| 29 | + normalize_array_len_calls(tcx, body) |
| 30 | + } |
| 31 | +} |
| 32 | + |
| 33 | +pub fn normalize_array_len_calls<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
| 34 | + let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut(); |
| 35 | + |
| 36 | + // do a preliminary analysis to see if we ever have locals of type `[T;N]` or `&[T;N]` |
| 37 | + let mut interesting_locals = BitSet::new_empty(local_decls.len()); |
| 38 | + for (local, decl) in local_decls.iter_enumerated() { |
| 39 | + match decl.ty.kind() { |
| 40 | + ty::Array(..) => { |
| 41 | + interesting_locals.insert(local); |
| 42 | + } |
| 43 | + ty::Ref(.., ty, Mutability::Not) => match ty.kind() { |
| 44 | + ty::Array(..) => { |
| 45 | + interesting_locals.insert(local); |
| 46 | + } |
| 47 | + _ => {} |
| 48 | + }, |
| 49 | + _ => {} |
| 50 | + } |
| 51 | + } |
| 52 | + if interesting_locals.is_empty() { |
| 53 | + // we have found nothing to analyze |
| 54 | + return; |
| 55 | + } |
| 56 | + let num_intesting_locals = interesting_locals.count(); |
| 57 | + let mut state = FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default()); |
| 58 | + let mut patches_scratchpad = |
| 59 | + FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default()); |
| 60 | + let mut replacements_scratchpad = |
| 61 | + FxIndexMap::with_capacity_and_hasher(num_intesting_locals, Default::default()); |
| 62 | + for block in basic_blocks { |
| 63 | + // make length calls for arrays [T; N] not to decay into length calls for &[T] |
| 64 | + // that forbids constant propagation |
| 65 | + normalize_array_len_call( |
| 66 | + tcx, |
| 67 | + block, |
| 68 | + local_decls, |
| 69 | + &interesting_locals, |
| 70 | + &mut state, |
| 71 | + &mut patches_scratchpad, |
| 72 | + &mut replacements_scratchpad, |
| 73 | + ); |
| 74 | + state.clear(); |
| 75 | + patches_scratchpad.clear(); |
| 76 | + replacements_scratchpad.clear(); |
| 77 | + } |
| 78 | +} |
| 79 | + |
| 80 | +struct Patcher<'a, 'tcx> { |
| 81 | + tcx: TyCtxt<'tcx>, |
| 82 | + patches_scratchpad: &'a FxIndexMap<usize, usize>, |
| 83 | + replacements_scratchpad: &'a mut FxIndexMap<usize, Local>, |
| 84 | + local_decls: &'a mut IndexVec<Local, LocalDecl<'tcx>>, |
| 85 | + statement_idx: usize, |
| 86 | +} |
| 87 | + |
| 88 | +impl<'a, 'tcx> Patcher<'a, 'tcx> { |
| 89 | + fn patch_expand_statement( |
| 90 | + &mut self, |
| 91 | + statement: &mut Statement<'tcx>, |
| 92 | + ) -> Option<std::vec::IntoIter<Statement<'tcx>>> { |
| 93 | + let idx = self.statement_idx; |
| 94 | + if let Some(len_statemnt_idx) = self.patches_scratchpad.get(&idx).copied() { |
| 95 | + let mut statements = Vec::with_capacity(2); |
| 96 | + |
| 97 | + // we are at statement that performs a cast. The only sound way is |
| 98 | + // to create another local that performs a similar copy without a cast and then |
| 99 | + // use this copy in the Len operation |
| 100 | + |
| 101 | + match &statement.kind { |
| 102 | + StatementKind::Assign(box ( |
| 103 | + .., |
| 104 | + Rvalue::Cast( |
| 105 | + CastKind::Pointer(ty::adjustment::PointerCast::Unsize), |
| 106 | + operand, |
| 107 | + _, |
| 108 | + ), |
| 109 | + )) => { |
| 110 | + match operand { |
| 111 | + Operand::Copy(place) | Operand::Move(place) => { |
| 112 | + // create new local |
| 113 | + let ty = operand.ty(self.local_decls, self.tcx); |
| 114 | + let local_decl = |
| 115 | + LocalDecl::with_source_info(ty, statement.source_info.clone()); |
| 116 | + let local = self.local_decls.push(local_decl); |
| 117 | + // make it live |
| 118 | + let mut make_live_statement = statement.clone(); |
| 119 | + make_live_statement.kind = StatementKind::StorageLive(local); |
| 120 | + statements.push(make_live_statement); |
| 121 | + // copy into it |
| 122 | + |
| 123 | + let operand = Operand::Copy(*place); |
| 124 | + let mut make_copy_statement = statement.clone(); |
| 125 | + let assign_to = Place::from(local); |
| 126 | + let rvalue = Rvalue::Use(operand); |
| 127 | + make_copy_statement.kind = |
| 128 | + StatementKind::Assign(box (assign_to, rvalue)); |
| 129 | + statements.push(make_copy_statement); |
| 130 | + |
| 131 | + // to reorder we have to copy and make NOP |
| 132 | + statements.push(statement.clone()); |
| 133 | + statement.make_nop(); |
| 134 | + |
| 135 | + self.replacements_scratchpad.insert(len_statemnt_idx, local); |
| 136 | + } |
| 137 | + _ => { |
| 138 | + unreachable!("it's a bug in the implementation") |
| 139 | + } |
| 140 | + } |
| 141 | + } |
| 142 | + _ => { |
| 143 | + unreachable!("it's a bug in the implementation") |
| 144 | + } |
| 145 | + } |
| 146 | + |
| 147 | + self.statement_idx += 1; |
| 148 | + |
| 149 | + Some(statements.into_iter()) |
| 150 | + } else if let Some(local) = self.replacements_scratchpad.get(&idx).copied() { |
| 151 | + let mut statements = Vec::with_capacity(2); |
| 152 | + |
| 153 | + match &statement.kind { |
| 154 | + StatementKind::Assign(box (into, Rvalue::Len(place))) => { |
| 155 | + let add_deref = if let Some(..) = place.as_local() { |
| 156 | + false |
| 157 | + } else if let Some(..) = place.local_or_deref_local() { |
| 158 | + true |
| 159 | + } else { |
| 160 | + unreachable!("it's a bug in the implementation") |
| 161 | + }; |
| 162 | + // replace len statement |
| 163 | + let mut len_statement = statement.clone(); |
| 164 | + let mut place = Place::from(local); |
| 165 | + if add_deref { |
| 166 | + place = self.tcx.mk_place_deref(place); |
| 167 | + } |
| 168 | + len_statement.kind = StatementKind::Assign(box (*into, Rvalue::Len(place))); |
| 169 | + statements.push(len_statement); |
| 170 | + |
| 171 | + // make temporary dead |
| 172 | + let mut make_dead_statement = statement.clone(); |
| 173 | + make_dead_statement.kind = StatementKind::StorageDead(local); |
| 174 | + statements.push(make_dead_statement); |
| 175 | + |
| 176 | + // make original statement NOP |
| 177 | + statement.make_nop(); |
| 178 | + } |
| 179 | + _ => { |
| 180 | + unreachable!("it's a bug in the implementation") |
| 181 | + } |
| 182 | + } |
| 183 | + |
| 184 | + self.statement_idx += 1; |
| 185 | + |
| 186 | + Some(statements.into_iter()) |
| 187 | + } else { |
| 188 | + self.statement_idx += 1; |
| 189 | + None |
| 190 | + } |
| 191 | + } |
| 192 | +} |
| 193 | + |
| 194 | +fn normalize_array_len_call<'tcx>( |
| 195 | + tcx: TyCtxt<'tcx>, |
| 196 | + block: &mut BasicBlockData<'tcx>, |
| 197 | + local_decls: &mut IndexVec<Local, LocalDecl<'tcx>>, |
| 198 | + interesting_locals: &BitSet<Local>, |
| 199 | + state: &mut FxIndexMap<Local, usize>, |
| 200 | + patches_scratchpad: &mut FxIndexMap<usize, usize>, |
| 201 | + replacements_scratchpad: &mut FxIndexMap<usize, Local>, |
| 202 | +) { |
| 203 | + for (statement_idx, statement) in block.statements.iter_mut().enumerate() { |
| 204 | + match &mut statement.kind { |
| 205 | + StatementKind::Assign(box (place, rvalue)) => { |
| 206 | + match rvalue { |
| 207 | + Rvalue::Cast( |
| 208 | + CastKind::Pointer(ty::adjustment::PointerCast::Unsize), |
| 209 | + operand, |
| 210 | + cast_ty, |
| 211 | + ) => { |
| 212 | + let local = if let Some(local) = place.as_local() { local } else { return }; |
| 213 | + match operand { |
| 214 | + Operand::Copy(place) | Operand::Move(place) => { |
| 215 | + let operand_local = |
| 216 | + if let Some(local) = place.local_or_deref_local() { |
| 217 | + local |
| 218 | + } else { |
| 219 | + return; |
| 220 | + }; |
| 221 | + if !interesting_locals.contains(operand_local) { |
| 222 | + return; |
| 223 | + } |
| 224 | + let operand_ty = local_decls[operand_local].ty; |
| 225 | + match (operand_ty.kind(), cast_ty.kind()) { |
| 226 | + (ty::Array(of_ty_src, ..), ty::Slice(of_ty_dst)) => { |
| 227 | + if of_ty_src == of_ty_dst { |
| 228 | + // this is a cast from [T; N] into [T], so we are good |
| 229 | + state.insert(local, statement_idx); |
| 230 | + } |
| 231 | + } |
| 232 | + // current way of patching doesn't allow to work with `mut` |
| 233 | + ( |
| 234 | + ty::Ref( |
| 235 | + ty::RegionKind::ReErased, |
| 236 | + operand_ty, |
| 237 | + Mutability::Not, |
| 238 | + ), |
| 239 | + ty::Ref(ty::RegionKind::ReErased, cast_ty, Mutability::Not), |
| 240 | + ) => { |
| 241 | + match (operand_ty.kind(), cast_ty.kind()) { |
| 242 | + // current way of patching doesn't allow to work with `mut` |
| 243 | + (ty::Array(of_ty_src, ..), ty::Slice(of_ty_dst)) => { |
| 244 | + if of_ty_src == of_ty_dst { |
| 245 | + // this is a cast from [T; N] into [T], so we are good |
| 246 | + state.insert(local, statement_idx); |
| 247 | + } |
| 248 | + } |
| 249 | + _ => {} |
| 250 | + } |
| 251 | + } |
| 252 | + _ => {} |
| 253 | + } |
| 254 | + } |
| 255 | + _ => {} |
| 256 | + } |
| 257 | + } |
| 258 | + Rvalue::Len(place) => { |
| 259 | + let local = if let Some(local) = place.local_or_deref_local() { |
| 260 | + local |
| 261 | + } else { |
| 262 | + return; |
| 263 | + }; |
| 264 | + if let Some(cast_statement_idx) = state.get(&local).copied() { |
| 265 | + patches_scratchpad.insert(cast_statement_idx, statement_idx); |
| 266 | + } |
| 267 | + } |
| 268 | + _ => { |
| 269 | + // invalidate |
| 270 | + state.remove(&place.local); |
| 271 | + } |
| 272 | + } |
| 273 | + } |
| 274 | + _ => {} |
| 275 | + } |
| 276 | + } |
| 277 | + |
| 278 | + let mut patcher = Patcher { |
| 279 | + tcx, |
| 280 | + patches_scratchpad: &*patches_scratchpad, |
| 281 | + replacements_scratchpad, |
| 282 | + local_decls, |
| 283 | + statement_idx: 0, |
| 284 | + }; |
| 285 | + |
| 286 | + block.expand_statements(|st| patcher.patch_expand_statement(st)); |
| 287 | +} |
0 commit comments