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::impls::{
53 MaybeBorrowedLocals, MaybeInitializedLocals, MaybeLiveLocals, MaybeStorageLive,
55 use crate::dataflow::{self, Analysis};
56 use crate::transform::no_landing_pads::no_landing_pads;
57 use crate::transform::simplify;
58 use crate::transform::{MirPass, MirSource};
59 use crate::util::dump_mir;
60 use crate::util::storage;
61 use rustc_data_structures::fx::FxHashMap;
63 use rustc_hir::def_id::DefId;
64 use rustc_hir::lang_items::{GeneratorStateLangItem, PinTypeLangItem};
65 use rustc_index::bit_set::{BitMatrix, BitSet};
66 use rustc_index::vec::{Idx, IndexVec};
67 use rustc_middle::mir::visit::{MutVisitor, PlaceContext};
68 use rustc_middle::mir::*;
69 use rustc_middle::ty::subst::SubstsRef;
70 use rustc_middle::ty::GeneratorSubsts;
71 use rustc_middle::ty::{self, AdtDef, Ty, TyCtxt};
72 use rustc_target::abi::VariantIdx;
73 use rustc_target::spec::PanicStrategy;
77 pub struct StateTransform;
79 struct RenameLocalVisitor<'tcx> {
85 impl<'tcx> MutVisitor<'tcx> for RenameLocalVisitor<'tcx> {
86 fn tcx(&self) -> TyCtxt<'tcx> {
90 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
91 if *local == self.from {
96 fn visit_terminator_kind(&mut self, kind: &mut TerminatorKind<'tcx>, location: Location) {
98 TerminatorKind::Return => {
99 // Do not replace the implicit `_0` access here, as that's not possible. The
100 // transform already handles `return` correctly.
102 _ => self.super_terminator_kind(kind, location),
107 struct DerefArgVisitor<'tcx> {
111 impl<'tcx> MutVisitor<'tcx> for DerefArgVisitor<'tcx> {
112 fn tcx(&self) -> TyCtxt<'tcx> {
116 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
117 assert_ne!(*local, SELF_ARG);
120 fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) {
121 if place.local == SELF_ARG {
126 projection: self.tcx().intern_place_elems(&[ProjectionElem::Deref]),
131 self.visit_local(&mut place.local, context, location);
133 for elem in place.projection.iter() {
134 if let PlaceElem::Index(local) = elem {
135 assert_ne!(local, SELF_ARG);
142 struct PinArgVisitor<'tcx> {
143 ref_gen_ty: Ty<'tcx>,
147 impl<'tcx> MutVisitor<'tcx> for PinArgVisitor<'tcx> {
148 fn tcx(&self) -> TyCtxt<'tcx> {
152 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
153 assert_ne!(*local, SELF_ARG);
156 fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) {
157 if place.local == SELF_ARG {
162 projection: self.tcx().intern_place_elems(&[ProjectionElem::Field(
170 self.visit_local(&mut place.local, context, location);
172 for elem in place.projection.iter() {
173 if let PlaceElem::Index(local) = elem {
174 assert_ne!(local, SELF_ARG);
181 fn replace_base<'tcx>(place: &mut Place<'tcx>, new_base: Place<'tcx>, tcx: TyCtxt<'tcx>) {
182 place.local = new_base.local;
184 let mut new_projection = new_base.projection.to_vec();
185 new_projection.append(&mut place.projection.to_vec());
187 place.projection = tcx.intern_place_elems(&new_projection);
190 const SELF_ARG: Local = Local::from_u32(1);
192 /// Generator has not been resumed yet.
193 const UNRESUMED: usize = GeneratorSubsts::UNRESUMED;
194 /// Generator has returned / is completed.
195 const RETURNED: usize = GeneratorSubsts::RETURNED;
196 /// Generator has panicked and is poisoned.
197 const POISONED: usize = GeneratorSubsts::POISONED;
199 /// A `yield` point in the generator.
200 struct SuspensionPoint<'tcx> {
201 /// State discriminant used when suspending or resuming at this point.
203 /// The block to jump to after resumption.
205 /// Where to move the resume argument after resumption.
206 resume_arg: Place<'tcx>,
207 /// Which block to jump to if the generator is dropped in this state.
208 drop: Option<BasicBlock>,
209 /// Set of locals that have live storage while at this suspension point.
210 storage_liveness: BitSet<Local>,
213 struct TransformVisitor<'tcx> {
215 state_adt_ref: &'tcx AdtDef,
216 state_substs: SubstsRef<'tcx>,
218 // The type of the discriminant in the generator struct
221 // Mapping from Local to (type of local, generator struct index)
222 // FIXME(eddyb) This should use `IndexVec<Local, Option<_>>`.
223 remap: FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
225 // A map from a suspension point in a block to the locals which have live storage at that point
226 storage_liveness: IndexVec<BasicBlock, Option<BitSet<Local>>>,
228 // A list of suspension points, generated during the transform
229 suspension_points: Vec<SuspensionPoint<'tcx>>,
231 // The set of locals that have no `StorageLive`/`StorageDead` annotations.
232 always_live_locals: storage::AlwaysLiveLocals,
234 // The original RETURN_PLACE local
235 new_ret_local: Local,
238 impl TransformVisitor<'tcx> {
239 // Make a GeneratorState rvalue
240 fn make_state(&self, idx: VariantIdx, val: Operand<'tcx>) -> Rvalue<'tcx> {
241 let adt = AggregateKind::Adt(self.state_adt_ref, idx, self.state_substs, None, None);
242 Rvalue::Aggregate(box adt, vec![val])
245 // Create a Place referencing a generator struct field
246 fn make_field(&self, variant_index: VariantIdx, idx: usize, ty: Ty<'tcx>) -> Place<'tcx> {
247 let self_place = Place::from(SELF_ARG);
248 let base = self.tcx.mk_place_downcast_unnamed(self_place, variant_index);
249 let mut projection = base.projection.to_vec();
250 projection.push(ProjectionElem::Field(Field::new(idx), ty));
252 Place { local: base.local, projection: self.tcx.intern_place_elems(&projection) }
255 // Create a statement which changes the discriminant
256 fn set_discr(&self, state_disc: VariantIdx, source_info: SourceInfo) -> Statement<'tcx> {
257 let self_place = Place::from(SELF_ARG);
260 kind: StatementKind::SetDiscriminant {
261 place: box self_place,
262 variant_index: state_disc,
267 // Create a statement which reads the discriminant into a temporary
268 fn get_discr(&self, body: &mut Body<'tcx>) -> (Statement<'tcx>, Place<'tcx>) {
269 let temp_decl = LocalDecl::new(self.discr_ty, body.span).internal();
270 let local_decls_len = body.local_decls.push(temp_decl);
271 let temp = Place::from(local_decls_len);
273 let self_place = Place::from(SELF_ARG);
274 let assign = Statement {
275 source_info: SourceInfo::outermost(body.span),
276 kind: StatementKind::Assign(box (temp, Rvalue::Discriminant(self_place))),
282 impl MutVisitor<'tcx> for TransformVisitor<'tcx> {
283 fn tcx(&self) -> TyCtxt<'tcx> {
287 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
288 assert_eq!(self.remap.get(local), None);
293 place: &mut Place<'tcx>,
294 _context: PlaceContext,
297 // Replace an Local in the remap with a generator struct access
298 if let Some(&(ty, variant_index, idx)) = self.remap.get(&place.local) {
299 replace_base(place, self.make_field(variant_index, idx, ty), self.tcx);
303 fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) {
304 // Remove StorageLive and StorageDead statements for remapped locals
305 data.retain_statements(|s| match s.kind {
306 StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => {
307 !self.remap.contains_key(&l)
312 let ret_val = match data.terminator().kind {
313 TerminatorKind::Return => Some((
316 Operand::Move(Place::from(self.new_ret_local)),
319 TerminatorKind::Yield { ref value, resume, resume_arg, drop } => {
320 Some((VariantIdx::new(0), Some((resume, resume_arg)), value.clone(), drop))
325 if let Some((state_idx, resume, v, drop)) = ret_val {
326 let source_info = data.terminator().source_info;
327 // We must assign the value first in case it gets declared dead below
328 data.statements.push(Statement {
330 kind: StatementKind::Assign(box (
331 Place::return_place(),
332 self.make_state(state_idx, v),
335 let state = if let Some((resume, resume_arg)) = resume {
337 let state = 3 + self.suspension_points.len();
339 // The resume arg target location might itself be remapped if its base local is
340 // live across a yield.
342 if let Some(&(ty, variant, idx)) = self.remap.get(&resume_arg.local) {
343 self.make_field(variant, idx, ty)
348 self.suspension_points.push(SuspensionPoint {
353 storage_liveness: self.storage_liveness[block].clone().unwrap(),
356 VariantIdx::new(state)
359 VariantIdx::new(RETURNED) // state for returned
361 data.statements.push(self.set_discr(state, source_info));
362 data.terminator_mut().kind = TerminatorKind::Return;
365 self.super_basic_block_data(block, data);
369 fn make_generator_state_argument_indirect<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
370 let gen_ty = body.local_decls.raw[1].ty;
373 tcx.mk_ref(tcx.lifetimes.re_erased, ty::TypeAndMut { ty: gen_ty, mutbl: Mutability::Mut });
375 // Replace the by value generator argument
376 body.local_decls.raw[1].ty = ref_gen_ty;
378 // Add a deref to accesses of the generator state
379 DerefArgVisitor { tcx }.visit_body(body);
382 fn make_generator_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
383 let ref_gen_ty = body.local_decls.raw[1].ty;
385 let pin_did = tcx.require_lang_item(PinTypeLangItem, Some(body.span));
386 let pin_adt_ref = tcx.adt_def(pin_did);
387 let substs = tcx.intern_substs(&[ref_gen_ty.into()]);
388 let pin_ref_gen_ty = tcx.mk_adt(pin_adt_ref, substs);
390 // Replace the by ref generator argument
391 body.local_decls.raw[1].ty = pin_ref_gen_ty;
393 // Add the Pin field access to accesses of the generator state
394 PinArgVisitor { ref_gen_ty, tcx }.visit_body(body);
397 /// Allocates a new local and replaces all references of `local` with it. Returns the new local.
399 /// `local` will be changed to a new local decl with type `ty`.
401 /// Note that the new local will be uninitialized. It is the caller's responsibility to assign some
402 /// valid value to it before its first use.
403 fn replace_local<'tcx>(
406 body: &mut Body<'tcx>,
409 let new_decl = LocalDecl::new(ty, body.span);
410 let new_local = body.local_decls.push(new_decl);
411 body.local_decls.swap(local, new_local);
413 RenameLocalVisitor { from: local, to: new_local, tcx }.visit_body(body);
418 struct LivenessInfo {
419 /// Which locals are live across any suspension point.
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: BitSet<Local>,
426 /// The set of saved locals live at each suspension point.
427 live_locals_at_suspension_points: Vec<BitSet<GeneratorSavedLocal>>,
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>,
434 /// For every suspending block, the locals which are storage-live across
435 /// that suspension point.
436 storage_liveness: IndexVec<BasicBlock, Option<BitSet<Local>>>,
439 fn locals_live_across_suspend_points(
442 source: MirSource<'tcx>,
443 always_live_locals: &storage::AlwaysLiveLocals,
446 let def_id = source.def_id();
448 // Calculate when MIR locals have live storage. This gives us an upper bound of their
450 let mut storage_live = MaybeStorageLive::new(always_live_locals.clone())
451 .into_engine(tcx, body, def_id)
452 .iterate_to_fixpoint()
453 .into_results_cursor(body);
455 let mut init = MaybeInitializedLocals
456 .into_engine(tcx, body, def_id)
457 .iterate_to_fixpoint()
458 .into_results_cursor(body);
460 let mut live = MaybeLiveLocals
461 .into_engine(tcx, body, def_id)
462 .iterate_to_fixpoint()
463 .into_results_cursor(body);
465 let mut borrowed = MaybeBorrowedLocals::all_borrows()
466 .into_engine(tcx, body, def_id)
467 .iterate_to_fixpoint()
468 .into_results_cursor(body);
470 // Liveness across yield points is determined by the following boolean equation, where `live`,
471 // `init` and `borrowed` come from dataflow and `movable` is a property of the generator.
472 // Movable generators do not allow borrows to live across yield points, so they don't need to
473 // store a local simply because it is borrowed.
475 // live_across_yield := (live & init) | (!movable & borrowed)
477 let mut locals_live_across_yield_point = |block| {
478 live.seek_to_block_end(block);
479 let mut live_locals = live.get().clone();
481 init.seek_to_block_end(block);
482 live_locals.intersect(init.get());
485 borrowed.seek_to_block_end(block);
486 live_locals.union(borrowed.get());
492 let mut storage_liveness_map = IndexVec::from_elem(None, body.basic_blocks());
493 let mut live_locals_at_suspension_points = Vec::new();
494 let mut live_locals_at_any_suspension_point = BitSet::new_empty(body.local_decls.len());
496 for (block, data) in body.basic_blocks().iter_enumerated() {
497 if !matches!(data.terminator().kind, TerminatorKind::Yield { .. }) {
501 // Store the storage liveness for later use so we can restore the state
502 // after a suspension point
503 storage_live.seek_to_block_end(block);
504 storage_liveness_map[block] = Some(storage_live.get().clone());
506 let mut live_locals = locals_live_across_yield_point(block);
508 // The combination of `MaybeInitializedLocals` and `MaybeBorrowedLocals` should be strictly
509 // more precise than `MaybeStorageLive` because they handle `StorageDead` themselves. This
510 // assumes that the MIR forbids locals from being initialized/borrowed before reaching
512 debug_assert!(storage_live.get().superset(&live_locals));
514 // Ignore the generator's `self` argument since it is handled seperately.
515 live_locals.remove(SELF_ARG);
516 debug!("block = {:?}, live_locals = {:?}", block, live_locals);
517 live_locals_at_any_suspension_point.union(&live_locals);
518 live_locals_at_suspension_points.push(live_locals);
521 debug!("live_locals_anywhere = {:?}", live_locals_at_any_suspension_point);
523 // Renumber our liveness_map bitsets to include only the locals we are
525 let live_locals_at_suspension_points = live_locals_at_suspension_points
527 .map(|live_here| renumber_bitset(&live_here, &live_locals_at_any_suspension_point))
530 let storage_conflicts = compute_storage_conflicts(
532 &live_locals_at_any_suspension_point,
533 always_live_locals.clone(),
539 live_locals: live_locals_at_any_suspension_point,
540 live_locals_at_suspension_points,
542 storage_liveness: storage_liveness_map,
546 /// Renumbers the items present in `stored_locals` and applies the renumbering
549 /// For example, if `stored_locals = [1, 3, 5]`, this would be renumbered to
550 /// `[0, 1, 2]`. Thus, if `input = [3, 5]` we would return `[1, 2]`.
552 input: &BitSet<Local>,
553 stored_locals: &BitSet<Local>,
554 ) -> BitSet<GeneratorSavedLocal> {
555 assert!(stored_locals.superset(&input), "{:?} not a superset of {:?}", stored_locals, input);
556 let mut out = BitSet::new_empty(stored_locals.count());
557 for (idx, local) in stored_locals.iter().enumerate() {
558 let saved_local = GeneratorSavedLocal::from(idx);
559 if input.contains(local) {
560 out.insert(saved_local);
563 debug!("renumber_bitset({:?}, {:?}) => {:?}", input, stored_locals, out);
567 /// Record conflicts between locals at the current dataflow cursor positions.
569 /// You need to seek the cursors before calling this function.
570 fn record_conflicts_at_curr_loc(
571 local_conflicts: &mut BitMatrix<Local, Local>,
572 init: &dataflow::ResultsCursor<'mir, 'tcx, MaybeInitializedLocals>,
573 borrowed: &dataflow::ResultsCursor<'mir, 'tcx, MaybeBorrowedLocals>,
575 // A local requires storage if it is initialized or borrowed. For now, a local
576 // becomes uninitialized if it is moved from, but is still considered "borrowed".
578 // requires_storage := init | borrowed
580 // Just like when determining what locals are live at yield points, there is no need
581 // to look at storage liveness here, since `init | borrowed` is strictly more precise.
583 // FIXME: This function is called in a loop, so it might be better to pass in a temporary
584 // bitset rather than cloning here.
585 let mut requires_storage = init.get().clone();
586 requires_storage.union(borrowed.get());
588 for local in requires_storage.iter() {
589 local_conflicts.union_row_with(&requires_storage, local);
592 // `>1` because the `self` argument always requires storage.
593 if requires_storage.count() > 1 {
594 trace!("requires_storage={:?}", requires_storage);
598 /// For every saved local, looks for which locals are StorageLive at the same
599 /// time. Generates a bitset for every local of all the other locals that may be
600 /// StorageLive simultaneously with that local. This is used in the layout
601 /// computation; see `GeneratorLayout` for more.
602 fn compute_storage_conflicts(
603 body: &'mir Body<'tcx>,
604 stored_locals: &BitSet<Local>,
605 always_live_locals: storage::AlwaysLiveLocals,
606 mut init: dataflow::ResultsCursor<'mir, 'tcx, MaybeInitializedLocals>,
607 mut borrowed: dataflow::ResultsCursor<'mir, 'tcx, MaybeBorrowedLocals>,
608 ) -> BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal> {
609 debug!("compute_storage_conflicts({:?})", body.span);
610 assert_eq!(body.local_decls.len(), stored_locals.domain_size());
612 // Locals that are always live conflict with all other locals.
614 // FIXME: Why do we need to handle locals without `Storage{Live,Dead}` specially here?
615 // Shouldn't it be enough to know whether they are initialized?
616 let always_live_locals = always_live_locals.into_inner();
617 let mut local_conflicts = BitMatrix::from_row_n(&always_live_locals, body.local_decls.len());
619 // Visit every reachable statement and terminator. The exact order does not matter. When two
620 // locals are live at the same point in time, add an entry in the conflict matrix.
621 for (block, data) in traversal::preorder(body) {
622 // Ignore unreachable blocks.
623 if data.terminator().kind == TerminatorKind::Unreachable {
627 // Observe the dataflow state *before* all possible locations (statement or terminator) in
628 // each basic block...
629 for statement_index in 0..=data.statements.len() {
630 let loc = Location { block, statement_index };
631 trace!("record conflicts at {:?}", loc);
632 init.seek_before_primary_effect(loc);
633 borrowed.seek_before_primary_effect(loc);
634 record_conflicts_at_curr_loc(&mut local_conflicts, &init, &borrowed);
637 // ...and then observe the state *after* the terminator effect is applied. As long as
638 // neither `init` nor `borrowed` has a "before" effect, we will observe all possible
639 // dataflow states here or in the loop above.
640 trace!("record conflicts at end of {:?}", block);
641 init.seek_to_block_end(block);
642 borrowed.seek_to_block_end(block);
643 record_conflicts_at_curr_loc(&mut local_conflicts, &init, &borrowed);
646 // Compress the matrix using only stored locals (Local -> GeneratorSavedLocal).
648 // NOTE: Today we store a full conflict bitset for every local. Technically
649 // this is twice as many bits as we need, since the relation is symmetric.
650 // However, in practice these bitsets are not usually large. The layout code
651 // also needs to keep track of how many conflicts each local has, so it's
652 // simpler to keep it this way for now.
653 let mut storage_conflicts = BitMatrix::new(stored_locals.count(), stored_locals.count());
654 for (idx_a, local_a) in stored_locals.iter().enumerate() {
655 let saved_local_a = GeneratorSavedLocal::new(idx_a);
656 if always_live_locals.contains(local_a) {
657 // Conflicts with everything.
658 storage_conflicts.insert_all_into_row(saved_local_a);
660 // Keep overlap information only for stored locals.
661 for (idx_b, local_b) in stored_locals.iter().enumerate() {
662 let saved_local_b = GeneratorSavedLocal::new(idx_b);
663 if local_conflicts.contains(local_a, local_b) {
664 storage_conflicts.insert(saved_local_a, saved_local_b);
672 /// Validates the typeck view of the generator against the actual set of types retained between
674 fn sanitize_witness<'tcx>(
679 upvars: &Vec<Ty<'tcx>>,
680 retained: &BitSet<Local>,
682 let allowed_upvars = tcx.erase_regions(upvars);
683 let allowed = match witness.kind {
684 ty::GeneratorWitness(s) => tcx.erase_late_bound_regions(&s),
686 tcx.sess.delay_span_bug(
688 &format!("unexpected generator witness type {:?}", witness.kind),
694 let param_env = tcx.param_env(did);
696 for (local, decl) in body.local_decls.iter_enumerated() {
697 // Ignore locals which are internal or not retained between yields.
698 if !retained.contains(local) || decl.internal {
701 let decl_ty = tcx.normalize_erasing_regions(param_env, decl.ty);
703 // Sanity check that typeck knows about the type of locals which are
704 // live across a suspension point
705 if !allowed.contains(&decl_ty) && !allowed_upvars.contains(&decl_ty) {
708 "Broken MIR: generator contains type {} in MIR, \
709 but typeck only knows about {}",
717 fn compute_layout<'tcx>(
719 source: MirSource<'tcx>,
720 upvars: &Vec<Ty<'tcx>>,
722 always_live_locals: &storage::AlwaysLiveLocals,
724 body: &mut Body<'tcx>,
726 FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
727 GeneratorLayout<'tcx>,
728 IndexVec<BasicBlock, Option<BitSet<Local>>>,
730 // Use a liveness analysis to compute locals which are live across a suspension point
733 live_locals_at_suspension_points,
736 } = locals_live_across_suspend_points(tcx, body, source, always_live_locals, movable);
738 sanitize_witness(tcx, body, source.def_id(), interior, upvars, &live_locals);
740 // Gather live local types and their indices.
741 let mut locals = IndexVec::<GeneratorSavedLocal, _>::new();
742 let mut tys = IndexVec::<GeneratorSavedLocal, _>::new();
743 for (idx, local) in live_locals.iter().enumerate() {
745 tys.push(body.local_decls[local].ty);
746 debug!("generator saved local {:?} => {:?}", GeneratorSavedLocal::from(idx), local);
749 // Leave empty variants for the UNRESUMED, RETURNED, and POISONED states.
750 const RESERVED_VARIANTS: usize = 3;
752 // Build the generator variant field list.
753 // Create a map from local indices to generator struct indices.
754 let mut variant_fields: IndexVec<VariantIdx, IndexVec<Field, GeneratorSavedLocal>> =
755 iter::repeat(IndexVec::new()).take(RESERVED_VARIANTS).collect();
756 let mut remap = FxHashMap::default();
757 for (suspension_point_idx, live_locals) in live_locals_at_suspension_points.iter().enumerate() {
758 let variant_index = VariantIdx::from(RESERVED_VARIANTS + suspension_point_idx);
759 let mut fields = IndexVec::new();
760 for (idx, saved_local) in live_locals.iter().enumerate() {
761 fields.push(saved_local);
762 // Note that if a field is included in multiple variants, we will
763 // just use the first one here. That's fine; fields do not move
764 // around inside generators, so it doesn't matter which variant
765 // index we access them by.
766 remap.entry(locals[saved_local]).or_insert((tys[saved_local], variant_index, idx));
768 variant_fields.push(fields);
770 debug!("generator variant_fields = {:?}", variant_fields);
771 debug!("generator storage_conflicts = {:#?}", storage_conflicts);
773 let layout = GeneratorLayout { field_tys: tys, variant_fields, storage_conflicts };
775 (remap, layout, storage_liveness)
778 /// Replaces the entry point of `body` with a block that switches on the generator discriminant and
779 /// dispatches to blocks according to `cases`.
781 /// After this function, the former entry point of the function will be bb1.
782 fn insert_switch<'tcx>(
783 body: &mut Body<'tcx>,
784 cases: Vec<(usize, BasicBlock)>,
785 transform: &TransformVisitor<'tcx>,
786 default: TerminatorKind<'tcx>,
788 let default_block = insert_term_block(body, default);
789 let (assign, discr) = transform.get_discr(body);
790 let switch = TerminatorKind::SwitchInt {
791 discr: Operand::Move(discr),
792 switch_ty: transform.discr_ty,
793 values: Cow::from(cases.iter().map(|&(i, _)| i as u128).collect::<Vec<_>>()),
794 targets: cases.iter().map(|&(_, d)| d).chain(iter::once(default_block)).collect(),
797 let source_info = SourceInfo::outermost(body.span);
798 body.basic_blocks_mut().raw.insert(
801 statements: vec![assign],
802 terminator: Some(Terminator { source_info, kind: switch }),
807 let blocks = body.basic_blocks_mut().iter_mut();
809 for target in blocks.flat_map(|b| b.terminator_mut().successors_mut()) {
810 *target = BasicBlock::new(target.index() + 1);
814 fn elaborate_generator_drops<'tcx>(tcx: TyCtxt<'tcx>, def_id: DefId, body: &mut Body<'tcx>) {
815 use crate::shim::DropShimElaborator;
816 use crate::util::elaborate_drops::{elaborate_drop, Unwind};
817 use crate::util::patch::MirPatch;
819 // Note that `elaborate_drops` only drops the upvars of a generator, and
820 // this is ok because `open_drop` can only be reached within that own
821 // generator's resume function.
823 let param_env = tcx.param_env(def_id);
825 let mut elaborator = DropShimElaborator { body, patch: MirPatch::new(body), tcx, param_env };
827 for (block, block_data) in body.basic_blocks().iter_enumerated() {
828 let (target, unwind, source_info) = match block_data.terminator() {
829 Terminator { source_info, kind: TerminatorKind::Drop { location, target, unwind } } => {
830 if let Some(local) = location.as_local() {
831 if local == SELF_ARG {
832 (target, unwind, source_info)
842 let unwind = if block_data.is_cleanup {
845 Unwind::To(unwind.unwrap_or_else(|| elaborator.patch.resume_block()))
850 Place::from(SELF_ARG),
857 elaborator.patch.apply(body);
860 fn create_generator_drop_shim<'tcx>(
862 transform: &TransformVisitor<'tcx>,
863 source: MirSource<'tcx>,
865 body: &mut Body<'tcx>,
866 drop_clean: BasicBlock,
868 let mut body = body.clone();
869 body.arg_count = 1; // make sure the resume argument is not included here
871 let source_info = SourceInfo::outermost(body.span);
873 let mut cases = create_cases(&mut body, transform, Operation::Drop);
875 cases.insert(0, (UNRESUMED, drop_clean));
877 // The returned state and the poisoned state fall through to the default
878 // case which is just to return
880 insert_switch(&mut body, cases, &transform, TerminatorKind::Return);
882 for block in body.basic_blocks_mut() {
883 let kind = &mut block.terminator_mut().kind;
884 if let TerminatorKind::GeneratorDrop = *kind {
885 *kind = TerminatorKind::Return;
889 // Replace the return variable
890 body.local_decls[RETURN_PLACE] = LocalDecl::with_source_info(tcx.mk_unit(), source_info);
892 make_generator_state_argument_indirect(tcx, &mut body);
894 // Change the generator argument from &mut to *mut
895 body.local_decls[SELF_ARG] = LocalDecl::with_source_info(
896 tcx.mk_ptr(ty::TypeAndMut { ty: gen_ty, mutbl: hir::Mutability::Mut }),
899 if tcx.sess.opts.debugging_opts.mir_emit_retag {
900 // Alias tracking must know we changed the type
901 body.basic_blocks_mut()[START_BLOCK].statements.insert(
905 kind: StatementKind::Retag(RetagKind::Raw, box Place::from(SELF_ARG)),
910 no_landing_pads(tcx, &mut body);
912 // Make sure we remove dead blocks to remove
913 // unrelated code from the resume part of the function
914 simplify::remove_dead_blocks(&mut body);
916 dump_mir(tcx, None, "generator_drop", &0, source, &body, |_, _| Ok(()));
921 fn insert_term_block<'tcx>(body: &mut Body<'tcx>, kind: TerminatorKind<'tcx>) -> BasicBlock {
922 let source_info = SourceInfo::outermost(body.span);
923 body.basic_blocks_mut().push(BasicBlockData {
924 statements: Vec::new(),
925 terminator: Some(Terminator { source_info, kind }),
930 fn insert_panic_block<'tcx>(
932 body: &mut Body<'tcx>,
933 message: AssertMessage<'tcx>,
935 let assert_block = BasicBlock::new(body.basic_blocks().len());
936 let term = TerminatorKind::Assert {
937 cond: Operand::Constant(box Constant {
940 literal: ty::Const::from_bool(tcx, false),
944 target: assert_block,
948 let source_info = SourceInfo::outermost(body.span);
949 body.basic_blocks_mut().push(BasicBlockData {
950 statements: Vec::new(),
951 terminator: Some(Terminator { source_info, kind: term }),
958 fn can_return<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) -> bool {
959 // Returning from a function with an uninhabited return type is undefined behavior.
960 if body.return_ty().conservative_is_privately_uninhabited(tcx) {
964 // If there's a return terminator the function may return.
965 for block in body.basic_blocks() {
966 if let TerminatorKind::Return = block.terminator().kind {
971 // Otherwise the function can't return.
975 fn can_unwind<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) -> bool {
976 // Nothing can unwind when landing pads are off.
977 if tcx.sess.panic_strategy() == PanicStrategy::Abort {
981 // Unwinds can only start at certain terminators.
982 for block in body.basic_blocks() {
983 match block.terminator().kind {
984 // These never unwind.
985 TerminatorKind::Goto { .. }
986 | TerminatorKind::SwitchInt { .. }
987 | TerminatorKind::Abort
988 | TerminatorKind::Return
989 | TerminatorKind::Unreachable
990 | TerminatorKind::GeneratorDrop
991 | TerminatorKind::FalseEdges { .. }
992 | TerminatorKind::FalseUnwind { .. }
993 | TerminatorKind::InlineAsm { .. } => {}
995 // Resume will *continue* unwinding, but if there's no other unwinding terminator it
996 // will never be reached.
997 TerminatorKind::Resume => {}
999 TerminatorKind::Yield { .. } => {
1000 unreachable!("`can_unwind` called before generator transform")
1003 // These may unwind.
1004 TerminatorKind::Drop { .. }
1005 | TerminatorKind::DropAndReplace { .. }
1006 | TerminatorKind::Call { .. }
1007 | TerminatorKind::Assert { .. } => return true,
1011 // If we didn't find an unwinding terminator, the function cannot unwind.
1015 fn create_generator_resume_function<'tcx>(
1017 transform: TransformVisitor<'tcx>,
1018 source: MirSource<'tcx>,
1019 body: &mut Body<'tcx>,
1022 let can_unwind = can_unwind(tcx, body);
1024 // Poison the generator when it unwinds
1026 let source_info = SourceInfo::outermost(body.span);
1027 let poison_block = body.basic_blocks_mut().push(BasicBlockData {
1028 statements: vec![transform.set_discr(VariantIdx::new(POISONED), source_info)],
1029 terminator: Some(Terminator { source_info, kind: TerminatorKind::Resume }),
1033 for (idx, block) in body.basic_blocks_mut().iter_enumerated_mut() {
1034 let source_info = block.terminator().source_info;
1036 if let TerminatorKind::Resume = block.terminator().kind {
1037 // An existing `Resume` terminator is redirected to jump to our dedicated
1038 // "poisoning block" above.
1039 if idx != poison_block {
1040 *block.terminator_mut() = Terminator {
1042 kind: TerminatorKind::Goto { target: poison_block },
1045 } else if !block.is_cleanup {
1046 // Any terminators that *can* unwind but don't have an unwind target set are also
1047 // pointed at our poisoning block (unless they're part of the cleanup path).
1048 if let Some(unwind @ None) = block.terminator_mut().unwind_mut() {
1049 *unwind = Some(poison_block);
1055 let mut cases = create_cases(body, &transform, Operation::Resume);
1057 use rustc_middle::mir::AssertKind::{ResumedAfterPanic, ResumedAfterReturn};
1059 // Jump to the entry point on the unresumed
1060 cases.insert(0, (UNRESUMED, BasicBlock::new(0)));
1062 // Panic when resumed on the returned or poisoned state
1063 let generator_kind = body.generator_kind.unwrap();
1068 (POISONED, insert_panic_block(tcx, body, ResumedAfterPanic(generator_kind))),
1075 (RETURNED, insert_panic_block(tcx, body, ResumedAfterReturn(generator_kind))),
1079 insert_switch(body, cases, &transform, TerminatorKind::Unreachable);
1081 make_generator_state_argument_indirect(tcx, body);
1082 make_generator_state_argument_pinned(tcx, body);
1084 no_landing_pads(tcx, body);
1086 // Make sure we remove dead blocks to remove
1087 // unrelated code from the drop part of the function
1088 simplify::remove_dead_blocks(body);
1090 dump_mir(tcx, None, "generator_resume", &0, source, body, |_, _| Ok(()));
1093 fn insert_clean_drop(body: &mut Body<'_>) -> BasicBlock {
1094 let return_block = insert_term_block(body, TerminatorKind::Return);
1096 let term = TerminatorKind::Drop {
1097 location: Place::from(SELF_ARG),
1098 target: return_block,
1101 let source_info = SourceInfo::outermost(body.span);
1103 // Create a block to destroy an unresumed generators. This can only destroy upvars.
1104 body.basic_blocks_mut().push(BasicBlockData {
1105 statements: Vec::new(),
1106 terminator: Some(Terminator { source_info, kind: term }),
1111 /// An operation that can be performed on a generator.
1112 #[derive(PartialEq, Copy, Clone)]
1119 fn target_block(self, point: &SuspensionPoint<'_>) -> Option<BasicBlock> {
1121 Operation::Resume => Some(point.resume),
1122 Operation::Drop => point.drop,
1127 fn create_cases<'tcx>(
1128 body: &mut Body<'tcx>,
1129 transform: &TransformVisitor<'tcx>,
1130 operation: Operation,
1131 ) -> Vec<(usize, BasicBlock)> {
1132 let source_info = SourceInfo::outermost(body.span);
1137 .filter_map(|point| {
1138 // Find the target for this suspension point, if applicable
1139 operation.target_block(point).map(|target| {
1140 let mut statements = Vec::new();
1142 // Create StorageLive instructions for locals with live storage
1143 for i in 0..(body.local_decls.len()) {
1145 // The resume argument is live on function entry. Don't insert a
1146 // `StorageLive`, or the following `Assign` will read from uninitialized
1151 let l = Local::new(i);
1152 let needs_storage_live = point.storage_liveness.contains(l)
1153 && !transform.remap.contains_key(&l)
1154 && !transform.always_live_locals.contains(l);
1155 if needs_storage_live {
1157 .push(Statement { source_info, kind: StatementKind::StorageLive(l) });
1161 if operation == Operation::Resume {
1162 // Move the resume argument to the destination place of the `Yield` terminator
1163 let resume_arg = Local::new(2); // 0 = return, 1 = self
1164 statements.push(Statement {
1166 kind: StatementKind::Assign(box (
1168 Rvalue::Use(Operand::Move(resume_arg.into())),
1173 // Then jump to the real target
1174 let block = body.basic_blocks_mut().push(BasicBlockData {
1176 terminator: Some(Terminator {
1178 kind: TerminatorKind::Goto { target },
1183 (point.state, block)
1189 impl<'tcx> MirPass<'tcx> for StateTransform {
1190 fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut Body<'tcx>) {
1191 let yield_ty = if let Some(yield_ty) = body.yield_ty {
1194 // This only applies to generators
1198 assert!(body.generator_drop.is_none());
1200 let def_id = source.def_id();
1202 // The first argument is the generator type passed by value
1203 let gen_ty = body.local_decls.raw[1].ty;
1205 // Get the interior types and substs which typeck computed
1206 let (upvars, interior, discr_ty, movable) = match gen_ty.kind {
1207 ty::Generator(_, substs, movability) => {
1208 let substs = substs.as_generator();
1210 substs.upvar_tys().collect(),
1212 substs.discr_ty(tcx),
1213 movability == hir::Movability::Movable,
1219 // Compute GeneratorState<yield_ty, return_ty>
1220 let state_did = tcx.require_lang_item(GeneratorStateLangItem, None);
1221 let state_adt_ref = tcx.adt_def(state_did);
1222 let state_substs = tcx.intern_substs(&[yield_ty.into(), body.return_ty().into()]);
1223 let ret_ty = tcx.mk_adt(state_adt_ref, state_substs);
1225 // We rename RETURN_PLACE which has type mir.return_ty to new_ret_local
1226 // RETURN_PLACE then is a fresh unused local with type ret_ty.
1227 let new_ret_local = replace_local(RETURN_PLACE, ret_ty, body, tcx);
1229 // We also replace the resume argument and insert an `Assign`.
1230 // This is needed because the resume argument `_2` might be live across a `yield`, in which
1231 // case there is no `Assign` to it that the transform can turn into a store to the generator
1232 // state. After the yield the slot in the generator state would then be uninitialized.
1233 let resume_local = Local::new(2);
1234 let new_resume_local =
1235 replace_local(resume_local, body.local_decls[resume_local].ty, body, tcx);
1237 // When first entering the generator, move the resume argument into its new local.
1238 let source_info = SourceInfo::outermost(body.span);
1239 let stmts = &mut body.basic_blocks_mut()[BasicBlock::new(0)].statements;
1244 kind: StatementKind::Assign(box (
1245 new_resume_local.into(),
1246 Rvalue::Use(Operand::Move(resume_local.into())),
1251 let always_live_locals = storage::AlwaysLiveLocals::new(&body);
1253 // Extract locals which are live across suspension point into `layout`
1254 // `remap` gives a mapping from local indices onto generator struct indices
1255 // `storage_liveness` tells us which locals have live storage at suspension points
1256 let (remap, layout, storage_liveness) =
1257 compute_layout(tcx, source, &upvars, interior, &always_live_locals, movable, body);
1259 let can_return = can_return(tcx, body);
1261 // Run the transformation which converts Places from Local to generator struct
1262 // accesses for locals in `remap`.
1263 // It also rewrites `return x` and `yield y` as writing a new generator state and returning
1264 // GeneratorState::Complete(x) and GeneratorState::Yielded(y) respectively.
1265 let mut transform = TransformVisitor {
1272 suspension_points: Vec::new(),
1276 transform.visit_body(body);
1278 // Update our MIR struct to reflect the changes we've made
1279 body.yield_ty = None;
1280 body.arg_count = 2; // self, resume arg
1281 body.spread_arg = None;
1282 body.generator_layout = Some(layout);
1284 // Insert `drop(generator_struct)` which is used to drop upvars for generators in
1285 // the unresumed state.
1286 // This is expanded to a drop ladder in `elaborate_generator_drops`.
1287 let drop_clean = insert_clean_drop(body);
1289 dump_mir(tcx, None, "generator_pre-elab", &0, source, body, |_, _| Ok(()));
1291 // Expand `drop(generator_struct)` to a drop ladder which destroys upvars.
1292 // If any upvars are moved out of, drop elaboration will handle upvar destruction.
1293 // However we need to also elaborate the code generated by `insert_clean_drop`.
1294 elaborate_generator_drops(tcx, def_id, body);
1296 dump_mir(tcx, None, "generator_post-transform", &0, source, body, |_, _| Ok(()));
1298 // Create a copy of our MIR and use it to create the drop shim for the generator
1300 create_generator_drop_shim(tcx, &transform, source, gen_ty, body, drop_clean);
1302 body.generator_drop = Some(box drop_shim);
1304 // Create the Generator::resume function
1305 create_generator_resume_function(tcx, transform, source, body, can_return);