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