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