src/librustc_mir/transform/generator.rs

   1 //! This is the implementation of the pass which transforms generators into state machines.
   2 //!
   3 //! MIR generation for generators creates a function which has a self argument which
   4 //! passes by value. This argument is effectively a generator type which only contains upvars and
   5 //! is only used for this argument inside the MIR for the generator.
   6 //! It is passed by value to enable upvars to be moved out of it. Drop elaboration runs on that
   7 //! MIR before this pass and creates drop flags for MIR locals.
   8 //! It will also drop the generator argument (which only consists of upvars) if any of the upvars
   9 //! are moved out of. This pass elaborates the drops of upvars / generator argument in the case
  10 //! that none of the upvars were moved out of. This is because we cannot have any drops of this
  11 //! generator in the MIR, since it is used to create the drop glue for the generator. We'd get
  12 //! infinite recursion otherwise.
  13 //!
  14 //! This pass creates the implementation for the Generator::resume function and the drop shim
  15 //! for the generator based on the MIR input. It converts the generator argument from Self to
  16 //! &mut Self adding derefs in the MIR as needed. It computes the final layout of the generator
  17 //! struct which looks like this:
  18 //!     First upvars are stored
  19 //!     It is followed by the generator state field.
  20 //!     Then finally the MIR locals which are live across a suspension point are stored.
  21 //!
  22 //!     struct Generator {
  23 //!         upvars...,
  24 //!         state: u32,
  25 //!         mir_locals...,
  26 //!     }
  27 //!
  28 //! This pass computes the meaning of the state field and the MIR locals which are live
  29 //! across a suspension point. There are however three hardcoded generator states:
  30 //!     0 - Generator have not been resumed yet
  31 //!     1 - Generator has returned / is completed
  32 //!     2 - Generator has been poisoned
  33 //!
  34 //! It also rewrites `return x` and `yield y` as setting a new generator state and returning
  35 //! GeneratorState::Complete(x) and GeneratorState::Yielded(y) respectively.
  36 //! MIR locals which are live across a suspension point are moved to the generator struct
  37 //! with references to them being updated with references to the generator struct.
  38 //!
  39 //! The pass creates two functions which have a switch on the generator state giving
  40 //! the action to take.
  41 //!
  42 //! One of them is the implementation of Generator::resume.
  43 //! For generators with state 0 (unresumed) it starts the execution of the generator.
  44 //! For generators with state 1 (returned) and state 2 (poisoned) it panics.
  45 //! Otherwise it continues the execution from the last suspension point.
  46 //!
  47 //! The other function is the drop glue for the generator.
  48 //! For generators with state 0 (unresumed) it drops the upvars of the generator.
  49 //! For generators with state 1 (returned) and state 2 (poisoned) it does nothing.
  50 //! Otherwise it drops all the values in scope at the last suspension point.
  51
  52 use rustc::hir;
  53 use rustc::hir::def_id::DefId;
  54 use rustc::mir::*;
  55 use rustc::mir::visit::{PlaceContext, Visitor, MutVisitor};
  56 use rustc::ty::{self, TyCtxt, AdtDef, Ty};
  57 use rustc::ty::GeneratorSubsts;
  58 use rustc::ty::layout::VariantIdx;
  59 use rustc::ty::subst::SubstsRef;
  60 use rustc_data_structures::fx::FxHashMap;
  61 use rustc_data_structures::indexed_vec::{Idx, IndexVec};
  62 use rustc_data_structures::bit_set::{BitSet, BitMatrix};
  63 use std::borrow::Cow;
  64 use std::iter;
  65 use std::mem;
  66 use crate::transform::{MirPass, MirSource};
  67 use crate::transform::simplify;
  68 use crate::transform::no_landing_pads::no_landing_pads;
  69 use crate::dataflow::{DataflowResults, DataflowResultsConsumer, FlowAtLocation};
  70 use crate::dataflow::{do_dataflow, DebugFormatted, state_for_location};
  71 use crate::dataflow::{MaybeStorageLive, HaveBeenBorrowedLocals, RequiresStorage};
  72 use crate::util::dump_mir;
  73 use crate::util::liveness;
  74
  75 pub struct StateTransform;
  76
  77 struct RenameLocalVisitor {
  78     from: Local,
  79     to: Local,
  80 }
  81
  82 impl<'tcx> MutVisitor<'tcx> for RenameLocalVisitor {
  83     fn visit_local(&mut self,
  84                    local: &mut Local,
  85                    _: PlaceContext,
  86                    _: Location) {
  87         if *local == self.from {
  88             *local = self.to;
  89         }
  90     }
  91 }
  92
  93 struct DerefArgVisitor;
  94
  95 impl<'tcx> MutVisitor<'tcx> for DerefArgVisitor {
  96     fn visit_local(&mut self,
  97                    local: &mut Local,
  98                    _: PlaceContext,
  99                    _: Location) {
 100         assert_ne!(*local, self_arg());
 101     }
 102
 103     fn visit_place(&mut self,
 104                     place: &mut Place<'tcx>,
 105                     context: PlaceContext,
 106                     location: Location) {
 107         if place.base == PlaceBase::Local(self_arg()) {
 108             replace_base(place, Place {
 109                 base: PlaceBase::Local(self_arg()),
 110                 projection: Some(Box::new(Projection {
 111                     base: None,
 112                     elem: ProjectionElem::Deref,
 113                 })),
 114             });
 115         } else {
 116             self.super_place(place, context, location);
 117         }
 118     }
 119 }
 120
 121 struct PinArgVisitor<'tcx> {
 122     ref_gen_ty: Ty<'tcx>,
 123 }
 124
 125 impl<'tcx> MutVisitor<'tcx> for PinArgVisitor<'tcx> {
 126     fn visit_local(&mut self,
 127                    local: &mut Local,
 128                    _: PlaceContext,
 129                    _: Location) {
 130         assert_ne!(*local, self_arg());
 131     }
 132
 133     fn visit_place(&mut self,
 134                     place: &mut Place<'tcx>,
 135                     context: PlaceContext,
 136                     location: Location) {
 137         if place.base == PlaceBase::Local(self_arg()) {
 138             replace_base(place, Place {
 139                 base: PlaceBase::Local(self_arg()),
 140                 projection: Some(Box::new(Projection {
 141                     base: None,
 142                     elem: ProjectionElem::Field(Field::new(0), self.ref_gen_ty),
 143                 })),
 144             });
 145         } else {
 146             self.super_place(place, context, location);
 147         }
 148     }
 149 }
 150
 151 fn replace_base(place: &mut Place<'tcx>, new_base: Place<'tcx>) {
 152     let mut projection = &mut place.projection;
 153     while let Some(box proj) = projection {
 154         projection = &mut proj.base;
 155     }
 156
 157     place.base = new_base.base;
 158     *projection = new_base.projection;
 159 }
 160
 161 fn self_arg() -> Local {
 162     Local::new(1)
 163 }
 164
 165 /// Generator have not been resumed yet
 166 const UNRESUMED: usize = GeneratorSubsts::UNRESUMED;
 167 /// Generator has returned / is completed
 168 const RETURNED: usize = GeneratorSubsts::RETURNED;
 169 /// Generator has been poisoned
 170 const POISONED: usize = GeneratorSubsts::POISONED;
 171
 172 struct SuspensionPoint {
 173     state: usize,
 174     resume: BasicBlock,
 175     drop: Option<BasicBlock>,
 176     storage_liveness: liveness::LiveVarSet,
 177 }
 178
 179 struct TransformVisitor<'tcx> {
 180     tcx: TyCtxt<'tcx>,
 181     state_adt_ref: &'tcx AdtDef,
 182     state_substs: SubstsRef<'tcx>,
 183
 184     // The type of the discriminant in the generator struct
 185     discr_ty: Ty<'tcx>,
 186
 187     // Mapping from Local to (type of local, generator struct index)
 188     // FIXME(eddyb) This should use `IndexVec<Local, Option<_>>`.
 189     remap: FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
 190
 191     // A map from a suspension point in a block to the locals which have live storage at that point
 192     // FIXME(eddyb) This should use `IndexVec<BasicBlock, Option<_>>`.
 193     storage_liveness: FxHashMap<BasicBlock, liveness::LiveVarSet>,
 194
 195     // A list of suspension points, generated during the transform
 196     suspension_points: Vec<SuspensionPoint>,
 197
 198     // The original RETURN_PLACE local
 199     new_ret_local: Local,
 200 }
 201
 202 impl TransformVisitor<'tcx> {
 203     // Make a GeneratorState rvalue
 204     fn make_state(&self, idx: VariantIdx, val: Operand<'tcx>) -> Rvalue<'tcx> {
 205         let adt = AggregateKind::Adt(self.state_adt_ref, idx, self.state_substs, None, None);
 206         Rvalue::Aggregate(box adt, vec![val])
 207     }
 208
 209     // Create a Place referencing a generator struct field
 210     fn make_field(&self, variant_index: VariantIdx, idx: usize, ty: Ty<'tcx>) -> Place<'tcx> {
 211         let self_place = Place::from(self_arg());
 212         let base = self_place.downcast_unnamed(variant_index);
 213         let field = Projection {
 214             base: base.projection,
 215             elem: ProjectionElem::Field(Field::new(idx), ty),
 216         };
 217         Place {
 218             base: base.base,
 219             projection: Some(Box::new(field)),
 220         }
 221     }
 222
 223     // Create a statement which changes the discriminant
 224     fn set_discr(&self, state_disc: VariantIdx, source_info: SourceInfo) -> Statement<'tcx> {
 225         let self_place = Place::from(self_arg());
 226         Statement {
 227             source_info,
 228             kind: StatementKind::SetDiscriminant { place: self_place, variant_index: state_disc },
 229         }
 230     }
 231
 232     // Create a statement which reads the discriminant into a temporary
 233     fn get_discr(&self, body: &mut Body<'tcx>) -> (Statement<'tcx>, Place<'tcx>) {
 234         let temp_decl = LocalDecl::new_internal(self.tcx.types.isize, body.span);
 235         let local_decls_len = body.local_decls.push(temp_decl);
 236         let temp = Place::from(local_decls_len);
 237
 238         let self_place = Place::from(self_arg());
 239         let assign = Statement {
 240             source_info: source_info(body),
 241             kind: StatementKind::Assign(temp.clone(), box Rvalue::Discriminant(self_place)),
 242         };
 243         (assign, temp)
 244     }
 245 }
 246
 247 impl MutVisitor<'tcx> for TransformVisitor<'tcx> {
 248     fn visit_local(&mut self,
 249                    local: &mut Local,
 250                    _: PlaceContext,
 251                    _: Location) {
 252         assert_eq!(self.remap.get(local), None);
 253     }
 254
 255     fn visit_place(&mut self,
 256                     place: &mut Place<'tcx>,
 257                     context: PlaceContext,
 258                     location: Location) {
 259         if let PlaceBase::Local(l) = place.base {
 260             // Replace an Local in the remap with a generator struct access
 261             if let Some(&(ty, variant_index, idx)) = self.remap.get(&l) {
 262                 replace_base(place, self.make_field(variant_index, idx, ty));
 263             }
 264         } else {
 265             self.super_place(place, context, location);
 266         }
 267     }
 268
 269     fn visit_basic_block_data(&mut self,
 270                               block: BasicBlock,
 271                               data: &mut BasicBlockData<'tcx>) {
 272         // Remove StorageLive and StorageDead statements for remapped locals
 273         data.retain_statements(|s| {
 274             match s.kind {
 275                 StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => {
 276                     !self.remap.contains_key(&l)
 277                 }
 278                 _ => true
 279             }
 280         });
 281
 282         let ret_val = match data.terminator().kind {
 283             TerminatorKind::Return => Some((VariantIdx::new(1),
 284                 None,
 285                 Operand::Move(Place::from(self.new_ret_local)),
 286                 None)),
 287             TerminatorKind::Yield { ref value, resume, drop } => Some((VariantIdx::new(0),
 288                 Some(resume),
 289                 value.clone(),
 290                 drop)),
 291             _ => None
 292         };
 293
 294         if let Some((state_idx, resume, v, drop)) = ret_val {
 295             let source_info = data.terminator().source_info;
 296             // We must assign the value first in case it gets declared dead below
 297             data.statements.push(Statement {
 298                 source_info,
 299                 kind: StatementKind::Assign(Place::RETURN_PLACE,
 300                                             box self.make_state(state_idx, v)),
 301             });
 302             let state = if let Some(resume) = resume { // Yield
 303                 let state = 3 + self.suspension_points.len();
 304
 305                 self.suspension_points.push(SuspensionPoint {
 306                     state,
 307                     resume,
 308                     drop,
 309                     storage_liveness: self.storage_liveness.get(&block).unwrap().clone(),
 310                 });
 311
 312                 VariantIdx::new(state)
 313             } else { // Return
 314                 VariantIdx::new(RETURNED) // state for returned
 315             };
 316             data.statements.push(self.set_discr(state, source_info));
 317             data.terminator.as_mut().unwrap().kind = TerminatorKind::Return;
 318         }
 319
 320         self.super_basic_block_data(block, data);
 321     }
 322 }
 323
 324 fn make_generator_state_argument_indirect<'tcx>(
 325     tcx: TyCtxt<'tcx>,
 326     def_id: DefId,
 327     body: &mut Body<'tcx>,
 328 ) {
 329     let gen_ty = body.local_decls.raw[1].ty;
 330
 331     let region = ty::ReFree(ty::FreeRegion {
 332         scope: def_id,
 333         bound_region: ty::BoundRegion::BrEnv,
 334     });
 335
 336     let region = tcx.mk_region(region);
 337
 338     let ref_gen_ty = tcx.mk_ref(region, ty::TypeAndMut {
 339         ty: gen_ty,
 340         mutbl: hir::MutMutable
 341     });
 342
 343     // Replace the by value generator argument
 344     body.local_decls.raw[1].ty = ref_gen_ty;
 345
 346     // Add a deref to accesses of the generator state
 347     DerefArgVisitor.visit_body(body);
 348 }
 349
 350 fn make_generator_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
 351     let ref_gen_ty = body.local_decls.raw[1].ty;
 352
 353     let pin_did = tcx.lang_items().pin_type().unwrap();
 354     let pin_adt_ref = tcx.adt_def(pin_did);
 355     let substs = tcx.intern_substs(&[ref_gen_ty.into()]);
 356     let pin_ref_gen_ty = tcx.mk_adt(pin_adt_ref, substs);
 357
 358     // Replace the by ref generator argument
 359     body.local_decls.raw[1].ty = pin_ref_gen_ty;
 360
 361     // Add the Pin field access to accesses of the generator state
 362     PinArgVisitor { ref_gen_ty }.visit_body(body);
 363 }
 364
 365 fn replace_result_variable<'tcx>(
 366     ret_ty: Ty<'tcx>,
 367     body: &mut Body<'tcx>,
 368 ) -> Local {
 369     let source_info = source_info(body);
 370     let new_ret = LocalDecl {
 371         mutability: Mutability::Mut,
 372         ty: ret_ty,
 373         user_ty: UserTypeProjections::none(),
 374         name: None,
 375         source_info,
 376         visibility_scope: source_info.scope,
 377         internal: false,
 378         is_block_tail: None,
 379         is_user_variable: None,
 380     };
 381     let new_ret_local = Local::new(body.local_decls.len());
 382     body.local_decls.push(new_ret);
 383     body.local_decls.swap(RETURN_PLACE, new_ret_local);
 384
 385     RenameLocalVisitor {
 386         from: RETURN_PLACE,
 387         to: new_ret_local,
 388     }.visit_body(body);
 389
 390     new_ret_local
 391 }
 392
 393 struct StorageIgnored(liveness::LiveVarSet);
 394
 395 impl<'tcx> Visitor<'tcx> for StorageIgnored {
 396     fn visit_statement(&mut self,
 397                        statement: &Statement<'tcx>,
 398                        _location: Location) {
 399         match statement.kind {
 400             StatementKind::StorageLive(l) |
 401             StatementKind::StorageDead(l) => { self.0.remove(l); }
 402             _ => (),
 403         }
 404     }
 405 }
 406
 407 struct LivenessInfo {
 408     /// Which locals are live across any suspension point.
 409     ///
 410     /// GeneratorSavedLocal is indexed in terms of the elements in this set;
 411     /// i.e. GeneratorSavedLocal::new(1) corresponds to the second local
 412     /// included in this set.
 413     live_locals: liveness::LiveVarSet,
 414
 415     /// The set of saved locals live at each suspension point.
 416     live_locals_at_suspension_points: Vec<BitSet<GeneratorSavedLocal>>,
 417
 418     /// For every saved local, the set of other saved locals that are
 419     /// storage-live at the same time as this local. We cannot overlap locals in
 420     /// the layout which have conflicting storage.
 421     storage_conflicts: BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal>,
 422
 423     /// For every suspending block, the locals which are storage-live across
 424     /// that suspension point.
 425     storage_liveness: FxHashMap<BasicBlock, liveness::LiveVarSet>,
 426 }
 427
 428 fn locals_live_across_suspend_points(
 429     tcx: TyCtxt<'tcx>,
 430     body: &Body<'tcx>,
 431     source: MirSource<'tcx>,
 432     movable: bool,
 433 ) -> LivenessInfo {
 434     let dead_unwinds = BitSet::new_empty(body.basic_blocks().len());
 435     let def_id = source.def_id();
 436
 437     // Calculate when MIR locals have live storage. This gives us an upper bound of their
 438     // lifetimes.
 439     let storage_live_analysis = MaybeStorageLive::new(body);
 440     let storage_live =
 441         do_dataflow(tcx, body, def_id, &[], &dead_unwinds, storage_live_analysis,
 442                     |bd, p| DebugFormatted::new(&bd.body().local_decls[p]));
 443
 444     // Find the MIR locals which do not use StorageLive/StorageDead statements.
 445     // The storage of these locals are always live.
 446     let mut ignored = StorageIgnored(BitSet::new_filled(body.local_decls.len()));
 447     ignored.visit_body(body);
 448
 449     // Calculate the MIR locals which have been previously
 450     // borrowed (even if they are still active).
 451     let borrowed_locals_analysis = HaveBeenBorrowedLocals::new(body);
 452     let borrowed_locals_result =
 453         do_dataflow(tcx, body, def_id, &[], &dead_unwinds, borrowed_locals_analysis,
 454                     |bd, p| DebugFormatted::new(&bd.body().local_decls[p]));
 455
 456     // Calculate the MIR locals that we actually need to keep storage around
 457     // for.
 458     let requires_storage_analysis = RequiresStorage::new(body, &borrowed_locals_result);
 459     let requires_storage =
 460         do_dataflow(tcx, body, def_id, &[], &dead_unwinds, requires_storage_analysis,
 461                     |bd, p| DebugFormatted::new(&bd.body().local_decls[p]));
 462     let requires_storage_analysis = RequiresStorage::new(body, &borrowed_locals_result);
 463
 464     // Calculate the liveness of MIR locals ignoring borrows.
 465     let mut live_locals = liveness::LiveVarSet::new_empty(body.local_decls.len());
 466     let mut liveness = liveness::liveness_of_locals(
 467         body,
 468     );
 469     liveness::dump_mir(
 470         tcx,
 471         "generator_liveness",
 472         source,
 473         body,
 474         &liveness,
 475     );
 476
 477     let mut storage_liveness_map = FxHashMap::default();
 478     let mut live_locals_at_suspension_points = Vec::new();
 479
 480     for (block, data) in body.basic_blocks().iter_enumerated() {
 481         if let TerminatorKind::Yield { .. } = data.terminator().kind {
 482             let loc = Location {
 483                 block: block,
 484                 statement_index: data.statements.len(),
 485             };
 486
 487             if !movable {
 488                 let borrowed_locals = state_for_location(loc,
 489                                                          &borrowed_locals_analysis,
 490                                                          &borrowed_locals_result,
 491                                                          body);
 492                 // The `liveness` variable contains the liveness of MIR locals ignoring borrows.
 493                 // This is correct for movable generators since borrows cannot live across
 494                 // suspension points. However for immovable generators we need to account for
 495                 // borrows, so we conseratively assume that all borrowed locals are live until
 496                 // we find a StorageDead statement referencing the locals.
 497                 // To do this we just union our `liveness` result with `borrowed_locals`, which
 498                 // contains all the locals which has been borrowed before this suspension point.
 499                 // If a borrow is converted to a raw reference, we must also assume that it lives
 500                 // forever. Note that the final liveness is still bounded by the storage liveness
 501                 // of the local, which happens using the `intersect` operation below.
 502                 liveness.outs[block].union(&borrowed_locals);
 503             }
 504
 505             let storage_liveness = state_for_location(loc,
 506                                                       &storage_live_analysis,
 507                                                       &storage_live,
 508                                                       body);
 509
 510             // Store the storage liveness for later use so we can restore the state
 511             // after a suspension point
 512             storage_liveness_map.insert(block, storage_liveness.clone());
 513
 514             let mut storage_required = state_for_location(loc,
 515                                                           &requires_storage_analysis,
 516                                                           &requires_storage,
 517                                                           body);
 518
 519             // Mark locals without storage statements as always requiring storage
 520             storage_required.union(&ignored.0);
 521
 522             // Locals live are live at this point only if they are used across
 523             // suspension points (the `liveness` variable)
 524             // and their storage is required (the `storage_required` variable)
 525             let mut live_locals_here = storage_required;
 526             live_locals_here.intersect(&liveness.outs[block]);
 527
 528             // The generator argument is ignored
 529             live_locals_here.remove(self_arg());
 530
 531             debug!("loc = {:?}, live_locals_here = {:?}", loc, live_locals_here);
 532
 533             // Add the locals live at this suspension point to the set of locals which live across
 534             // any suspension points
 535             live_locals.union(&live_locals_here);
 536
 537             live_locals_at_suspension_points.push(live_locals_here);
 538         }
 539     }
 540     debug!("live_locals = {:?}", live_locals);
 541
 542     // Renumber our liveness_map bitsets to include only the locals we are
 543     // saving.
 544     let live_locals_at_suspension_points = live_locals_at_suspension_points
 545         .iter()
 546         .map(|live_here| renumber_bitset(&live_here, &live_locals))
 547         .collect();
 548
 549     let storage_conflicts = compute_storage_conflicts(
 550         body,
 551         &live_locals,
 552         &ignored,
 553         requires_storage,
 554         requires_storage_analysis);
 555
 556     LivenessInfo {
 557         live_locals,
 558         live_locals_at_suspension_points,
 559         storage_conflicts,
 560         storage_liveness: storage_liveness_map,
 561     }
 562 }
 563
 564 /// Renumbers the items present in `stored_locals` and applies the renumbering
 565 /// to 'input`.
 566 ///
 567 /// For example, if `stored_locals = [1, 3, 5]`, this would be renumbered to
 568 /// `[0, 1, 2]`. Thus, if `input = [3, 5]` we would return `[1, 2]`.
 569 fn renumber_bitset(input: &BitSet<Local>, stored_locals: &liveness::LiveVarSet)
 570 -> BitSet<GeneratorSavedLocal> {
 571     assert!(stored_locals.superset(&input), "{:?} not a superset of {:?}", stored_locals, input);
 572     let mut out = BitSet::new_empty(stored_locals.count());
 573     for (idx, local) in stored_locals.iter().enumerate() {
 574         let saved_local = GeneratorSavedLocal::from(idx);
 575         if input.contains(local) {
 576             out.insert(saved_local);
 577         }
 578     }
 579     debug!("renumber_bitset({:?}, {:?}) => {:?}", input, stored_locals, out);
 580     out
 581 }
 582
 583 /// For every saved local, looks for which locals are StorageLive at the same
 584 /// time. Generates a bitset for every local of all the other locals that may be
 585 /// StorageLive simultaneously with that local. This is used in the layout
 586 /// computation; see `GeneratorLayout` for more.
 587 fn compute_storage_conflicts(
 588     body: &'mir Body<'tcx>,
 589     stored_locals: &liveness::LiveVarSet,
 590     ignored: &StorageIgnored,
 591     requires_storage: DataflowResults<'tcx, RequiresStorage<'mir, 'tcx>>,
 592     _requires_storage_analysis: RequiresStorage<'mir, 'tcx>,
 593 ) -> BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal> {
 594     assert_eq!(body.local_decls.len(), ignored.0.domain_size());
 595     assert_eq!(body.local_decls.len(), stored_locals.domain_size());
 596     debug!("compute_storage_conflicts({:?})", body.span);
 597     debug!("ignored = {:?}", ignored.0);
 598
 599     // Storage ignored locals are not eligible for overlap, since their storage
 600     // is always live.
 601     let mut ineligible_locals = ignored.0.clone();
 602     ineligible_locals.intersect(&stored_locals);
 603
 604     // Compute the storage conflicts for all eligible locals.
 605     let mut visitor = StorageConflictVisitor {
 606         body,
 607         stored_locals: &stored_locals,
 608         local_conflicts: BitMatrix::from_row_n(&ineligible_locals, body.local_decls.len()),
 609     };
 610     let mut state = FlowAtLocation::new(requires_storage);
 611     visitor.analyze_results(&mut state);
 612     let local_conflicts = visitor.local_conflicts;
 613
 614     // Compress the matrix using only stored locals (Local -> GeneratorSavedLocal).
 615     //
 616     // NOTE: Today we store a full conflict bitset for every local. Technically
 617     // this is twice as many bits as we need, since the relation is symmetric.
 618     // However, in practice these bitsets are not usually large. The layout code
 619     // also needs to keep track of how many conflicts each local has, so it's
 620     // simpler to keep it this way for now.
 621     let mut storage_conflicts = BitMatrix::new(stored_locals.count(), stored_locals.count());
 622     for (idx_a, local_a) in stored_locals.iter().enumerate() {
 623         let saved_local_a = GeneratorSavedLocal::new(idx_a);
 624         if ineligible_locals.contains(local_a) {
 625             // Conflicts with everything.
 626             storage_conflicts.insert_all_into_row(saved_local_a);
 627         } else {
 628             // Keep overlap information only for stored locals.
 629             for (idx_b, local_b) in stored_locals.iter().enumerate() {
 630                 let saved_local_b = GeneratorSavedLocal::new(idx_b);
 631                 if local_conflicts.contains(local_a, local_b) {
 632                     storage_conflicts.insert(saved_local_a, saved_local_b);
 633                 }
 634             }
 635         }
 636     }
 637     storage_conflicts
 638 }
 639
 640 struct StorageConflictVisitor<'body, 'tcx, 's> {
 641     body: &'body Body<'tcx>,
 642     stored_locals: &'s liveness::LiveVarSet,
 643     // FIXME(tmandry): Consider using sparse bitsets here once we have good
 644     // benchmarks for generators.
 645     local_conflicts: BitMatrix<Local, Local>,
 646 }
 647
 648 impl<'body, 'tcx, 's> DataflowResultsConsumer<'body, 'tcx>
 649     for StorageConflictVisitor<'body, 'tcx, 's>
 650 {
 651     type FlowState = FlowAtLocation<'tcx, RequiresStorage<'body, 'tcx>>;
 652
 653     fn body(&self) -> &'body Body<'tcx> {
 654         self.body
 655     }
 656
 657     fn visit_block_entry(&mut self,
 658                          block: BasicBlock,
 659                          flow_state: &Self::FlowState) {
 660         // statement_index is only used for logging, so this is fine.
 661         self.apply_state(flow_state, Location { block, statement_index: 0 });
 662     }
 663
 664     fn visit_statement_entry(&mut self,
 665                              loc: Location,
 666                              _stmt: &Statement<'tcx>,
 667                              flow_state: &Self::FlowState) {
 668         self.apply_state(flow_state, loc);
 669     }
 670
 671     fn visit_terminator_entry(&mut self,
 672                               loc: Location,
 673                               _term: &Terminator<'tcx>,
 674                               flow_state: &Self::FlowState) {
 675         self.apply_state(flow_state, loc);
 676     }
 677 }
 678
 679 impl<'body, 'tcx, 's> StorageConflictVisitor<'body, 'tcx, 's> {
 680     fn apply_state(&mut self,
 681                    flow_state: &FlowAtLocation<'tcx, RequiresStorage<'body, 'tcx>>,
 682                    loc: Location) {
 683         // Ignore unreachable blocks.
 684         match self.body.basic_blocks()[loc.block].terminator().kind {
 685             TerminatorKind::Unreachable => return,
 686             _ => (),
 687         };
 688
 689         let mut eligible_storage_live = flow_state.as_dense().clone();
 690         eligible_storage_live.intersect(&self.stored_locals);
 691
 692         for local in eligible_storage_live.iter() {
 693             self.local_conflicts.union_row_with(&eligible_storage_live, local);
 694         }
 695
 696         if eligible_storage_live.count() > 1 {
 697             trace!("at {:?}, eligible_storage_live={:?}", loc, eligible_storage_live);
 698         }
 699     }
 700 }
 701
 702 fn compute_layout<'tcx>(
 703     tcx: TyCtxt<'tcx>,
 704     source: MirSource<'tcx>,
 705     upvars: &Vec<Ty<'tcx>>,
 706     interior: Ty<'tcx>,
 707     movable: bool,
 708     body: &mut Body<'tcx>,
 709 ) -> (
 710     FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
 711     GeneratorLayout<'tcx>,
 712     FxHashMap<BasicBlock, liveness::LiveVarSet>,
 713 ) {
 714     // Use a liveness analysis to compute locals which are live across a suspension point
 715     let LivenessInfo {
 716         live_locals, live_locals_at_suspension_points, storage_conflicts, storage_liveness
 717     } = locals_live_across_suspend_points(tcx, body, source, movable);
 718
 719     // Erase regions from the types passed in from typeck so we can compare them with
 720     // MIR types
 721     let allowed_upvars = tcx.erase_regions(upvars);
 722     let allowed = match interior.sty {
 723         ty::GeneratorWitness(s) => tcx.erase_late_bound_regions(&s),
 724         _ => bug!(),
 725     };
 726
 727     for (local, decl) in body.local_decls.iter_enumerated() {
 728         // Ignore locals which are internal or not live
 729         if !live_locals.contains(local) || decl.internal {
 730             continue;
 731         }
 732
 733         // Sanity check that typeck knows about the type of locals which are
 734         // live across a suspension point
 735         if !allowed.contains(&decl.ty) && !allowed_upvars.contains(&decl.ty) {
 736             span_bug!(body.span,
 737                       "Broken MIR: generator contains type {} in MIR, \
 738                        but typeck only knows about {}",
 739                       decl.ty,
 740                       interior);
 741         }
 742     }
 743
 744     let dummy_local = LocalDecl::new_internal(tcx.mk_unit(), body.span);
 745
 746     // Gather live locals and their indices replacing values in body.local_decls
 747     // with a dummy to avoid changing local indices.
 748     let mut locals = IndexVec::<GeneratorSavedLocal, _>::new();
 749     let mut tys = IndexVec::<GeneratorSavedLocal, _>::new();
 750     let mut decls = IndexVec::<GeneratorSavedLocal, _>::new();
 751     for (idx, local) in live_locals.iter().enumerate() {
 752         let var = mem::replace(&mut body.local_decls[local], dummy_local.clone());
 753         locals.push(local);
 754         tys.push(var.ty);
 755         decls.push(var);
 756         debug!("generator saved local {:?} => {:?}", GeneratorSavedLocal::from(idx), local);
 757     }
 758
 759     // Leave empty variants for the UNRESUMED, RETURNED, and POISONED states.
 760     const RESERVED_VARIANTS: usize = 3;
 761
 762     // Build the generator variant field list.
 763     // Create a map from local indices to generator struct indices.
 764     let mut variant_fields: IndexVec<VariantIdx, IndexVec<Field, GeneratorSavedLocal>> =
 765         iter::repeat(IndexVec::new()).take(RESERVED_VARIANTS).collect();
 766     let mut remap = FxHashMap::default();
 767     for (suspension_point_idx, live_locals) in live_locals_at_suspension_points.iter().enumerate() {
 768         let variant_index = VariantIdx::from(RESERVED_VARIANTS + suspension_point_idx);
 769         let mut fields = IndexVec::new();
 770         for (idx, saved_local) in live_locals.iter().enumerate() {
 771             fields.push(saved_local);
 772             // Note that if a field is included in multiple variants, we will
 773             // just use the first one here. That's fine; fields do not move
 774             // around inside generators, so it doesn't matter which variant
 775             // index we access them by.
 776             remap.entry(locals[saved_local]).or_insert((tys[saved_local], variant_index, idx));
 777         }
 778         variant_fields.push(fields);
 779     }
 780     debug!("generator variant_fields = {:?}", variant_fields);
 781     debug!("generator storage_conflicts = {:#?}", storage_conflicts);
 782
 783     let layout = GeneratorLayout {
 784         field_tys: tys,
 785         variant_fields,
 786         storage_conflicts,
 787         __local_debuginfo_codegen_only_do_not_use: decls,
 788     };
 789
 790     (remap, layout, storage_liveness)
 791 }
 792
 793 fn insert_switch<'tcx>(
 794     body: &mut Body<'tcx>,
 795     cases: Vec<(usize, BasicBlock)>,
 796     transform: &TransformVisitor<'tcx>,
 797     default: TerminatorKind<'tcx>,
 798 ) {
 799     let default_block = insert_term_block(body, default);
 800     let (assign, discr) = transform.get_discr(body);
 801     let switch = TerminatorKind::SwitchInt {
 802         discr: Operand::Move(discr),
 803         switch_ty: transform.discr_ty,
 804         values: Cow::from(cases.iter().map(|&(i, _)| i as u128).collect::<Vec<_>>()),
 805         targets: cases.iter().map(|&(_, d)| d).chain(iter::once(default_block)).collect(),
 806     };
 807
 808     let source_info = source_info(body);
 809     body.basic_blocks_mut().raw.insert(0, BasicBlockData {
 810         statements: vec![assign],
 811         terminator: Some(Terminator {
 812             source_info,
 813             kind: switch,
 814         }),
 815         is_cleanup: false,
 816     });
 817
 818     let blocks = body.basic_blocks_mut().iter_mut();
 819
 820     for target in blocks.flat_map(|b| b.terminator_mut().successors_mut()) {
 821         *target = BasicBlock::new(target.index() + 1);
 822     }
 823 }
 824
 825 fn elaborate_generator_drops<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, body: &mut Body<'tcx>) {
 826     use crate::util::elaborate_drops::{elaborate_drop, Unwind};
 827     use crate::util::patch::MirPatch;
 828     use crate::shim::DropShimElaborator;
 829
 830     // Note that `elaborate_drops` only drops the upvars of a generator, and
 831     // this is ok because `open_drop` can only be reached within that own
 832     // generator's resume function.
 833
 834     let param_env = tcx.param_env(def_id);
 835     let gen = self_arg();
 836
 837     let mut elaborator = DropShimElaborator {
 838         body: body,
 839         patch: MirPatch::new(body),
 840         tcx,
 841         param_env
 842     };
 843
 844     for (block, block_data) in body.basic_blocks().iter_enumerated() {
 845         let (target, unwind, source_info) = match block_data.terminator() {
 846             &Terminator {
 847                 source_info,
 848                 kind: TerminatorKind::Drop {
 849                     location: Place {
 850                         base: PlaceBase::Local(local),
 851                         projection: None,
 852                     },
 853                     target,
 854                     unwind
 855                 }
 856             } if local == gen => (target, unwind, source_info),
 857             _ => continue,
 858         };
 859         let unwind = if block_data.is_cleanup {
 860             Unwind::InCleanup
 861         } else {
 862             Unwind::To(unwind.unwrap_or_else(|| elaborator.patch.resume_block()))
 863         };
 864         elaborate_drop(
 865             &mut elaborator,
 866             source_info,
 867             &Place::from(gen),
 868             (),
 869             target,
 870             unwind,
 871             block,
 872         );
 873     }
 874     elaborator.patch.apply(body);
 875 }
 876
 877 fn create_generator_drop_shim<'tcx>(
 878     tcx: TyCtxt<'tcx>,
 879     transform: &TransformVisitor<'tcx>,
 880     def_id: DefId,
 881     source: MirSource<'tcx>,
 882     gen_ty: Ty<'tcx>,
 883     body: &Body<'tcx>,
 884     drop_clean: BasicBlock,
 885 ) -> Body<'tcx> {
 886     let mut body = body.clone();
 887
 888     let source_info = source_info(&body);
 889
 890     let mut cases = create_cases(&mut body, transform, |point| point.drop);
 891
 892     cases.insert(0, (UNRESUMED, drop_clean));
 893
 894     // The returned state and the poisoned state fall through to the default
 895     // case which is just to return
 896
 897     insert_switch(&mut body, cases, &transform, TerminatorKind::Return);
 898
 899     for block in body.basic_blocks_mut() {
 900         let kind = &mut block.terminator_mut().kind;
 901         if let TerminatorKind::GeneratorDrop = *kind {
 902             *kind = TerminatorKind::Return;
 903         }
 904     }
 905
 906     // Replace the return variable
 907     body.local_decls[RETURN_PLACE] = LocalDecl {
 908         mutability: Mutability::Mut,
 909         ty: tcx.mk_unit(),
 910         user_ty: UserTypeProjections::none(),
 911         name: None,
 912         source_info,
 913         visibility_scope: source_info.scope,
 914         internal: false,
 915         is_block_tail: None,
 916         is_user_variable: None,
 917     };
 918
 919     make_generator_state_argument_indirect(tcx, def_id, &mut body);
 920
 921     // Change the generator argument from &mut to *mut
 922     body.local_decls[self_arg()] = LocalDecl {
 923         mutability: Mutability::Mut,
 924         ty: tcx.mk_ptr(ty::TypeAndMut {
 925             ty: gen_ty,
 926             mutbl: hir::Mutability::MutMutable,
 927         }),
 928         user_ty: UserTypeProjections::none(),
 929         name: None,
 930         source_info,
 931         visibility_scope: source_info.scope,
 932         internal: false,
 933         is_block_tail: None,
 934         is_user_variable: None,
 935     };
 936     if tcx.sess.opts.debugging_opts.mir_emit_retag {
 937         // Alias tracking must know we changed the type
 938         body.basic_blocks_mut()[START_BLOCK].statements.insert(0, Statement {
 939             source_info,
 940             kind: StatementKind::Retag(RetagKind::Raw, Place::from(self_arg())),
 941         })
 942     }
 943
 944     no_landing_pads(tcx, &mut body);
 945
 946     // Make sure we remove dead blocks to remove
 947     // unrelated code from the resume part of the function
 948     simplify::remove_dead_blocks(&mut body);
 949
 950     dump_mir(tcx, None, "generator_drop", &0, source, &mut body, |_, _| Ok(()) );
 951
 952     body
 953 }
 954
 955 fn insert_term_block<'tcx>(body: &mut Body<'tcx>, kind: TerminatorKind<'tcx>) -> BasicBlock {
 956     let term_block = BasicBlock::new(body.basic_blocks().len());
 957     let source_info = source_info(body);
 958     body.basic_blocks_mut().push(BasicBlockData {
 959         statements: Vec::new(),
 960         terminator: Some(Terminator {
 961             source_info,
 962             kind,
 963         }),
 964         is_cleanup: false,
 965     });
 966     term_block
 967 }
 968
 969 fn insert_panic_block<'tcx>(
 970     tcx: TyCtxt<'tcx>,
 971     body: &mut Body<'tcx>,
 972     message: AssertMessage<'tcx>,
 973 ) -> BasicBlock {
 974     let assert_block = BasicBlock::new(body.basic_blocks().len());
 975     let term = TerminatorKind::Assert {
 976         cond: Operand::Constant(box Constant {
 977             span: body.span,
 978             ty: tcx.types.bool,
 979             user_ty: None,
 980             literal: ty::Const::from_bool(tcx, false),
 981         }),
 982         expected: true,
 983         msg: message,
 984         target: assert_block,
 985         cleanup: None,
 986     };
 987
 988     let source_info = source_info(body);
 989     body.basic_blocks_mut().push(BasicBlockData {
 990         statements: Vec::new(),
 991         terminator: Some(Terminator {
 992             source_info,
 993             kind: term,
 994         }),
 995         is_cleanup: false,
 996     });
 997
 998     assert_block
 999 }
