1 //! This is the implementation of the pass which transforms generators into state machines.
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.
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.
22 //! struct Generator {
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
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.
39 //! The pass creates two functions which have a switch on the generator state giving
40 //! the action to take.
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.
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.
52 use crate::dataflow::generic::{self as dataflow, Analysis};
53 use crate::dataflow::{MaybeBorrowedLocals, MaybeRequiresStorage, MaybeStorageLive};
54 use crate::transform::no_landing_pads::no_landing_pads;
55 use crate::transform::simplify;
56 use crate::transform::{MirPass, MirSource};
57 use crate::util::dump_mir;
58 use crate::util::liveness;
59 use rustc::mir::visit::{MutVisitor, PlaceContext, Visitor};
61 use rustc::ty::layout::VariantIdx;
62 use rustc::ty::subst::SubstsRef;
63 use rustc::ty::GeneratorSubsts;
64 use rustc::ty::{self, AdtDef, Ty, TyCtxt};
65 use rustc_data_structures::fx::FxHashMap;
67 use rustc_hir::def_id::DefId;
68 use rustc_index::bit_set::{BitMatrix, BitSet};
69 use rustc_index::vec::{Idx, IndexVec};
73 pub struct StateTransform;
75 struct RenameLocalVisitor<'tcx> {
81 impl<'tcx> MutVisitor<'tcx> for RenameLocalVisitor<'tcx> {
82 fn tcx(&self) -> TyCtxt<'tcx> {
86 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
87 if *local == self.from {
92 fn process_projection_elem(&mut self, elem: &PlaceElem<'tcx>) -> Option<PlaceElem<'tcx>> {
94 PlaceElem::Index(local) if *local == self.from => Some(PlaceElem::Index(self.to)),
100 struct DerefArgVisitor<'tcx> {
104 impl<'tcx> MutVisitor<'tcx> for DerefArgVisitor<'tcx> {
105 fn tcx(&self) -> TyCtxt<'tcx> {
109 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
110 assert_ne!(*local, SELF_ARG);
113 fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) {
114 if place.local == SELF_ARG {
119 projection: self.tcx().intern_place_elems(&[ProjectionElem::Deref]),
124 self.visit_place_base(&mut place.local, context, location);
126 for elem in place.projection.iter() {
127 if let PlaceElem::Index(local) = elem {
128 assert_ne!(*local, SELF_ARG);
135 struct PinArgVisitor<'tcx> {
136 ref_gen_ty: Ty<'tcx>,
140 impl<'tcx> MutVisitor<'tcx> for PinArgVisitor<'tcx> {
141 fn tcx(&self) -> TyCtxt<'tcx> {
145 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
146 assert_ne!(*local, SELF_ARG);
149 fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) {
150 if place.local == SELF_ARG {
155 projection: self.tcx().intern_place_elems(&[ProjectionElem::Field(
163 self.visit_place_base(&mut place.local, context, location);
165 for elem in place.projection.iter() {
166 if let PlaceElem::Index(local) = elem {
167 assert_ne!(*local, SELF_ARG);
174 fn replace_base<'tcx>(place: &mut Place<'tcx>, new_base: Place<'tcx>, tcx: TyCtxt<'tcx>) {
175 place.local = new_base.local;
177 let mut new_projection = new_base.projection.to_vec();
178 new_projection.append(&mut place.projection.to_vec());
180 place.projection = tcx.intern_place_elems(&new_projection);
183 const SELF_ARG: Local = Local::from_u32(1);
185 /// Generator has not been resumed yet.
186 const UNRESUMED: usize = GeneratorSubsts::UNRESUMED;
187 /// Generator has returned / is completed.
188 const RETURNED: usize = GeneratorSubsts::RETURNED;
189 /// Generator has panicked and is poisoned.
190 const POISONED: usize = GeneratorSubsts::POISONED;
192 /// A `yield` point in the generator.
193 struct SuspensionPoint<'tcx> {
194 /// State discriminant used when suspending or resuming at this point.
196 /// The block to jump to after resumption.
198 /// Where to move the resume argument after resumption.
199 resume_arg: Place<'tcx>,
200 /// Which block to jump to if the generator is dropped in this state.
201 drop: Option<BasicBlock>,
202 /// Set of locals that have live storage while at this suspension point.
203 storage_liveness: liveness::LiveVarSet,
206 struct TransformVisitor<'tcx> {
208 state_adt_ref: &'tcx AdtDef,
209 state_substs: SubstsRef<'tcx>,
211 // The type of the discriminant in the generator struct
214 // Mapping from Local to (type of local, generator struct index)
215 // FIXME(eddyb) This should use `IndexVec<Local, Option<_>>`.
216 remap: FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
218 // A map from a suspension point in a block to the locals which have live storage at that point
219 // FIXME(eddyb) This should use `IndexVec<BasicBlock, Option<_>>`.
220 storage_liveness: FxHashMap<BasicBlock, liveness::LiveVarSet>,
222 // A list of suspension points, generated during the transform
223 suspension_points: Vec<SuspensionPoint<'tcx>>,
225 // The original RETURN_PLACE local
226 new_ret_local: Local,
229 impl TransformVisitor<'tcx> {
230 // Make a GeneratorState rvalue
231 fn make_state(&self, idx: VariantIdx, val: Operand<'tcx>) -> Rvalue<'tcx> {
232 let adt = AggregateKind::Adt(self.state_adt_ref, idx, self.state_substs, None, None);
233 Rvalue::Aggregate(box adt, vec![val])
236 // Create a Place referencing a generator struct field
237 fn make_field(&self, variant_index: VariantIdx, idx: usize, ty: Ty<'tcx>) -> Place<'tcx> {
238 let self_place = Place::from(SELF_ARG);
239 let base = self.tcx.mk_place_downcast_unnamed(self_place, variant_index);
240 let mut projection = base.projection.to_vec();
241 projection.push(ProjectionElem::Field(Field::new(idx), ty));
243 Place { local: base.local, projection: self.tcx.intern_place_elems(&projection) }
246 // Create a statement which changes the discriminant
247 fn set_discr(&self, state_disc: VariantIdx, source_info: SourceInfo) -> Statement<'tcx> {
248 let self_place = Place::from(SELF_ARG);
251 kind: StatementKind::SetDiscriminant {
252 place: box self_place,
253 variant_index: state_disc,
258 // Create a statement which reads the discriminant into a temporary
259 fn get_discr(&self, body: &mut Body<'tcx>) -> (Statement<'tcx>, Place<'tcx>) {
260 let temp_decl = LocalDecl::new_internal(self.tcx.types.isize, body.span);
261 let local_decls_len = body.local_decls.push(temp_decl);
262 let temp = Place::from(local_decls_len);
264 let self_place = Place::from(SELF_ARG);
265 let assign = Statement {
266 source_info: source_info(body),
267 kind: StatementKind::Assign(box (temp, Rvalue::Discriminant(self_place))),
273 impl MutVisitor<'tcx> for TransformVisitor<'tcx> {
274 fn tcx(&self) -> TyCtxt<'tcx> {
278 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
279 assert_eq!(self.remap.get(local), None);
284 place: &mut Place<'tcx>,
285 _context: PlaceContext,
288 // Replace an Local in the remap with a generator struct access
289 if let Some(&(ty, variant_index, idx)) = self.remap.get(&place.local) {
290 replace_base(place, self.make_field(variant_index, idx, ty), self.tcx);
294 fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) {
295 // Remove StorageLive and StorageDead statements for remapped locals
296 data.retain_statements(|s| match s.kind {
297 StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => {
298 !self.remap.contains_key(&l)
303 let ret_val = match data.terminator().kind {
304 TerminatorKind::Return => Some((
307 Operand::Move(Place::from(self.new_ret_local)),
310 TerminatorKind::Yield { ref value, resume, resume_arg, drop } => {
311 Some((VariantIdx::new(0), Some((resume, resume_arg)), value.clone(), drop))
316 if let Some((state_idx, resume, v, drop)) = ret_val {
317 let source_info = data.terminator().source_info;
318 // We must assign the value first in case it gets declared dead below
319 data.statements.push(Statement {
321 kind: StatementKind::Assign(box (
322 Place::return_place(),
323 self.make_state(state_idx, v),
326 let state = if let Some((resume, resume_arg)) = resume {
328 let state = 3 + self.suspension_points.len();
330 // The resume arg target location might itself be remapped if its base local is
331 // live across a yield.
333 if let Some(&(ty, variant, idx)) = self.remap.get(&resume_arg.local) {
334 self.make_field(variant, idx, ty)
339 self.suspension_points.push(SuspensionPoint {
344 storage_liveness: self.storage_liveness.get(&block).unwrap().clone(),
347 VariantIdx::new(state)
350 VariantIdx::new(RETURNED) // state for returned
352 data.statements.push(self.set_discr(state, source_info));
353 data.terminator_mut().kind = TerminatorKind::Return;
356 self.super_basic_block_data(block, data);
360 fn make_generator_state_argument_indirect<'tcx>(
363 body: &mut BodyAndCache<'tcx>,
365 let gen_ty = body.local_decls.raw[1].ty;
367 let region = ty::ReFree(ty::FreeRegion { scope: def_id, bound_region: ty::BoundRegion::BrEnv });
369 let region = tcx.mk_region(region);
371 let ref_gen_ty = tcx.mk_ref(region, ty::TypeAndMut { ty: gen_ty, mutbl: hir::Mutability::Mut });
373 // Replace the by value generator argument
374 body.local_decls.raw[1].ty = ref_gen_ty;
376 // Add a deref to accesses of the generator state
377 DerefArgVisitor { tcx }.visit_body(body);
380 fn make_generator_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut BodyAndCache<'tcx>) {
381 let ref_gen_ty = body.local_decls.raw[1].ty;
383 let pin_did = tcx.lang_items().pin_type().unwrap();
384 let pin_adt_ref = tcx.adt_def(pin_did);
385 let substs = tcx.intern_substs(&[ref_gen_ty.into()]);
386 let pin_ref_gen_ty = tcx.mk_adt(pin_adt_ref, substs);
388 // Replace the by ref generator argument
389 body.local_decls.raw[1].ty = pin_ref_gen_ty;
391 // Add the Pin field access to accesses of the generator state
392 PinArgVisitor { ref_gen_ty, tcx }.visit_body(body);
395 /// Allocates a new local and replaces all references of `local` with it. Returns the new local.
397 /// `local` will be changed to a new local decl with type `ty`.
399 /// Note that the new local will be uninitialized. It is the caller's responsibility to assign some
400 /// valid value to it before its first use.
401 fn replace_local<'tcx>(
404 body: &mut BodyAndCache<'tcx>,
407 let source_info = source_info(body);
408 let new_decl = LocalDecl {
409 mutability: Mutability::Mut,
411 user_ty: UserTypeProjections::none(),
415 local_info: LocalInfo::Other,
417 let new_local = Local::new(body.local_decls.len());
418 body.local_decls.push(new_decl);
419 body.local_decls.swap(local, new_local);
421 RenameLocalVisitor { from: local, to: new_local, tcx }.visit_body(body);
426 struct StorageIgnored(liveness::LiveVarSet);
428 impl<'tcx> Visitor<'tcx> for StorageIgnored {
429 fn visit_statement(&mut self, statement: &Statement<'tcx>, _location: Location) {
430 match statement.kind {
431 StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => {
439 struct LivenessInfo {
440 /// Which locals are live across any suspension point.
442 /// GeneratorSavedLocal is indexed in terms of the elements in this set;
443 /// i.e. GeneratorSavedLocal::new(1) corresponds to the second local
444 /// included in this set.
445 live_locals: liveness::LiveVarSet,
447 /// The set of saved locals live at each suspension point.
448 live_locals_at_suspension_points: Vec<BitSet<GeneratorSavedLocal>>,
450 /// For every saved local, the set of other saved locals that are
451 /// storage-live at the same time as this local. We cannot overlap locals in
452 /// the layout which have conflicting storage.
453 storage_conflicts: BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal>,
455 /// For every suspending block, the locals which are storage-live across
456 /// that suspension point.
457 storage_liveness: FxHashMap<BasicBlock, liveness::LiveVarSet>,
460 fn locals_live_across_suspend_points(
462 body: ReadOnlyBodyAndCache<'_, 'tcx>,
463 source: MirSource<'tcx>,
466 let def_id = source.def_id();
467 let body_ref: &Body<'_> = &body;
469 // Calculate when MIR locals have live storage. This gives us an upper bound of their
471 let mut storage_live = MaybeStorageLive
472 .into_engine(tcx, body_ref, def_id)
473 .iterate_to_fixpoint()
474 .into_results_cursor(body_ref);
476 // Find the MIR locals which do not use StorageLive/StorageDead statements.
477 // The storage of these locals are always live.
478 let mut ignored = StorageIgnored(BitSet::new_filled(body.local_decls.len()));
479 ignored.visit_body(body);
481 // Calculate the MIR locals which have been previously
482 // borrowed (even if they are still active).
483 let borrowed_locals_results =
484 MaybeBorrowedLocals::all_borrows().into_engine(tcx, body_ref, def_id).iterate_to_fixpoint();
486 let mut borrowed_locals_cursor =
487 dataflow::ResultsCursor::new(body_ref, &borrowed_locals_results);
489 // Calculate the MIR locals that we actually need to keep storage around
491 let requires_storage_results = MaybeRequiresStorage::new(body, &borrowed_locals_results)
492 .into_engine(tcx, body_ref, def_id)
493 .iterate_to_fixpoint();
494 let mut requires_storage_cursor =
495 dataflow::ResultsCursor::new(body_ref, &requires_storage_results);
497 // Calculate the liveness of MIR locals ignoring borrows.
498 let mut live_locals = liveness::LiveVarSet::new_empty(body.local_decls.len());
499 let mut liveness = liveness::liveness_of_locals(body);
500 liveness::dump_mir(tcx, "generator_liveness", source, body_ref, &liveness);
502 let mut storage_liveness_map = FxHashMap::default();
503 let mut live_locals_at_suspension_points = Vec::new();
505 for (block, data) in body.basic_blocks().iter_enumerated() {
506 if let TerminatorKind::Yield { .. } = data.terminator().kind {
507 let loc = Location { block, statement_index: data.statements.len() };
510 // The `liveness` variable contains the liveness of MIR locals ignoring borrows.
511 // This is correct for movable generators since borrows cannot live across
512 // suspension points. However for immovable generators we need to account for
513 // borrows, so we conseratively assume that all borrowed locals are live until
514 // we find a StorageDead statement referencing the locals.
515 // To do this we just union our `liveness` result with `borrowed_locals`, which
516 // contains all the locals which has been borrowed before this suspension point.
517 // If a borrow is converted to a raw reference, we must also assume that it lives
518 // forever. Note that the final liveness is still bounded by the storage liveness
519 // of the local, which happens using the `intersect` operation below.
520 borrowed_locals_cursor.seek_before(loc);
521 liveness.outs[block].union(borrowed_locals_cursor.get());
524 storage_live.seek_before(loc);
525 let storage_liveness = storage_live.get();
527 // Store the storage liveness for later use so we can restore the state
528 // after a suspension point
529 storage_liveness_map.insert(block, storage_liveness.clone());
531 requires_storage_cursor.seek_before(loc);
532 let storage_required = requires_storage_cursor.get().clone();
534 // Locals live are live at this point only if they are used across
535 // suspension points (the `liveness` variable)
536 // and their storage is required (the `storage_required` variable)
537 let mut live_locals_here = storage_required;
538 live_locals_here.intersect(&liveness.outs[block]);
540 // The generator argument is ignored.
541 live_locals_here.remove(SELF_ARG);
543 debug!("loc = {:?}, live_locals_here = {:?}", loc, live_locals_here);
545 // Add the locals live at this suspension point to the set of locals which live across
546 // any suspension points
547 live_locals.union(&live_locals_here);
549 live_locals_at_suspension_points.push(live_locals_here);
552 debug!("live_locals = {:?}", live_locals);
554 // Renumber our liveness_map bitsets to include only the locals we are
556 let live_locals_at_suspension_points = live_locals_at_suspension_points
558 .map(|live_here| renumber_bitset(&live_here, &live_locals))
561 let storage_conflicts =
562 compute_storage_conflicts(body_ref, &live_locals, &ignored, requires_storage_results);
566 live_locals_at_suspension_points,
568 storage_liveness: storage_liveness_map,
572 /// Renumbers the items present in `stored_locals` and applies the renumbering
575 /// For example, if `stored_locals = [1, 3, 5]`, this would be renumbered to
576 /// `[0, 1, 2]`. Thus, if `input = [3, 5]` we would return `[1, 2]`.
578 input: &BitSet<Local>,
579 stored_locals: &liveness::LiveVarSet,
580 ) -> BitSet<GeneratorSavedLocal> {
581 assert!(stored_locals.superset(&input), "{:?} not a superset of {:?}", stored_locals, input);
582 let mut out = BitSet::new_empty(stored_locals.count());
583 for (idx, local) in stored_locals.iter().enumerate() {
584 let saved_local = GeneratorSavedLocal::from(idx);
585 if input.contains(local) {
586 out.insert(saved_local);
589 debug!("renumber_bitset({:?}, {:?}) => {:?}", input, stored_locals, out);
593 /// For every saved local, looks for which locals are StorageLive at the same
594 /// time. Generates a bitset for every local of all the other locals that may be
595 /// StorageLive simultaneously with that local. This is used in the layout
596 /// computation; see `GeneratorLayout` for more.
597 fn compute_storage_conflicts(
598 body: &'mir Body<'tcx>,
599 stored_locals: &liveness::LiveVarSet,
600 ignored: &StorageIgnored,
601 requires_storage: dataflow::Results<'tcx, MaybeRequiresStorage<'mir, 'tcx>>,
602 ) -> BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal> {
603 assert_eq!(body.local_decls.len(), ignored.0.domain_size());
604 assert_eq!(body.local_decls.len(), stored_locals.domain_size());
605 debug!("compute_storage_conflicts({:?})", body.span);
606 debug!("ignored = {:?}", ignored.0);
608 // Storage ignored locals are not eligible for overlap, since their storage
610 let mut ineligible_locals = ignored.0.clone();
611 ineligible_locals.intersect(&stored_locals);
613 // Compute the storage conflicts for all eligible locals.
614 let mut visitor = StorageConflictVisitor {
616 stored_locals: &stored_locals,
617 local_conflicts: BitMatrix::from_row_n(&ineligible_locals, body.local_decls.len()),
620 // Visit only reachable basic blocks. The exact order is not important.
621 let reachable_blocks = traversal::preorder(body).map(|(bb, _)| bb);
622 requires_storage.visit_with(body, reachable_blocks, &mut visitor);
624 let local_conflicts = visitor.local_conflicts;
626 // Compress the matrix using only stored locals (Local -> GeneratorSavedLocal).
628 // NOTE: Today we store a full conflict bitset for every local. Technically
629 // this is twice as many bits as we need, since the relation is symmetric.
630 // However, in practice these bitsets are not usually large. The layout code
631 // also needs to keep track of how many conflicts each local has, so it's
632 // simpler to keep it this way for now.
633 let mut storage_conflicts = BitMatrix::new(stored_locals.count(), stored_locals.count());
634 for (idx_a, local_a) in stored_locals.iter().enumerate() {
635 let saved_local_a = GeneratorSavedLocal::new(idx_a);
636 if ineligible_locals.contains(local_a) {
637 // Conflicts with everything.
638 storage_conflicts.insert_all_into_row(saved_local_a);
640 // Keep overlap information only for stored locals.
641 for (idx_b, local_b) in stored_locals.iter().enumerate() {
642 let saved_local_b = GeneratorSavedLocal::new(idx_b);
643 if local_conflicts.contains(local_a, local_b) {
644 storage_conflicts.insert(saved_local_a, saved_local_b);
652 struct StorageConflictVisitor<'mir, 'tcx, 's> {
653 body: &'mir Body<'tcx>,
654 stored_locals: &'s liveness::LiveVarSet,
655 // FIXME(tmandry): Consider using sparse bitsets here once we have good
656 // benchmarks for generators.
657 local_conflicts: BitMatrix<Local, Local>,
660 impl dataflow::ResultsVisitor<'mir, 'tcx> for StorageConflictVisitor<'mir, 'tcx, '_> {
661 type FlowState = BitSet<Local>;
665 state: &Self::FlowState,
666 _statement: &'mir Statement<'tcx>,
669 self.apply_state(state, loc);
674 state: &Self::FlowState,
675 _terminator: &'mir Terminator<'tcx>,
678 self.apply_state(state, loc);
682 impl<'body, 'tcx, 's> StorageConflictVisitor<'body, 'tcx, 's> {
683 fn apply_state(&mut self, flow_state: &BitSet<Local>, loc: Location) {
684 // Ignore unreachable blocks.
685 match self.body.basic_blocks()[loc.block].terminator().kind {
686 TerminatorKind::Unreachable => return,
690 let mut eligible_storage_live = flow_state.clone();
691 eligible_storage_live.intersect(&self.stored_locals);
693 for local in eligible_storage_live.iter() {
694 self.local_conflicts.union_row_with(&eligible_storage_live, local);
697 if eligible_storage_live.count() > 1 {
698 trace!("at {:?}, eligible_storage_live={:?}", loc, eligible_storage_live);
703 fn compute_layout<'tcx>(
705 source: MirSource<'tcx>,
706 upvars: &Vec<Ty<'tcx>>,
709 body: &mut BodyAndCache<'tcx>,
711 FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
712 GeneratorLayout<'tcx>,
713 FxHashMap<BasicBlock, liveness::LiveVarSet>,
715 // Use a liveness analysis to compute locals which are live across a suspension point
718 live_locals_at_suspension_points,
721 } = locals_live_across_suspend_points(tcx, read_only!(body), source, movable);
723 // Erase regions from the types passed in from typeck so we can compare them with
725 let allowed_upvars = tcx.erase_regions(upvars);
726 let allowed = match interior.kind {
727 ty::GeneratorWitness(s) => tcx.erase_late_bound_regions(&s),
731 for (local, decl) in body.local_decls.iter_enumerated() {
732 // Ignore locals which are internal or not live
733 if !live_locals.contains(local) || decl.internal {
737 // Sanity check that typeck knows about the type of locals which are
738 // live across a suspension point
739 if !allowed.contains(&decl.ty) && !allowed_upvars.contains(&decl.ty) {
742 "Broken MIR: generator contains type {} in MIR, \
743 but typeck only knows about {}",
750 // Gather live local types and their indices.
751 let mut locals = IndexVec::<GeneratorSavedLocal, _>::new();
752 let mut tys = IndexVec::<GeneratorSavedLocal, _>::new();
753 for (idx, local) in live_locals.iter().enumerate() {
755 tys.push(body.local_decls[local].ty);
756 debug!("generator saved local {:?} => {:?}", GeneratorSavedLocal::from(idx), local);
759 // Leave empty variants for the UNRESUMED, RETURNED, and POISONED states.
760 const RESERVED_VARIANTS: usize = 3;
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));
778 variant_fields.push(fields);
780 debug!("generator variant_fields = {:?}", variant_fields);
781 debug!("generator storage_conflicts = {:#?}", storage_conflicts);
783 let layout = GeneratorLayout { field_tys: tys, variant_fields, storage_conflicts };
785 (remap, layout, storage_liveness)
788 /// Replaces the entry point of `body` with a block that switches on the generator discriminant and
789 /// dispatches to blocks according to `cases`.
791 /// After this function, the former entry point of the function will be bb1.
792 fn insert_switch<'tcx>(
793 body: &mut BodyAndCache<'tcx>,
794 cases: Vec<(usize, BasicBlock)>,
795 transform: &TransformVisitor<'tcx>,
796 default: TerminatorKind<'tcx>,
798 let default_block = insert_term_block(body, default);
799 let (assign, discr) = transform.get_discr(body);
800 let switch = TerminatorKind::SwitchInt {
801 discr: Operand::Move(discr),
802 switch_ty: transform.discr_ty,
803 values: Cow::from(cases.iter().map(|&(i, _)| i as u128).collect::<Vec<_>>()),
804 targets: cases.iter().map(|&(_, d)| d).chain(iter::once(default_block)).collect(),
807 let source_info = source_info(body);
808 body.basic_blocks_mut().raw.insert(
811 statements: vec![assign],
812 terminator: Some(Terminator { source_info, kind: switch }),
817 let blocks = body.basic_blocks_mut().iter_mut();
819 for target in blocks.flat_map(|b| b.terminator_mut().successors_mut()) {
820 *target = BasicBlock::new(target.index() + 1);
824 fn elaborate_generator_drops<'tcx>(
827 body: &mut BodyAndCache<'tcx>,
829 use crate::shim::DropShimElaborator;
830 use crate::util::elaborate_drops::{elaborate_drop, Unwind};
831 use crate::util::patch::MirPatch;
833 // Note that `elaborate_drops` only drops the upvars of a generator, and
834 // this is ok because `open_drop` can only be reached within that own
835 // generator's resume function.
837 let param_env = tcx.param_env(def_id);
839 let mut elaborator = DropShimElaborator { body, patch: MirPatch::new(body), tcx, param_env };
841 for (block, block_data) in body.basic_blocks().iter_enumerated() {
842 let (target, unwind, source_info) = match block_data.terminator() {
843 Terminator { source_info, kind: TerminatorKind::Drop { location, target, unwind } } => {
844 if let Some(local) = location.as_local() {
845 if local == SELF_ARG {
846 (target, unwind, source_info)
856 let unwind = if block_data.is_cleanup {
859 Unwind::To(unwind.unwrap_or_else(|| elaborator.patch.resume_block()))
864 &Place::from(SELF_ARG),
871 elaborator.patch.apply(body);
874 fn create_generator_drop_shim<'tcx>(
876 transform: &TransformVisitor<'tcx>,
878 source: MirSource<'tcx>,
880 body: &mut BodyAndCache<'tcx>,
881 drop_clean: BasicBlock,
882 ) -> BodyAndCache<'tcx> {
883 let mut body = body.clone();
884 body.arg_count = 1; // make sure the resume argument is not included here
886 let source_info = source_info(&body);
888 let mut cases = create_cases(&mut body, transform, Operation::Drop);
890 cases.insert(0, (UNRESUMED, drop_clean));
892 // The returned state and the poisoned state fall through to the default
893 // case which is just to return
895 insert_switch(&mut body, cases, &transform, TerminatorKind::Return);
897 for block in body.basic_blocks_mut() {
898 let kind = &mut block.terminator_mut().kind;
899 if let TerminatorKind::GeneratorDrop = *kind {
900 *kind = TerminatorKind::Return;
904 // Replace the return variable
905 body.local_decls[RETURN_PLACE] = LocalDecl {
906 mutability: Mutability::Mut,
908 user_ty: UserTypeProjections::none(),
912 local_info: LocalInfo::Other,
915 make_generator_state_argument_indirect(tcx, def_id, &mut body);
917 // Change the generator argument from &mut to *mut
918 body.local_decls[SELF_ARG] = LocalDecl {
919 mutability: Mutability::Mut,
920 ty: tcx.mk_ptr(ty::TypeAndMut { ty: gen_ty, mutbl: hir::Mutability::Mut }),
921 user_ty: UserTypeProjections::none(),
925 local_info: LocalInfo::Other,
927 if tcx.sess.opts.debugging_opts.mir_emit_retag {
928 // Alias tracking must know we changed the type
929 body.basic_blocks_mut()[START_BLOCK].statements.insert(
933 kind: StatementKind::Retag(RetagKind::Raw, box Place::from(SELF_ARG)),
938 no_landing_pads(tcx, &mut body);
940 // Make sure we remove dead blocks to remove
941 // unrelated code from the resume part of the function
942 simplify::remove_dead_blocks(&mut body);
944 dump_mir(tcx, None, "generator_drop", &0, source, &body, |_, _| Ok(()));
949 fn insert_term_block<'tcx>(
950 body: &mut BodyAndCache<'tcx>,
951 kind: TerminatorKind<'tcx>,
953 let term_block = BasicBlock::new(body.basic_blocks().len());
954 let source_info = source_info(body);
955 body.basic_blocks_mut().push(BasicBlockData {
956 statements: Vec::new(),
957 terminator: Some(Terminator { source_info, kind }),
963 fn insert_panic_block<'tcx>(
965 body: &mut BodyAndCache<'tcx>,
966 message: AssertMessage<'tcx>,
968 let assert_block = BasicBlock::new(body.basic_blocks().len());
969 let term = TerminatorKind::Assert {
970 cond: Operand::Constant(box Constant {
973 literal: ty::Const::from_bool(tcx, false),
977 target: assert_block,
981 let source_info = source_info(body);
982 body.basic_blocks_mut().push(BasicBlockData {
983 statements: Vec::new(),
984 terminator: Some(Terminator { source_info, kind: term }),
991 fn can_return<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) -> bool {
992 // Returning from a function with an uninhabited return type is undefined behavior.
993 if body.return_ty().conservative_is_privately_uninhabited(tcx) {
997 // If there's a return terminator the function may return.
998 for block in body.basic_blocks() {
999 if let TerminatorKind::Return = block.terminator().kind {
1004 // Otherwise the function can't return.
1008 fn can_unwind<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) -> bool {
1009 // Nothing can unwind when landing pads are off.
1010 if tcx.sess.no_landing_pads() {
1014 // Unwinds can only start at certain terminators.
1015 for block in body.basic_blocks() {
1016 match block.terminator().kind {
1017 // These never unwind.
1018 TerminatorKind::Goto { .. }
1019 | TerminatorKind::SwitchInt { .. }
1020 | TerminatorKind::Abort
1021 | TerminatorKind::Return
1022 | TerminatorKind::Unreachable
1023 | TerminatorKind::GeneratorDrop
1024 | TerminatorKind::FalseEdges { .. }
1025 | TerminatorKind::FalseUnwind { .. } => {}
1027 // Resume will *continue* unwinding, but if there's no other unwinding terminator it
1028 // will never be reached.
1029 TerminatorKind::Resume => {}
1031 TerminatorKind::Yield { .. } => {
1032 unreachable!("`can_unwind` called before generator transform")
1035 // These may unwind.
1036 TerminatorKind::Drop { .. }
1037 | TerminatorKind::DropAndReplace { .. }
1038 | TerminatorKind::Call { .. }
1039 | TerminatorKind::Assert { .. } => return true,
1043 // If we didn't find an unwinding terminator, the function cannot unwind.
1047 fn create_generator_resume_function<'tcx>(
1049 transform: TransformVisitor<'tcx>,
1051 source: MirSource<'tcx>,
1052 body: &mut BodyAndCache<'tcx>,
1055 let can_unwind = can_unwind(tcx, body);
1057 // Poison the generator when it unwinds
1059 let poison_block = BasicBlock::new(body.basic_blocks().len());
1060 let source_info = source_info(body);
1061 body.basic_blocks_mut().push(BasicBlockData {
1062 statements: vec![transform.set_discr(VariantIdx::new(POISONED), source_info)],
1063 terminator: Some(Terminator { source_info, kind: TerminatorKind::Resume }),
1067 for (idx, block) in body.basic_blocks_mut().iter_enumerated_mut() {
1068 let source_info = block.terminator().source_info;
1070 if let TerminatorKind::Resume = block.terminator().kind {
1071 // An existing `Resume` terminator is redirected to jump to our dedicated
1072 // "poisoning block" above.
1073 if idx != poison_block {
1074 *block.terminator_mut() = Terminator {
1076 kind: TerminatorKind::Goto { target: poison_block },
1079 } else if !block.is_cleanup {
1080 // Any terminators that *can* unwind but don't have an unwind target set are also
1081 // pointed at our poisoning block (unless they're part of the cleanup path).
1082 if let Some(unwind @ None) = block.terminator_mut().unwind_mut() {
1083 *unwind = Some(poison_block);
1089 let mut cases = create_cases(body, &transform, Operation::Resume);
1091 use rustc::mir::AssertKind::{ResumedAfterPanic, ResumedAfterReturn};
1093 // Jump to the entry point on the unresumed
1094 cases.insert(0, (UNRESUMED, BasicBlock::new(0)));
1096 // Panic when resumed on the returned or poisoned state
1097 let generator_kind = body.generator_kind.unwrap();
1102 (POISONED, insert_panic_block(tcx, body, ResumedAfterPanic(generator_kind))),
1109 (RETURNED, insert_panic_block(tcx, body, ResumedAfterReturn(generator_kind))),
1113 insert_switch(body, cases, &transform, TerminatorKind::Unreachable);
1115 make_generator_state_argument_indirect(tcx, def_id, body);
1116 make_generator_state_argument_pinned(tcx, body);
1118 no_landing_pads(tcx, body);
1120 // Make sure we remove dead blocks to remove
1121 // unrelated code from the drop part of the function
1122 simplify::remove_dead_blocks(body);
1124 dump_mir(tcx, None, "generator_resume", &0, source, body, |_, _| Ok(()));
1127 fn source_info(body: &Body<'_>) -> SourceInfo {
1128 SourceInfo { span: body.span, scope: OUTERMOST_SOURCE_SCOPE }
1131 fn insert_clean_drop(body: &mut BodyAndCache<'_>) -> BasicBlock {
1132 let return_block = insert_term_block(body, TerminatorKind::Return);
1134 // Create a block to destroy an unresumed generators. This can only destroy upvars.
1135 let drop_clean = BasicBlock::new(body.basic_blocks().len());
1136 let term = TerminatorKind::Drop {
1137 location: Place::from(SELF_ARG),
1138 target: return_block,
1141 let source_info = source_info(body);
1142 body.basic_blocks_mut().push(BasicBlockData {
1143 statements: Vec::new(),
1144 terminator: Some(Terminator { source_info, kind: term }),
1151 /// An operation that can be performed on a generator.
1152 #[derive(PartialEq, Copy, Clone)]
1159 fn target_block(self, point: &SuspensionPoint<'_>) -> Option<BasicBlock> {
1161 Operation::Resume => Some(point.resume),
1162 Operation::Drop => point.drop,
1167 fn create_cases<'tcx>(
1168 body: &mut BodyAndCache<'tcx>,
1169 transform: &TransformVisitor<'tcx>,
1170 operation: Operation,
1171 ) -> Vec<(usize, BasicBlock)> {
1172 let source_info = source_info(body);
1177 .filter_map(|point| {
1178 // Find the target for this suspension point, if applicable
1179 operation.target_block(point).map(|target| {
1180 let block = BasicBlock::new(body.basic_blocks().len());
1181 let mut statements = Vec::new();
1183 // Create StorageLive instructions for locals with live storage
1184 for i in 0..(body.local_decls.len()) {
1186 // The resume argument is live on function entry. Don't insert a
1187 // `StorageLive`, or the following `Assign` will read from uninitialized
1192 let l = Local::new(i);
1193 if point.storage_liveness.contains(l) && !transform.remap.contains_key(&l) {
1195 .push(Statement { source_info, kind: StatementKind::StorageLive(l) });
1199 if operation == Operation::Resume {
1200 // Move the resume argument to the destination place of the `Yield` terminator
1201 let resume_arg = Local::new(2); // 0 = return, 1 = self
1202 statements.push(Statement {
1204 kind: StatementKind::Assign(box (
1206 Rvalue::Use(Operand::Move(resume_arg.into())),
1211 // Then jump to the real target
1212 body.basic_blocks_mut().push(BasicBlockData {
1214 terminator: Some(Terminator {
1216 kind: TerminatorKind::Goto { target },
1221 (point.state, block)
1227 impl<'tcx> MirPass<'tcx> for StateTransform {
1228 fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut BodyAndCache<'tcx>) {
1229 let yield_ty = if let Some(yield_ty) = body.yield_ty {
1232 // This only applies to generators
1236 assert!(body.generator_drop.is_none());
1238 let def_id = source.def_id();
1240 // The first argument is the generator type passed by value
1241 let gen_ty = body.local_decls.raw[1].ty;
1243 // Get the interior types and substs which typeck computed
1244 let (upvars, interior, discr_ty, movable) = match gen_ty.kind {
1245 ty::Generator(_, substs, movability) => {
1246 let substs = substs.as_generator();
1248 substs.upvar_tys(def_id, tcx).collect(),
1249 substs.witness(def_id, tcx),
1250 substs.discr_ty(tcx),
1251 movability == hir::Movability::Movable,
1257 // Compute GeneratorState<yield_ty, return_ty>
1258 let state_did = tcx.lang_items().gen_state().unwrap();
1259 let state_adt_ref = tcx.adt_def(state_did);
1260 let state_substs = tcx.intern_substs(&[yield_ty.into(), body.return_ty().into()]);
1261 let ret_ty = tcx.mk_adt(state_adt_ref, state_substs);
1263 // We rename RETURN_PLACE which has type mir.return_ty to new_ret_local
1264 // RETURN_PLACE then is a fresh unused local with type ret_ty.
1265 let new_ret_local = replace_local(RETURN_PLACE, ret_ty, body, tcx);
1267 // We also replace the resume argument and insert an `Assign`.
1268 // This is needed because the resume argument `_2` might be live across a `yield`, in which
1269 // case there is no `Assign` to it that the transform can turn into a store to the generator
1270 // state. After the yield the slot in the generator state would then be uninitialized.
1271 let resume_local = Local::new(2);
1272 let new_resume_local =
1273 replace_local(resume_local, body.local_decls[resume_local].ty, body, tcx);
1275 // When first entering the generator, move the resume argument into its new local.
1276 let source_info = source_info(body);
1277 let stmts = &mut body.basic_blocks_mut()[BasicBlock::new(0)].statements;
1282 kind: StatementKind::Assign(box (
1283 new_resume_local.into(),
1284 Rvalue::Use(Operand::Move(resume_local.into())),
1289 // Extract locals which are live across suspension point into `layout`
1290 // `remap` gives a mapping from local indices onto generator struct indices
1291 // `storage_liveness` tells us which locals have live storage at suspension points
1292 let (remap, layout, storage_liveness) =
1293 compute_layout(tcx, source, &upvars, interior, movable, body);
1295 let can_return = can_return(tcx, body);
1297 // Run the transformation which converts Places from Local to generator struct
1298 // accesses for locals in `remap`.
1299 // It also rewrites `return x` and `yield y` as writing a new generator state and returning
1300 // GeneratorState::Complete(x) and GeneratorState::Yielded(y) respectively.
1301 let mut transform = TransformVisitor {
1307 suspension_points: Vec::new(),
1311 transform.visit_body(body);
1313 // Update our MIR struct to reflect the changes we've made
1314 body.yield_ty = None;
1315 body.arg_count = 2; // self, resume arg
1316 body.spread_arg = None;
1317 body.generator_layout = Some(layout);
1319 // Insert `drop(generator_struct)` which is used to drop upvars for generators in
1320 // the unresumed state.
1321 // This is expanded to a drop ladder in `elaborate_generator_drops`.
1322 let drop_clean = insert_clean_drop(body);
1324 dump_mir(tcx, None, "generator_pre-elab", &0, source, body, |_, _| Ok(()));
1326 // Expand `drop(generator_struct)` to a drop ladder which destroys upvars.
1327 // If any upvars are moved out of, drop elaboration will handle upvar destruction.
1328 // However we need to also elaborate the code generated by `insert_clean_drop`.
1329 elaborate_generator_drops(tcx, def_id, body);
1331 dump_mir(tcx, None, "generator_post-transform", &0, source, body, |_, _| Ok(()));
1333 // Create a copy of our MIR and use it to create the drop shim for the generator
1335 create_generator_drop_shim(tcx, &transform, def_id, source, gen_ty, body, drop_clean);
1337 body.generator_drop = Some(box drop_shim);
1339 // Create the Generator::resume function
1340 create_generator_resume_function(tcx, transform, def_id, source, body, can_return);