compiler/rustc_mir_transform/src/normalize_array_len.rs

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