1000
1001 fn create_generator_resume_function<'tcx>(
1002     tcx: TyCtxt<'tcx>,
1003     transform: TransformVisitor<'tcx>,
1004     def_id: DefId,
1005     source: MirSource<'tcx>,
1006     body: &mut Body<'tcx>,
1007 ) {
1008     // Poison the generator when it unwinds
1009     for block in body.basic_blocks_mut() {
1010         let source_info = block.terminator().source_info;
1011         if let &TerminatorKind::Resume = &block.terminator().kind {
1012             block.statements.push(
1013                 transform.set_discr(VariantIdx::new(POISONED), source_info));
1014         }
1015     }
1016
1017     let mut cases = create_cases(body, &transform, |point| Some(point.resume));
1018
1019     use rustc::mir::interpret::InterpError::{
1020         GeneratorResumedAfterPanic,
1021         GeneratorResumedAfterReturn,
1022     };
1023
1024     // Jump to the entry point on the unresumed
1025     cases.insert(0, (UNRESUMED, BasicBlock::new(0)));
1026     // Panic when resumed on the returned state
1027     cases.insert(1, (RETURNED, insert_panic_block(tcx, body, GeneratorResumedAfterReturn)));
1028     // Panic when resumed on the poisoned state
1029     cases.insert(2, (POISONED, insert_panic_block(tcx, body, GeneratorResumedAfterPanic)));
1030
1031     insert_switch(body, cases, &transform, TerminatorKind::Unreachable);
1032
1033     make_generator_state_argument_indirect(tcx, def_id, body);
1034     make_generator_state_argument_pinned(tcx, body);
1035
1036     no_landing_pads(tcx, body);
1037
1038     // Make sure we remove dead blocks to remove
1039     // unrelated code from the drop part of the function
1040     simplify::remove_dead_blocks(body);
1041
1042     dump_mir(tcx, None, "generator_resume", &0, source, body, |_, _| Ok(()) );
1043 }
1044
1045 fn source_info(body: &Body<'_>) -> SourceInfo {
1046     SourceInfo {
1047         span: body.span,
1048         scope: OUTERMOST_SOURCE_SCOPE,
1049     }
1050 }
1051
1052 fn insert_clean_drop(body: &mut Body<'_>) -> BasicBlock {
1053     let return_block = insert_term_block(body, TerminatorKind::Return);
1054
1055     // Create a block to destroy an unresumed generators. This can only destroy upvars.
1056     let drop_clean = BasicBlock::new(body.basic_blocks().len());
1057     let term = TerminatorKind::Drop {
1058         location: Place::from(self_arg()),
1059         target: return_block,
1060         unwind: None,
1061     };
1062     let source_info = source_info(body);
1063     body.basic_blocks_mut().push(BasicBlockData {
1064         statements: Vec::new(),
1065         terminator: Some(Terminator {
1066             source_info,
1067             kind: term,
1068         }),
1069         is_cleanup: false,
1070     });
1071
1072     drop_clean
1073 }
1074
1075 fn create_cases<'tcx, F>(
1076     body: &mut Body<'tcx>,
1077     transform: &TransformVisitor<'tcx>,
1078     target: F,
1079 ) -> Vec<(usize, BasicBlock)>
1080 where
1081     F: Fn(&SuspensionPoint) -> Option<BasicBlock>,
1082 {
1083     let source_info = source_info(body);
1084
1085     transform.suspension_points.iter().filter_map(|point| {
1086         // Find the target for this suspension point, if applicable
1087         target(point).map(|target| {
1088             let block = BasicBlock::new(body.basic_blocks().len());
1089             let mut statements = Vec::new();
1090
1091             // Create StorageLive instructions for locals with live storage
1092             for i in 0..(body.local_decls.len()) {
1093                 let l = Local::new(i);
1094                 if point.storage_liveness.contains(l) && !transform.remap.contains_key(&l) {
1095                     statements.push(Statement {
1096                         source_info,
1097                         kind: StatementKind::StorageLive(l),
1098                     });
1099                 }
1100             }
1101
1102             // Then jump to the real target
1103             body.basic_blocks_mut().push(BasicBlockData {
1104                 statements,
1105                 terminator: Some(Terminator {
1106                     source_info,
1107                     kind: TerminatorKind::Goto {
1108                         target,
1109                     },
1110                 }),
1111                 is_cleanup: false,
1112             });
1113
1114             (point.state, block)
1115         })
1116     }).collect()
1117 }
1118
1119 impl MirPass for StateTransform {
1120     fn run_pass<'tcx>(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut Body<'tcx>) {
1121         let yield_ty = if let Some(yield_ty) = body.yield_ty {
1122             yield_ty
1123         } else {
1124             // This only applies to generators
1125             return
1126         };
1127
1128         assert!(body.generator_drop.is_none());
1129
1130         let def_id = source.def_id();
1131
1132         // The first argument is the generator type passed by value
1133         let gen_ty = body.local_decls.raw[1].ty;
1134
1135         // Get the interior types and substs which typeck computed
1136         let (upvars, interior, discr_ty, movable) = match gen_ty.sty {
1137             ty::Generator(_, substs, movability) => {
1138                 (substs.upvar_tys(def_id, tcx).collect(),
1139                  substs.witness(def_id, tcx),
1140                  substs.discr_ty(tcx),
1141                  movability == hir::GeneratorMovability::Movable)
1142             }
1143             _ => bug!(),
1144         };
1145
1146         // Compute GeneratorState<yield_ty, return_ty>
1147         let state_did = tcx.lang_items().gen_state().unwrap();
1148         let state_adt_ref = tcx.adt_def(state_did);
1149         let state_substs = tcx.intern_substs(&[
1150             yield_ty.into(),
1151             body.return_ty().into(),
1152         ]);
1153         let ret_ty = tcx.mk_adt(state_adt_ref, state_substs);
1154
1155         // We rename RETURN_PLACE which has type mir.return_ty to new_ret_local
1156         // RETURN_PLACE then is a fresh unused local with type ret_ty.
1157         let new_ret_local = replace_result_variable(ret_ty, body);
1158
1159         // Extract locals which are live across suspension point into `layout`
1160         // `remap` gives a mapping from local indices onto generator struct indices
1161         // `storage_liveness` tells us which locals have live storage at suspension points
1162         let (remap, layout, storage_liveness) = compute_layout(
1163             tcx,
1164             source,
1165             &upvars,
1166             interior,
1167             movable,
1168             body);
1169
1170         // Run the transformation which converts Places from Local to generator struct
1171         // accesses for locals in `remap`.
1172         // It also rewrites `return x` and `yield y` as writing a new generator state and returning
1173         // GeneratorState::Complete(x) and GeneratorState::Yielded(y) respectively.
1174         let mut transform = TransformVisitor {
1175             tcx,
1176             state_adt_ref,
1177             state_substs,
1178             remap,
1179             storage_liveness,
1180             suspension_points: Vec::new(),
1181             new_ret_local,
1182             discr_ty,
1183         };
1184         transform.visit_body(body);
1185
1186         // Update our MIR struct to reflect the changed we've made
1187         body.yield_ty = None;
1188         body.arg_count = 1;
1189         body.spread_arg = None;
1190         body.generator_layout = Some(layout);
1191
1192         // Insert `drop(generator_struct)` which is used to drop upvars for generators in
1193         // the unresumed state.
1194         // This is expanded to a drop ladder in `elaborate_generator_drops`.
1195         let drop_clean = insert_clean_drop(body);
1196
1197         dump_mir(tcx, None, "generator_pre-elab", &0, source, body, |_, _| Ok(()) );
1198
1199         // Expand `drop(generator_struct)` to a drop ladder which destroys upvars.
1200         // If any upvars are moved out of, drop elaboration will handle upvar destruction.
1201         // However we need to also elaborate the code generated by `insert_clean_drop`.
1202         elaborate_generator_drops(tcx, def_id, body);
1203
1204         dump_mir(tcx, None, "generator_post-transform", &0, source, body, |_, _| Ok(()) );
1205
1206         // Create a copy of our MIR and use it to create the drop shim for the generator
1207         let drop_shim = create_generator_drop_shim(tcx,
1208             &transform,
1209             def_id,
1210             source,
1211             gen_ty,
1212             &body,
1213             drop_clean);
1214
1215         body.generator_drop = Some(box drop_shim);
1216
1217         // Create the Generator::resume function
1218         create_generator_resume_function(tcx, transform, def_id, source, body);
1219     }
1220 }