src/librustc_mir/transform/uniform_array_move_out.rs

   1 // This pass converts move out from array by Subslice and
   2 // ConstIndex{.., from_end: true} to ConstIndex move out(s) from begin
   3 // of array. It allows detect error by mir borrowck and elaborate
   4 // drops for array without additional work.
   5 //
   6 // Example:
   7 //
   8 // let a = [ box 1,box 2, box 3];
   9 // if b {
  10 //  let [_a.., _] = a;
  11 // } else {
  12 //  let [.., _b] = a;
  13 // }
  14 //
  15 //  mir statement _10 = move _2[:-1]; replaced by:
  16 //  StorageLive(_12);
  17 //  _12 = move _2[0 of 3];
  18 //  StorageLive(_13);
  19 //  _13 = move _2[1 of 3];
  20 //  _10 = [move _12, move _13]
  21 //  StorageDead(_12);
  22 //  StorageDead(_13);
  23 //
  24 //  and mir statement _11 = move _2[-1 of 1]; replaced by:
  25 //  _11 = move _2[2 of 3];
  26 //
  27 // FIXME: integrate this transformation to the mir build
  28
  29 use rustc::ty;
  30 use rustc::ty::TyCtxt;
  31 use rustc::mir::*;
  32 use rustc::mir::visit::{Visitor, PlaceContext, NonUseContext};
  33 use rustc_data_structures::indexed_vec::{IndexVec};
  34 use crate::transform::{MirPass, MirSource};
  35 use crate::util::patch::MirPatch;
  36
  37 pub struct UniformArrayMoveOut;
  38
  39 impl MirPass for UniformArrayMoveOut {
  40     fn run_pass<'tcx>(&self, tcx: TyCtxt<'tcx>, src: MirSource<'tcx>, body: &mut Body<'tcx>) {
  41         let mut patch = MirPatch::new(body);
  42         let param_env = tcx.param_env(src.def_id());
  43         {
  44             let mut visitor = UniformArrayMoveOutVisitor{body, patch: &mut patch, tcx, param_env};
  45             visitor.visit_body(body);
  46         }
  47         patch.apply(body);
  48     }
  49 }
  50
  51 struct UniformArrayMoveOutVisitor<'a, 'tcx> {
  52     body: &'a Body<'tcx>,
  53     patch: &'a mut MirPatch<'tcx>,
  54     tcx: TyCtxt<'tcx>,
  55     param_env: ty::ParamEnv<'tcx>,
  56 }
  57
  58 impl<'a, 'tcx> Visitor<'tcx> for UniformArrayMoveOutVisitor<'a, 'tcx> {
  59     fn visit_assign(&mut self,
  60                     dst_place: &Place<'tcx>,
  61                     rvalue: &Rvalue<'tcx>,
  62                     location: Location) {
  63         if let Rvalue::Use(Operand::Move(ref src_place)) = rvalue {
  64             if let Some(ref proj) = src_place.projection {
  65                 if let ProjectionElem::ConstantIndex{offset: _,
  66                                                      min_length: _,
  67                                                      from_end: false} = proj.elem {
  68                     // no need to transformation
  69                 } else {
  70                     let place_ty =
  71                         Place::ty_from(&src_place.base, &proj.base, self.body, self.tcx).ty;
  72                     if let ty::Array(item_ty, const_size) = place_ty.sty {
  73                         if let Some(size) = const_size.try_eval_usize(self.tcx, self.param_env) {
  74                             assert!(size <= u32::max_value() as u64,
  75                                     "uniform array move out doesn't supported
  76                                      for array bigger then u32");
  77                             self.uniform(
  78                                 location,
  79                                 dst_place,
  80                                 &src_place.base,
  81                                 proj,
  82                                 item_ty,
  83                                 size as u32,
  84                             );
  85                         }
  86                     }
  87
  88                 }
  89             }
  90         }
  91         self.super_assign(dst_place, rvalue, location)
  92     }
  93 }
  94
  95 impl<'a, 'tcx> UniformArrayMoveOutVisitor<'a, 'tcx> {
  96     fn uniform(&mut self,
  97                location: Location,
  98                dst_place: &Place<'tcx>,
  99                base: &PlaceBase<'tcx>,
 100                proj: &Projection<'tcx>,
 101                item_ty: &'tcx ty::TyS<'tcx>,
 102                size: u32) {
 103         match proj.elem {
 104             // uniforms statements like_10 = move _2[:-1];
 105             ProjectionElem::Subslice{from, to} => {
 106                 self.patch.make_nop(location);
 107                 let temps : Vec<_> = (from..(size-to)).map(|i| {
 108                     let temp = self.patch.new_temp(item_ty, self.body.source_info(location).span);
 109                     self.patch.add_statement(location, StatementKind::StorageLive(temp));
 110                     self.patch.add_assign(location,
 111                                           Place::from(temp),
 112                                           Rvalue::Use(
 113                                               Operand::Move(
 114                                                   Place {
 115                                                       base: base.clone(),
 116                                                       projection: Some(box Projection {
 117                                                           base: proj.base.clone(),
 118                                                           elem: ProjectionElem::ConstantIndex {
 119                                                               offset: i,
 120                                                               min_length: size,
 121                                                               from_end: false,
 122                                                           }
 123                                                       }),
 124                                                   }
 125                                               )
 126                                           )
 127                     );
 128                     temp
 129                 }).collect();
 130                 self.patch.add_assign(
 131                     location,
 132                     dst_place.clone(),
 133                     Rvalue::Aggregate(
 134                         box AggregateKind::Array(item_ty),
 135                         temps.iter().map(
 136                             |x| Operand::Move(Place::from(*x))
 137                         ).collect()
 138                     )
 139                 );
 140                 for temp in temps {
 141                     self.patch.add_statement(location, StatementKind::StorageDead(temp));
 142                 }
 143             }
 144             // uniforms statements like _11 = move _2[-1 of 1];
 145             ProjectionElem::ConstantIndex{offset, min_length: _, from_end: true} => {
 146                 self.patch.make_nop(location);
 147                 self.patch.add_assign(location,
 148                                       dst_place.clone(),
 149                                       Rvalue::Use(
 150                                           Operand::Move(
 151                                               Place {
 152                                                   base: base.clone(),
 153                                                   projection: Some(box Projection {
 154                                                       base: proj.base.clone(),
 155                                                       elem: ProjectionElem::ConstantIndex {
 156                                                           offset: size - offset,
 157                                                           min_length: size,
 158                                                           from_end: false,
 159                                                       },
 160                                                   }),
 161                                               }
 162                                           )
 163                                       )
 164                 );
 165             }
 166             _ => {}
 167         }
 168     }
 169 }
 170
 171 // Restore Subslice move out after analysis
 172 // Example:
 173 //
 174 //  next statements:
 175 //   StorageLive(_12);
 176 //   _12 = move _2[0 of 3];
 177 //   StorageLive(_13);
 178 //   _13 = move _2[1 of 3];
 179 //   _10 = [move _12, move _13]
 180 //   StorageDead(_12);
 181 //   StorageDead(_13);
 182 //
 183 // replaced by _10 = move _2[:-1];
 184
 185 pub struct RestoreSubsliceArrayMoveOut;
 186
 187 impl MirPass for RestoreSubsliceArrayMoveOut {
 188     fn run_pass<'tcx>(&self, tcx: TyCtxt<'tcx>, src: MirSource<'tcx>, body: &mut Body<'tcx>) {
 189         let mut patch = MirPatch::new(body);
 190         let param_env = tcx.param_env(src.def_id());
 191         {
 192             let mut visitor = RestoreDataCollector {
 193                 locals_use: IndexVec::from_elem(LocalUse::new(), &body.local_decls),
 194                 candidates: vec![],
 195             };
 196             visitor.visit_body(body);
 197
 198             for candidate in &visitor.candidates {
 199                 let statement = &body[candidate.block].statements[candidate.statement_index];
 200                 if let StatementKind::Assign(ref dst_place, ref rval) = statement.kind {
 201                     if let Rvalue::Aggregate(box AggregateKind::Array(_), ref items) = **rval {
 202                         let items : Vec<_> = items.iter().map(|item| {
 203                             if let Operand::Move(Place {
 204                                 base: PlaceBase::Local(local),
 205                                 projection: None,
 206                             }) = item {
 207                                 let local_use = &visitor.locals_use[*local];
 208                                 let opt_index_and_place =
 209                                     Self::try_get_item_source(local_use, body);
 210                                 // each local should be used twice:
 211                                 //  in assign and in aggregate statements
 212                                 if local_use.use_count == 2 && opt_index_and_place.is_some() {
 213                                     let (index, src_place) = opt_index_and_place.unwrap();
 214                                     return Some((local_use, index, src_place));
 215                                 }
 216                             }
 217                             None
 218                         }).collect();
 219
 220                         let opt_src_place = items.first().and_then(|x| *x).map(|x| x.2);
 221                         let opt_size = opt_src_place.and_then(|src_place| {
 222                             let src_ty =
 223                                 Place::ty_from(src_place.base, src_place.projection, body, tcx).ty;
 224                             if let ty::Array(_, ref size_o) = src_ty.sty {
 225                                 size_o.try_eval_usize(tcx, param_env)
 226                             } else {
 227                                 None
 228                             }
 229                         });
 230                         Self::check_and_patch(*candidate, &items, opt_size, &mut patch, dst_place);
 231                     }
 232                 }
 233             }
 234         }
 235         patch.apply(body);
 236     }
 237 }
 238
 239 impl RestoreSubsliceArrayMoveOut {
 240     // Checks that source has size, all locals are inited from same source place and
 241     // indices is an integer interval. If all checks pass do the replacent.
 242     // items are Vec<Option<LocalUse, index in source array, source place for init local>>
 243     fn check_and_patch<'tcx>(candidate: Location,
 244                              items: &[Option<(&LocalUse, u32, PlaceRef<'_, 'tcx>)>],
 245                              opt_size: Option<u64>,
 246                              patch: &mut MirPatch<'tcx>,
 247                              dst_place: &Place<'tcx>) {
 248         let opt_src_place = items.first().and_then(|x| *x).map(|x| x.2);
 249
 250         if opt_size.is_some() && items.iter().all(
 251             |l| l.is_some() && l.unwrap().2 == opt_src_place.unwrap()) {
 252             let src_place = opt_src_place.unwrap();
 253
 254             let indices: Vec<_> = items.iter().map(|x| x.unwrap().1).collect();
 255             for i in 1..indices.len() {
 256                 if indices[i - 1] + 1 != indices[i] {
 257                     return;
 258                 }
 259             }
 260
 261             let min = *indices.first().unwrap();
 262             let max = *indices.last().unwrap();
 263
 264             for item in items {
 265                 let locals_use = item.unwrap().0;
 266                 patch.make_nop(locals_use.alive.unwrap());
 267                 patch.make_nop(locals_use.dead.unwrap());
 268                 patch.make_nop(locals_use.first_use.unwrap());
 269             }
 270             patch.make_nop(candidate);
 271             let size = opt_size.unwrap() as u32;
 272             patch.add_assign(candidate,
 273                              dst_place.clone(),
 274                              Rvalue::Use(
 275                                  Operand::Move(
 276                                      Place {
 277                                          base: src_place.base.clone(),
 278                                          projection: Some(box Projection {
 279                                              base: src_place.projection.clone(),
 280                                              elem: ProjectionElem::Subslice{
 281                                                  from: min, to: size - max - 1}})})));
 282         }
 283     }
 284
 285     fn try_get_item_source<'a, 'tcx>(local_use: &LocalUse,
 286                                      body: &'a Body<'tcx>) -> Option<(u32, PlaceRef<'a, 'tcx>)> {
 287         if let Some(location) = local_use.first_use {
 288             let block = &body[location.block];
 289             if block.statements.len() > location.statement_index {
 290                 let statement = &block.statements[location.statement_index];
 291                 if let StatementKind::Assign(
 292                     Place {
 293                         base: PlaceBase::Local(_),
 294                         projection: None,
 295                     },
 296                     box Rvalue::Use(Operand::Move(Place {
 297                         base,
 298                         projection: Some(box Projection {
 299                             base: proj_base,
 300                             elem: ProjectionElem::ConstantIndex {
 301                                 offset, min_length: _, from_end: false
 302                             }
 303                         }),
 304                     }))) = &statement.kind {
 305                     return Some((*offset, PlaceRef {
 306                         base,
 307                         projection: proj_base,
 308                     }))
 309                 }
 310             }
 311         }
 312         None
 313     }
 314 }
 315
 316 #[derive(Copy, Clone, Debug)]
 317 struct LocalUse {
 318     alive: Option<Location>,
 319     dead: Option<Location>,
 320     use_count: u32,
 321     first_use: Option<Location>,
 322 }
 323
 324 impl LocalUse {
 325     pub fn new() -> Self {
 326         LocalUse{alive: None, dead: None, use_count: 0, first_use: None}
 327     }
 328 }
 329
 330 struct RestoreDataCollector {
 331     locals_use: IndexVec<Local, LocalUse>,
 332     candidates: Vec<Location>,
 333 }
 334
 335 impl<'tcx> Visitor<'tcx> for RestoreDataCollector {
 336     fn visit_assign(&mut self,
 337                     place: &Place<'tcx>,
 338                     rvalue: &Rvalue<'tcx>,
 339                     location: Location) {
 340         if let Rvalue::Aggregate(box AggregateKind::Array(_), _) = *rvalue {
 341             self.candidates.push(location);
 342         }
 343         self.super_assign(place, rvalue, location)
 344     }
 345
 346     fn visit_local(&mut self,
 347                    local: &Local,
 348                    context: PlaceContext,
 349                    location: Location) {
 350         let local_use = &mut self.locals_use[*local];
 351         match context {
 352             PlaceContext::NonUse(NonUseContext::StorageLive) => local_use.alive = Some(location),
 353             PlaceContext::NonUse(NonUseContext::StorageDead) => local_use.dead = Some(location),
 354             _ => {
 355                 local_use.use_count += 1;
 356                 if local_use.first_use.is_none() {
 357                     local_use.first_use = Some(location);
 358                 }
 359             }
 360         }
 361     }
 362 }