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::simplify;
53 use crate::util::expand_aggregate;
54 use crate::MirPass;
55 use rustc_data_structures::fx::FxHashMap;
56 use rustc_hir as hir;
57 use rustc_hir::lang_items::LangItem;
58 use rustc_index::bit_set::{BitMatrix, BitSet};
59 use rustc_index::vec::{Idx, IndexVec};
60 use rustc_middle::mir::dump_mir;
61 use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
62 use rustc_middle::mir::*;
63 use rustc_middle::ty::subst::{Subst, SubstsRef};
64 use rustc_middle::ty::GeneratorSubsts;
65 use rustc_middle::ty::{self, AdtDef, Ty, TyCtxt};
66 use rustc_mir_dataflow::impls::{
67 MaybeBorrowedLocals, MaybeLiveLocals, MaybeRequiresStorage, MaybeStorageLive,
68 };
69 use rustc_mir_dataflow::storage;
70 use rustc_mir_dataflow::{self, Analysis};
71 use rustc_target::abi::VariantIdx;
72 use rustc_target::spec::PanicStrategy;
73 use std::{iter, ops};
74
75 pub struct StateTransform;
76
77 struct RenameLocalVisitor<'tcx> {
78 from: Local,
79 to: Local,
80 tcx: TyCtxt<'tcx>,
81 }
82
83 impl<'tcx> MutVisitor<'tcx> for RenameLocalVisitor<'tcx> {
tcx(&self) -> TyCtxt<'tcx>84 fn tcx(&self) -> TyCtxt<'tcx> {
85 self.tcx
86 }
87
visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location)88 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
89 if *local == self.from {
90 *local = self.to;
91 }
92 }
93
visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location)94 fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) {
95 match terminator.kind {
96 TerminatorKind::Return => {
97 // Do not replace the implicit `_0` access here, as that's not possible. The
98 // transform already handles `return` correctly.
99 }
100 _ => self.super_terminator(terminator, location),
101 }
102 }
103 }
104
105 struct DerefArgVisitor<'tcx> {
106 tcx: TyCtxt<'tcx>,
107 }
108
109 impl<'tcx> MutVisitor<'tcx> for DerefArgVisitor<'tcx> {
tcx(&self) -> TyCtxt<'tcx>110 fn tcx(&self) -> TyCtxt<'tcx> {
111 self.tcx
112 }
113
visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location)114 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
115 assert_ne!(*local, SELF_ARG);
116 }
117
visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location)118 fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) {
119 if place.local == SELF_ARG {
120 replace_base(
121 place,
122 Place {
123 local: SELF_ARG,
124 projection: self.tcx().intern_place_elems(&[ProjectionElem::Deref]),
125 },
126 self.tcx,
127 );
128 } else {
129 self.visit_local(&mut place.local, context, location);
130
131 for elem in place.projection.iter() {
132 if let PlaceElem::Index(local) = elem {
133 assert_ne!(local, SELF_ARG);
134 }
135 }
136 }
137 }
138 }
139
140 struct PinArgVisitor<'tcx> {
141 ref_gen_ty: Ty<'tcx>,
142 tcx: TyCtxt<'tcx>,
143 }
144
145 impl<'tcx> MutVisitor<'tcx> for PinArgVisitor<'tcx> {
tcx(&self) -> TyCtxt<'tcx>146 fn tcx(&self) -> TyCtxt<'tcx> {
147 self.tcx
148 }
149
visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location)150 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
151 assert_ne!(*local, SELF_ARG);
152 }
153
visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location)154 fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) {
155 if place.local == SELF_ARG {
156 replace_base(
157 place,
158 Place {
159 local: SELF_ARG,
160 projection: self.tcx().intern_place_elems(&[ProjectionElem::Field(
161 Field::new(0),
162 self.ref_gen_ty,
163 )]),
164 },
165 self.tcx,
166 );
167 } else {
168 self.visit_local(&mut place.local, context, location);
169
170 for elem in place.projection.iter() {
171 if let PlaceElem::Index(local) = elem {
172 assert_ne!(local, SELF_ARG);
173 }
174 }
175 }
176 }
177 }
178
replace_base<'tcx>(place: &mut Place<'tcx>, new_base: Place<'tcx>, tcx: TyCtxt<'tcx>)179 fn replace_base<'tcx>(place: &mut Place<'tcx>, new_base: Place<'tcx>, tcx: TyCtxt<'tcx>) {
180 place.local = new_base.local;
181
182 let mut new_projection = new_base.projection.to_vec();
183 new_projection.append(&mut place.projection.to_vec());
184
185 place.projection = tcx.intern_place_elems(&new_projection);
186 }
187
188 const SELF_ARG: Local = Local::from_u32(1);
189
190 /// Generator has not been resumed yet.
191 const UNRESUMED: usize = GeneratorSubsts::UNRESUMED;
192 /// Generator has returned / is completed.
193 const RETURNED: usize = GeneratorSubsts::RETURNED;
194 /// Generator has panicked and is poisoned.
195 const POISONED: usize = GeneratorSubsts::POISONED;
196
197 /// A `yield` point in the generator.
198 struct SuspensionPoint<'tcx> {
199 /// State discriminant used when suspending or resuming at this point.
200 state: usize,
201 /// The block to jump to after resumption.
202 resume: BasicBlock,
203 /// Where to move the resume argument after resumption.
204 resume_arg: Place<'tcx>,
205 /// Which block to jump to if the generator is dropped in this state.
206 drop: Option<BasicBlock>,
207 /// Set of locals that have live storage while at this suspension point.
208 storage_liveness: BitSet<Local>,
209 }
210
211 struct TransformVisitor<'tcx> {
212 tcx: TyCtxt<'tcx>,
213 state_adt_ref: &'tcx AdtDef,
214 state_substs: SubstsRef<'tcx>,
215
216 // The type of the discriminant in the generator struct
217 discr_ty: Ty<'tcx>,
218
219 // Mapping from Local to (type of local, generator struct index)
220 // FIXME(eddyb) This should use `IndexVec<Local, Option<_>>`.
221 remap: FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
222
223 // A map from a suspension point in a block to the locals which have live storage at that point
224 storage_liveness: IndexVec<BasicBlock, Option<BitSet<Local>>>,
225
226 // A list of suspension points, generated during the transform
227 suspension_points: Vec<SuspensionPoint<'tcx>>,
228
229 // The set of locals that have no `StorageLive`/`StorageDead` annotations.
230 always_live_locals: storage::AlwaysLiveLocals,
231
232 // The original RETURN_PLACE local
233 new_ret_local: Local,
234 }
235
236 impl TransformVisitor<'tcx> {
237 // Make a GeneratorState variant assignment. `core::ops::GeneratorState` only has single
238 // element tuple variants, so we can just write to the downcasted first field and then set the
239 // discriminant to the appropriate variant.
make_state( &self, idx: VariantIdx, val: Operand<'tcx>, source_info: SourceInfo, ) -> impl Iterator<Item = Statement<'tcx>>240 fn make_state(
241 &self,
242 idx: VariantIdx,
243 val: Operand<'tcx>,
244 source_info: SourceInfo,
245 ) -> impl Iterator<Item = Statement<'tcx>> {
246 let kind = AggregateKind::Adt(self.state_adt_ref, idx, self.state_substs, None, None);
247 assert_eq!(self.state_adt_ref.variants[idx].fields.len(), 1);
248 let ty = self
249 .tcx
250 .type_of(self.state_adt_ref.variants[idx].fields[0].did)
251 .subst(self.tcx, self.state_substs);
252 expand_aggregate(
253 Place::return_place(),
254 std::iter::once((val, ty)),
255 kind,
256 source_info,
257 self.tcx,
258 )
259 }
260
261 // Create a Place referencing a generator struct field
make_field(&self, variant_index: VariantIdx, idx: usize, ty: Ty<'tcx>) -> Place<'tcx>262 fn make_field(&self, variant_index: VariantIdx, idx: usize, ty: Ty<'tcx>) -> Place<'tcx> {
263 let self_place = Place::from(SELF_ARG);
264 let base = self.tcx.mk_place_downcast_unnamed(self_place, variant_index);
265 let mut projection = base.projection.to_vec();
266 projection.push(ProjectionElem::Field(Field::new(idx), ty));
267
268 Place { local: base.local, projection: self.tcx.intern_place_elems(&projection) }
269 }
270
271 // Create a statement which changes the discriminant
set_discr(&self, state_disc: VariantIdx, source_info: SourceInfo) -> Statement<'tcx>272 fn set_discr(&self, state_disc: VariantIdx, source_info: SourceInfo) -> Statement<'tcx> {
273 let self_place = Place::from(SELF_ARG);
274 Statement {
275 source_info,
276 kind: StatementKind::SetDiscriminant {
277 place: Box::new(self_place),
278 variant_index: state_disc,
279 },
280 }
281 }
282
283 // Create a statement which reads the discriminant into a temporary
get_discr(&self, body: &mut Body<'tcx>) -> (Statement<'tcx>, Place<'tcx>)284 fn get_discr(&self, body: &mut Body<'tcx>) -> (Statement<'tcx>, Place<'tcx>) {
285 let temp_decl = LocalDecl::new(self.discr_ty, body.span).internal();
286 let local_decls_len = body.local_decls.push(temp_decl);
287 let temp = Place::from(local_decls_len);
288
289 let self_place = Place::from(SELF_ARG);
290 let assign = Statement {
291 source_info: SourceInfo::outermost(body.span),
292 kind: StatementKind::Assign(Box::new((temp, Rvalue::Discriminant(self_place)))),
293 };
294 (assign, temp)
295 }
296 }
297
298 impl MutVisitor<'tcx> for TransformVisitor<'tcx> {
tcx(&self) -> TyCtxt<'tcx>299 fn tcx(&self) -> TyCtxt<'tcx> {
300 self.tcx
301 }
302
visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location)303 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
304 assert_eq!(self.remap.get(local), None);
305 }
306
visit_place( &mut self, place: &mut Place<'tcx>, _context: PlaceContext, _location: Location, )307 fn visit_place(
308 &mut self,
309 place: &mut Place<'tcx>,
310 _context: PlaceContext,
311 _location: Location,
312 ) {
313 // Replace an Local in the remap with a generator struct access
314 if let Some(&(ty, variant_index, idx)) = self.remap.get(&place.local) {
315 replace_base(place, self.make_field(variant_index, idx, ty), self.tcx);
316 }
317 }
318
visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>)319 fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) {
320 // Remove StorageLive and StorageDead statements for remapped locals
321 data.retain_statements(|s| match s.kind {
322 StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => {
323 !self.remap.contains_key(&l)
324 }
325 _ => true,
326 });
327
328 let ret_val = match data.terminator().kind {
329 TerminatorKind::Return => Some((
330 VariantIdx::new(1),
331 None,
332 Operand::Move(Place::from(self.new_ret_local)),
333 None,
334 )),
335 TerminatorKind::Yield { ref value, resume, resume_arg, drop } => {
336 Some((VariantIdx::new(0), Some((resume, resume_arg)), value.clone(), drop))
337 }
338 _ => None,
339 };
340
341 if let Some((state_idx, resume, v, drop)) = ret_val {
342 let source_info = data.terminator().source_info;
343 // We must assign the value first in case it gets declared dead below
344 data.statements.extend(self.make_state(state_idx, v, source_info));
345 let state = if let Some((resume, mut resume_arg)) = resume {
346 // Yield
347 let state = 3 + self.suspension_points.len();
348
349 // The resume arg target location might itself be remapped if its base local is
350 // live across a yield.
351 let resume_arg =
352 if let Some(&(ty, variant, idx)) = self.remap.get(&resume_arg.local) {
353 replace_base(&mut resume_arg, self.make_field(variant, idx, ty), self.tcx);
354 resume_arg
355 } else {
356 resume_arg
357 };
358
359 self.suspension_points.push(SuspensionPoint {
360 state,
361 resume,
362 resume_arg,
363 drop,
364 storage_liveness: self.storage_liveness[block].clone().unwrap(),
365 });
366
367 VariantIdx::new(state)
368 } else {
369 // Return
370 VariantIdx::new(RETURNED) // state for returned
371 };
372 data.statements.push(self.set_discr(state, source_info));
373 data.terminator_mut().kind = TerminatorKind::Return;
374 }
375
376 self.super_basic_block_data(block, data);
377 }
378 }
379
make_generator_state_argument_indirect<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>)380 fn make_generator_state_argument_indirect<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
381 let gen_ty = body.local_decls.raw[1].ty;
382
383 let ref_gen_ty =
384 tcx.mk_ref(tcx.lifetimes.re_erased, ty::TypeAndMut { ty: gen_ty, mutbl: Mutability::Mut });
385
386 // Replace the by value generator argument
387 body.local_decls.raw[1].ty = ref_gen_ty;
388
389 // Add a deref to accesses of the generator state
390 DerefArgVisitor { tcx }.visit_body(body);
391 }
392
make_generator_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>)393 fn make_generator_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
394 let ref_gen_ty = body.local_decls.raw[1].ty;
395
396 let pin_did = tcx.require_lang_item(LangItem::Pin, Some(body.span));
397 let pin_adt_ref = tcx.adt_def(pin_did);
398 let substs = tcx.intern_substs(&[ref_gen_ty.into()]);
399 let pin_ref_gen_ty = tcx.mk_adt(pin_adt_ref, substs);
400
401 // Replace the by ref generator argument
402 body.local_decls.raw[1].ty = pin_ref_gen_ty;
403
404 // Add the Pin field access to accesses of the generator state
405 PinArgVisitor { ref_gen_ty, tcx }.visit_body(body);
406 }
407
408 /// Allocates a new local and replaces all references of `local` with it. Returns the new local.
409 ///
410 /// `local` will be changed to a new local decl with type `ty`.
411 ///
412 /// Note that the new local will be uninitialized. It is the caller's responsibility to assign some
413 /// valid value to it before its first use.
replace_local<'tcx>( local: Local, ty: Ty<'tcx>, body: &mut Body<'tcx>, tcx: TyCtxt<'tcx>, ) -> Local414 fn replace_local<'tcx>(
415 local: Local,
416 ty: Ty<'tcx>,
417 body: &mut Body<'tcx>,
418 tcx: TyCtxt<'tcx>,
419 ) -> Local {
420 let new_decl = LocalDecl::new(ty, body.span);
421 let new_local = body.local_decls.push(new_decl);
422 body.local_decls.swap(local, new_local);
423
424 RenameLocalVisitor { from: local, to: new_local, tcx }.visit_body(body);
425
426 new_local
427 }
428
429 struct LivenessInfo {
430 /// Which locals are live across any suspension point.
431 saved_locals: GeneratorSavedLocals,
432
433 /// The set of saved locals live at each suspension point.
434 live_locals_at_suspension_points: Vec<BitSet<GeneratorSavedLocal>>,
435
436 /// Parallel vec to the above with SourceInfo for each yield terminator.
437 source_info_at_suspension_points: Vec<SourceInfo>,
438
439 /// For every saved local, the set of other saved locals that are
440 /// storage-live at the same time as this local. We cannot overlap locals in
441 /// the layout which have conflicting storage.
442 storage_conflicts: BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal>,
443
444 /// For every suspending block, the locals which are storage-live across
445 /// that suspension point.
446 storage_liveness: IndexVec<BasicBlock, Option<BitSet<Local>>>,
447 }
448
locals_live_across_suspend_points( tcx: TyCtxt<'tcx>, body: &Body<'tcx>, always_live_locals: &storage::AlwaysLiveLocals, movable: bool, ) -> LivenessInfo449 fn locals_live_across_suspend_points(
450 tcx: TyCtxt<'tcx>,
451 body: &Body<'tcx>,
452 always_live_locals: &storage::AlwaysLiveLocals,
453 movable: bool,
454 ) -> LivenessInfo {
455 let body_ref: &Body<'_> = &body;
456
457 // Calculate when MIR locals have live storage. This gives us an upper bound of their
458 // lifetimes.
459 let mut storage_live = MaybeStorageLive::new(always_live_locals.clone())
460 .into_engine(tcx, body_ref)
461 .iterate_to_fixpoint()
462 .into_results_cursor(body_ref);
463
464 // Calculate the MIR locals which have been previously
465 // borrowed (even if they are still active).
466 let borrowed_locals_results = MaybeBorrowedLocals::all_borrows()
467 .into_engine(tcx, body_ref)
468 .pass_name("generator")
469 .iterate_to_fixpoint();
470
471 let mut borrowed_locals_cursor =
472 rustc_mir_dataflow::ResultsCursor::new(body_ref, &borrowed_locals_results);
473
474 // Calculate the MIR locals that we actually need to keep storage around
475 // for.
476 let requires_storage_results = MaybeRequiresStorage::new(body, &borrowed_locals_results)
477 .into_engine(tcx, body_ref)
478 .iterate_to_fixpoint();
479 let mut requires_storage_cursor =
480 rustc_mir_dataflow::ResultsCursor::new(body_ref, &requires_storage_results);
481
482 // Calculate the liveness of MIR locals ignoring borrows.
483 let mut liveness = MaybeLiveLocals
484 .into_engine(tcx, body_ref)
485 .pass_name("generator")
486 .iterate_to_fixpoint()
487 .into_results_cursor(body_ref);
488
489 let mut storage_liveness_map = IndexVec::from_elem(None, body.basic_blocks());
490 let mut live_locals_at_suspension_points = Vec::new();
491 let mut source_info_at_suspension_points = Vec::new();
492 let mut live_locals_at_any_suspension_point = BitSet::new_empty(body.local_decls.len());
493
494 for (block, data) in body.basic_blocks().iter_enumerated() {
495 if let TerminatorKind::Yield { .. } = data.terminator().kind {
496 let loc = Location { block, statement_index: data.statements.len() };
497
498 liveness.seek_to_block_end(block);
499 let mut live_locals = liveness.get().clone();
500
501 if !movable {
502 // The `liveness` variable contains the liveness of MIR locals ignoring borrows.
503 // This is correct for movable generators since borrows cannot live across
504 // suspension points. However for immovable generators we need to account for
505 // borrows, so we conseratively assume that all borrowed locals are live until
506 // we find a StorageDead statement referencing the locals.
507 // To do this we just union our `liveness` result with `borrowed_locals`, which
508 // contains all the locals which has been borrowed before this suspension point.
509 // If a borrow is converted to a raw reference, we must also assume that it lives
510 // forever. Note that the final liveness is still bounded by the storage liveness
511 // of the local, which happens using the `intersect` operation below.
512 borrowed_locals_cursor.seek_before_primary_effect(loc);
513 live_locals.union(borrowed_locals_cursor.get());
514 }
515
516 // Store the storage liveness for later use so we can restore the state
517 // after a suspension point
518 storage_live.seek_before_primary_effect(loc);
519 storage_liveness_map[block] = Some(storage_live.get().clone());
520
521 // Locals live are live at this point only if they are used across
522 // suspension points (the `liveness` variable)
523 // and their storage is required (the `storage_required` variable)
524 requires_storage_cursor.seek_before_primary_effect(loc);
525 live_locals.intersect(requires_storage_cursor.get());
526
527 // The generator argument is ignored.
528 live_locals.remove(SELF_ARG);
529
530 debug!("loc = {:?}, live_locals = {:?}", loc, live_locals);
531
532 // Add the locals live at this suspension point to the set of locals which live across
533 // any suspension points
534 live_locals_at_any_suspension_point.union(&live_locals);
535
536 live_locals_at_suspension_points.push(live_locals);
537 source_info_at_suspension_points.push(data.terminator().source_info);
538 }
539 }
540
541 debug!("live_locals_anywhere = {:?}", live_locals_at_any_suspension_point);
542 let saved_locals = GeneratorSavedLocals(live_locals_at_any_suspension_point);
543
544 // Renumber our liveness_map bitsets to include only the locals we are
545 // saving.
546 let live_locals_at_suspension_points = live_locals_at_suspension_points
547 .iter()
548 .map(|live_here| saved_locals.renumber_bitset(&live_here))
549 .collect();
550
551 let storage_conflicts = compute_storage_conflicts(
552 body_ref,
553 &saved_locals,
554 always_live_locals.clone(),
555 requires_storage_results,
556 );
557
558 LivenessInfo {
559 saved_locals,
560 live_locals_at_suspension_points,
561 source_info_at_suspension_points,
562 storage_conflicts,
563 storage_liveness: storage_liveness_map,
564 }
565 }
566
567 /// The set of `Local`s that must be saved across yield points.
568 ///
569 /// `GeneratorSavedLocal` is indexed in terms of the elements in this set;
570 /// i.e. `GeneratorSavedLocal::new(1)` corresponds to the second local
571 /// included in this set.
572 struct GeneratorSavedLocals(BitSet<Local>);
573
574 impl GeneratorSavedLocals {
575 /// Returns an iterator over each `GeneratorSavedLocal` along with the `Local` it corresponds
576 /// to.
iter_enumerated(&self) -> impl '_ + Iterator<Item = (GeneratorSavedLocal, Local)>577 fn iter_enumerated(&self) -> impl '_ + Iterator<Item = (GeneratorSavedLocal, Local)> {
578 self.iter().enumerate().map(|(i, l)| (GeneratorSavedLocal::from(i), l))
579 }
580
581 /// Transforms a `BitSet<Local>` that contains only locals saved across yield points to the
582 /// equivalent `BitSet<GeneratorSavedLocal>`.
renumber_bitset(&self, input: &BitSet<Local>) -> BitSet<GeneratorSavedLocal>583 fn renumber_bitset(&self, input: &BitSet<Local>) -> BitSet<GeneratorSavedLocal> {
584 assert!(self.superset(&input), "{:?} not a superset of {:?}", self.0, input);
585 let mut out = BitSet::new_empty(self.count());
586 for (saved_local, local) in self.iter_enumerated() {
587 if input.contains(local) {
588 out.insert(saved_local);
589 }
590 }
591 out
592 }
593
get(&self, local: Local) -> Option<GeneratorSavedLocal>594 fn get(&self, local: Local) -> Option<GeneratorSavedLocal> {
595 if !self.contains(local) {
596 return None;
597 }
598
599 let idx = self.iter().take_while(|&l| l < local).count();
600 Some(GeneratorSavedLocal::new(idx))
601 }
602 }
603
604 impl ops::Deref for GeneratorSavedLocals {
605 type Target = BitSet<Local>;
606
deref(&self) -> &Self::Target607 fn deref(&self) -> &Self::Target {
608 &self.0
609 }
610 }
611
612 /// For every saved local, looks for which locals are StorageLive at the same
613 /// time. Generates a bitset for every local of all the other locals that may be
614 /// StorageLive simultaneously with that local. This is used in the layout
615 /// computation; see `GeneratorLayout` for more.
compute_storage_conflicts( body: &'mir Body<'tcx>, saved_locals: &GeneratorSavedLocals, always_live_locals: storage::AlwaysLiveLocals, requires_storage: rustc_mir_dataflow::Results<'tcx, MaybeRequiresStorage<'mir, 'tcx>>, ) -> BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal>616 fn compute_storage_conflicts(
617 body: &'mir Body<'tcx>,
618 saved_locals: &GeneratorSavedLocals,
619 always_live_locals: storage::AlwaysLiveLocals,
620 requires_storage: rustc_mir_dataflow::Results<'tcx, MaybeRequiresStorage<'mir, 'tcx>>,
621 ) -> BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal> {
622 assert_eq!(body.local_decls.len(), saved_locals.domain_size());
623
624 debug!("compute_storage_conflicts({:?})", body.span);
625 debug!("always_live = {:?}", always_live_locals);
626
627 // Locals that are always live or ones that need to be stored across
628 // suspension points are not eligible for overlap.
629 let mut ineligible_locals = always_live_locals.into_inner();
630 ineligible_locals.intersect(&**saved_locals);
631
632 // Compute the storage conflicts for all eligible locals.
633 let mut visitor = StorageConflictVisitor {
634 body,
635 saved_locals: &saved_locals,
636 local_conflicts: BitMatrix::from_row_n(&ineligible_locals, body.local_decls.len()),
637 };
638
639 requires_storage.visit_reachable_with(body, &mut visitor);
640
641 let local_conflicts = visitor.local_conflicts;
642
643 // Compress the matrix using only stored locals (Local -> GeneratorSavedLocal).
644 //
645 // NOTE: Today we store a full conflict bitset for every local. Technically
646 // this is twice as many bits as we need, since the relation is symmetric.
647 // However, in practice these bitsets are not usually large. The layout code
648 // also needs to keep track of how many conflicts each local has, so it's
649 // simpler to keep it this way for now.
650 let mut storage_conflicts = BitMatrix::new(saved_locals.count(), saved_locals.count());
651 for (saved_local_a, local_a) in saved_locals.iter_enumerated() {
652 if ineligible_locals.contains(local_a) {
653 // Conflicts with everything.
654 storage_conflicts.insert_all_into_row(saved_local_a);
655 } else {
656 // Keep overlap information only for stored locals.
657 for (saved_local_b, local_b) in saved_locals.iter_enumerated() {
658 if local_conflicts.contains(local_a, local_b) {
659 storage_conflicts.insert(saved_local_a, saved_local_b);
660 }
661 }
662 }
663 }
664 storage_conflicts
665 }
666
667 struct StorageConflictVisitor<'mir, 'tcx, 's> {
668 body: &'mir Body<'tcx>,
669 saved_locals: &'s GeneratorSavedLocals,
670 // FIXME(tmandry): Consider using sparse bitsets here once we have good
671 // benchmarks for generators.
672 local_conflicts: BitMatrix<Local, Local>,
673 }
674
675 impl rustc_mir_dataflow::ResultsVisitor<'mir, 'tcx> for StorageConflictVisitor<'mir, 'tcx, '_> {
676 type FlowState = BitSet<Local>;
677
visit_statement_before_primary_effect( &mut self, state: &Self::FlowState, _statement: &'mir Statement<'tcx>, loc: Location, )678 fn visit_statement_before_primary_effect(
679 &mut self,
680 state: &Self::FlowState,
681 _statement: &'mir Statement<'tcx>,
682 loc: Location,
683 ) {
684 self.apply_state(state, loc);
685 }
686
visit_terminator_before_primary_effect( &mut self, state: &Self::FlowState, _terminator: &'mir Terminator<'tcx>, loc: Location, )687 fn visit_terminator_before_primary_effect(
688 &mut self,
689 state: &Self::FlowState,
690 _terminator: &'mir Terminator<'tcx>,
691 loc: Location,
692 ) {
693 self.apply_state(state, loc);
694 }
695 }
696
697 impl<'body, 'tcx, 's> StorageConflictVisitor<'body, 'tcx, 's> {
apply_state(&mut self, flow_state: &BitSet<Local>, loc: Location)698 fn apply_state(&mut self, flow_state: &BitSet<Local>, loc: Location) {
699 // Ignore unreachable blocks.
700 if self.body.basic_blocks()[loc.block].terminator().kind == TerminatorKind::Unreachable {
701 return;
702 }
703
704 let mut eligible_storage_live = flow_state.clone();
705 eligible_storage_live.intersect(&**self.saved_locals);
706
707 for local in eligible_storage_live.iter() {
708 self.local_conflicts.union_row_with(&eligible_storage_live, local);
709 }
710
711 if eligible_storage_live.count() > 1 {
712 trace!("at {:?}, eligible_storage_live={:?}", loc, eligible_storage_live);
713 }
714 }
715 }
716
717 /// Validates the typeck view of the generator against the actual set of types saved between
718 /// yield points.
sanitize_witness<'tcx>( tcx: TyCtxt<'tcx>, body: &Body<'tcx>, witness: Ty<'tcx>, upvars: Vec<Ty<'tcx>>, saved_locals: &GeneratorSavedLocals, )719 fn sanitize_witness<'tcx>(
720 tcx: TyCtxt<'tcx>,
721 body: &Body<'tcx>,
722 witness: Ty<'tcx>,
723 upvars: Vec<Ty<'tcx>>,
724 saved_locals: &GeneratorSavedLocals,
725 ) {
726 let did = body.source.def_id();
727 let allowed_upvars = tcx.erase_regions(upvars);
728 let allowed = match witness.kind() {
729 &ty::GeneratorWitness(s) => tcx.erase_late_bound_regions(s),
730 _ => {
731 tcx.sess.delay_span_bug(
732 body.span,
733 &format!("unexpected generator witness type {:?}", witness.kind()),
734 );
735 return;
736 }
737 };
738
739 let param_env = tcx.param_env(did);
740
741 for (local, decl) in body.local_decls.iter_enumerated() {
742 // Ignore locals which are internal or not saved between yields.
743 if !saved_locals.contains(local) || decl.internal {
744 continue;
745 }
746 let decl_ty = tcx.normalize_erasing_regions(param_env, decl.ty);
747
748 // Sanity check that typeck knows about the type of locals which are
749 // live across a suspension point
750 if !allowed.contains(&decl_ty) && !allowed_upvars.contains(&decl_ty) {
751 span_bug!(
752 body.span,
753 "Broken MIR: generator contains type {} in MIR, \
754 but typeck only knows about {} and {:?}",
755 decl_ty,
756 allowed,
757 allowed_upvars
758 );
759 }
760 }
761 }
762
compute_layout<'tcx>( liveness: LivenessInfo, body: &mut Body<'tcx>, ) -> ( FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>, GeneratorLayout<'tcx>, IndexVec<BasicBlock, Option<BitSet<Local>>>, )763 fn compute_layout<'tcx>(
764 liveness: LivenessInfo,
765 body: &mut Body<'tcx>,
766 ) -> (
767 FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
768 GeneratorLayout<'tcx>,
769 IndexVec<BasicBlock, Option<BitSet<Local>>>,
770 ) {
771 let LivenessInfo {
772 saved_locals,
773 live_locals_at_suspension_points,
774 source_info_at_suspension_points,
775 storage_conflicts,
776 storage_liveness,
777 } = liveness;
778
779 // Gather live local types and their indices.
780 let mut locals = IndexVec::<GeneratorSavedLocal, _>::new();
781 let mut tys = IndexVec::<GeneratorSavedLocal, _>::new();
782 for (saved_local, local) in saved_locals.iter_enumerated() {
783 locals.push(local);
784 tys.push(body.local_decls[local].ty);
785 debug!("generator saved local {:?} => {:?}", saved_local, local);
786 }
787
788 // Leave empty variants for the UNRESUMED, RETURNED, and POISONED states.
789 // In debuginfo, these will correspond to the beginning (UNRESUMED) or end
790 // (RETURNED, POISONED) of the function.
791 const RESERVED_VARIANTS: usize = 3;
792 let body_span = body.source_scopes[OUTERMOST_SOURCE_SCOPE].span;
793 let mut variant_source_info: IndexVec<VariantIdx, SourceInfo> = [
794 SourceInfo::outermost(body_span.shrink_to_lo()),
795 SourceInfo::outermost(body_span.shrink_to_hi()),
796 SourceInfo::outermost(body_span.shrink_to_hi()),
797 ]
798 .iter()
799 .copied()
800 .collect();
801
802 // Build the generator variant field list.
803 // Create a map from local indices to generator struct indices.
804 let mut variant_fields: IndexVec<VariantIdx, IndexVec<Field, GeneratorSavedLocal>> =
805 iter::repeat(IndexVec::new()).take(RESERVED_VARIANTS).collect();
806 let mut remap = FxHashMap::default();
807 for (suspension_point_idx, live_locals) in live_locals_at_suspension_points.iter().enumerate() {
808 let variant_index = VariantIdx::from(RESERVED_VARIANTS + suspension_point_idx);
809 let mut fields = IndexVec::new();
810 for (idx, saved_local) in live_locals.iter().enumerate() {
811 fields.push(saved_local);
812 // Note that if a field is included in multiple variants, we will
813 // just use the first one here. That's fine; fields do not move
814 // around inside generators, so it doesn't matter which variant
815 // index we access them by.
816 remap.entry(locals[saved_local]).or_insert((tys[saved_local], variant_index, idx));
817 }
818 variant_fields.push(fields);
819 variant_source_info.push(source_info_at_suspension_points[suspension_point_idx]);
820 }
821 debug!("generator variant_fields = {:?}", variant_fields);
822 debug!("generator storage_conflicts = {:#?}", storage_conflicts);
823
824 let layout =
825 GeneratorLayout { field_tys: tys, variant_fields, variant_source_info, storage_conflicts };
826
827 (remap, layout, storage_liveness)
828 }
829
830 /// Replaces the entry point of `body` with a block that switches on the generator discriminant and
831 /// dispatches to blocks according to `cases`.
832 ///
833 /// After this function, the former entry point of the function will be bb1.
insert_switch<'tcx>( body: &mut Body<'tcx>, cases: Vec<(usize, BasicBlock)>, transform: &TransformVisitor<'tcx>, default: TerminatorKind<'tcx>, )834 fn insert_switch<'tcx>(
835 body: &mut Body<'tcx>,
836 cases: Vec<(usize, BasicBlock)>,
837 transform: &TransformVisitor<'tcx>,
838 default: TerminatorKind<'tcx>,
839 ) {
840 let default_block = insert_term_block(body, default);
841 let (assign, discr) = transform.get_discr(body);
842 let switch_targets =
843 SwitchTargets::new(cases.iter().map(|(i, bb)| ((*i) as u128, *bb)), default_block);
844 let switch = TerminatorKind::SwitchInt {
845 discr: Operand::Move(discr),
846 switch_ty: transform.discr_ty,
847 targets: switch_targets,
848 };
849
850 let source_info = SourceInfo::outermost(body.span);
851 body.basic_blocks_mut().raw.insert(
852 0,
853 BasicBlockData {
854 statements: vec![assign],
855 terminator: Some(Terminator { source_info, kind: switch }),
856 is_cleanup: false,
857 },
858 );
859
860 let blocks = body.basic_blocks_mut().iter_mut();
861
862 for target in blocks.flat_map(|b| b.terminator_mut().successors_mut()) {
863 *target = BasicBlock::new(target.index() + 1);
864 }
865 }
866
elaborate_generator_drops<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>)867 fn elaborate_generator_drops<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
868 use crate::shim::DropShimElaborator;
869 use rustc_middle::mir::patch::MirPatch;
870 use rustc_mir_dataflow::elaborate_drops::{elaborate_drop, Unwind};
871
872 // Note that `elaborate_drops` only drops the upvars of a generator, and
873 // this is ok because `open_drop` can only be reached within that own
874 // generator's resume function.
875
876 let def_id = body.source.def_id();
877 let param_env = tcx.param_env(def_id);
878
879 let mut elaborator = DropShimElaborator { body, patch: MirPatch::new(body), tcx, param_env };
880
881 for (block, block_data) in body.basic_blocks().iter_enumerated() {
882 let (target, unwind, source_info) = match block_data.terminator() {
883 Terminator { source_info, kind: TerminatorKind::Drop { place, target, unwind } } => {
884 if let Some(local) = place.as_local() {
885 if local == SELF_ARG {
886 (target, unwind, source_info)
887 } else {
888 continue;
889 }
890 } else {
891 continue;
892 }
893 }
894 _ => continue,
895 };
896 let unwind = if block_data.is_cleanup {
897 Unwind::InCleanup
898 } else {
899 Unwind::To(unwind.unwrap_or_else(|| elaborator.patch.resume_block()))
900 };
901 elaborate_drop(
902 &mut elaborator,
903 *source_info,
904 Place::from(SELF_ARG),
905 (),
906 *target,
907 unwind,
908 block,
909 );
910 }
911 elaborator.patch.apply(body);
912 }
913
create_generator_drop_shim<'tcx>( tcx: TyCtxt<'tcx>, transform: &TransformVisitor<'tcx>, gen_ty: Ty<'tcx>, body: &mut Body<'tcx>, drop_clean: BasicBlock, ) -> Body<'tcx>914 fn create_generator_drop_shim<'tcx>(
915 tcx: TyCtxt<'tcx>,
916 transform: &TransformVisitor<'tcx>,
917 gen_ty: Ty<'tcx>,
918 body: &mut Body<'tcx>,
919 drop_clean: BasicBlock,
920 ) -> Body<'tcx> {
921 let mut body = body.clone();
922 body.arg_count = 1; // make sure the resume argument is not included here
923
924 let source_info = SourceInfo::outermost(body.span);
925
926 let mut cases = create_cases(&mut body, transform, Operation::Drop);
927
928 cases.insert(0, (UNRESUMED, drop_clean));
929
930 // The returned state and the poisoned state fall through to the default
931 // case which is just to return
932
933 insert_switch(&mut body, cases, &transform, TerminatorKind::Return);
934
935 for block in body.basic_blocks_mut() {
936 let kind = &mut block.terminator_mut().kind;
937 if let TerminatorKind::GeneratorDrop = *kind {
938 *kind = TerminatorKind::Return;
939 }
940 }
941
942 // Replace the return variable
943 body.local_decls[RETURN_PLACE] = LocalDecl::with_source_info(tcx.mk_unit(), source_info);
944
945 make_generator_state_argument_indirect(tcx, &mut body);
946
947 // Change the generator argument from &mut to *mut
948 body.local_decls[SELF_ARG] = LocalDecl::with_source_info(
949 tcx.mk_ptr(ty::TypeAndMut { ty: gen_ty, mutbl: hir::Mutability::Mut }),
950 source_info,
951 );
952 if tcx.sess.opts.debugging_opts.mir_emit_retag {
953 // Alias tracking must know we changed the type
954 body.basic_blocks_mut()[START_BLOCK].statements.insert(
955 0,
956 Statement {
957 source_info,
958 kind: StatementKind::Retag(RetagKind::Raw, Box::new(Place::from(SELF_ARG))),
959 },
960 )
961 }
962
963 // Make sure we remove dead blocks to remove
964 // unrelated code from the resume part of the function
965 simplify::remove_dead_blocks(tcx, &mut body);
966
967 dump_mir(tcx, None, "generator_drop", &0, &body, |_, _| Ok(()));
968
969 body
970 }
971
insert_term_block<'tcx>(body: &mut Body<'tcx>, kind: TerminatorKind<'tcx>) -> BasicBlock972 fn insert_term_block<'tcx>(body: &mut Body<'tcx>, kind: TerminatorKind<'tcx>) -> BasicBlock {
973 let source_info = SourceInfo::outermost(body.span);
974 body.basic_blocks_mut().push(BasicBlockData {
975 statements: Vec::new(),
976 terminator: Some(Terminator { source_info, kind }),
977 is_cleanup: false,
978 })
979 }
980
insert_panic_block<'tcx>( tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>, message: AssertMessage<'tcx>, ) -> BasicBlock981 fn insert_panic_block<'tcx>(
982 tcx: TyCtxt<'tcx>,
983 body: &mut Body<'tcx>,
984 message: AssertMessage<'tcx>,
985 ) -> BasicBlock {
986 let assert_block = BasicBlock::new(body.basic_blocks().len());
987 let term = TerminatorKind::Assert {
988 cond: Operand::Constant(Box::new(Constant {
989 span: body.span,
990 user_ty: None,
991 literal: ty::Const::from_bool(tcx, false).into(),
992 })),
993 expected: true,
994 msg: message,
995 target: assert_block,
996 cleanup: None,
997 };
998
999 let source_info = SourceInfo::outermost(body.span);
1000 body.basic_blocks_mut().push(BasicBlockData {
1001 statements: Vec::new(),
1002 terminator: Some(Terminator { source_info, kind: term }),
1003 is_cleanup: false,
1004 });
1005
1006 assert_block
1007 }
1008
can_return<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, param_env: ty::ParamEnv<'tcx>) -> bool1009 fn can_return<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, param_env: ty::ParamEnv<'tcx>) -> bool {
1010 // Returning from a function with an uninhabited return type is undefined behavior.
1011 if tcx.conservative_is_privately_uninhabited(param_env.and(body.return_ty())) {
1012 return false;
1013 }
1014
1015 // If there's a return terminator the function may return.
1016 for block in body.basic_blocks() {
1017 if let TerminatorKind::Return = block.terminator().kind {
1018 return true;
1019 }
1020 }
1021
1022 // Otherwise the function can't return.
1023 false
1024 }
1025
can_unwind<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) -> bool1026 fn can_unwind<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) -> bool {
1027 // Nothing can unwind when landing pads are off.
1028 if tcx.sess.panic_strategy() == PanicStrategy::Abort {
1029 return false;
1030 }
1031
1032 // Unwinds can only start at certain terminators.
1033 for block in body.basic_blocks() {
1034 match block.terminator().kind {
1035 // These never unwind.
1036 TerminatorKind::Goto { .. }
1037 | TerminatorKind::SwitchInt { .. }
1038 | TerminatorKind::Abort
1039 | TerminatorKind::Return
1040 | TerminatorKind::Unreachable
1041 | TerminatorKind::GeneratorDrop
1042 | TerminatorKind::FalseEdge { .. }
1043 | TerminatorKind::FalseUnwind { .. }
1044 | TerminatorKind::InlineAsm { .. } => {}
1045
1046 // Resume will *continue* unwinding, but if there's no other unwinding terminator it
1047 // will never be reached.
1048 TerminatorKind::Resume => {}
1049
1050 TerminatorKind::Yield { .. } => {
1051 unreachable!("`can_unwind` called before generator transform")
1052 }
1053
1054 // These may unwind.
1055 TerminatorKind::Drop { .. }
1056 | TerminatorKind::DropAndReplace { .. }
1057 | TerminatorKind::Call { .. }
1058 | TerminatorKind::Assert { .. } => return true,
1059 }
1060 }
1061
1062 // If we didn't find an unwinding terminator, the function cannot unwind.
1063 false
1064 }
1065
create_generator_resume_function<'tcx>( tcx: TyCtxt<'tcx>, transform: TransformVisitor<'tcx>, body: &mut Body<'tcx>, can_return: bool, )1066 fn create_generator_resume_function<'tcx>(
1067 tcx: TyCtxt<'tcx>,
1068 transform: TransformVisitor<'tcx>,
1069 body: &mut Body<'tcx>,
1070 can_return: bool,
1071 ) {
1072 let can_unwind = can_unwind(tcx, body);
1073
1074 // Poison the generator when it unwinds
1075 if can_unwind {
1076 let source_info = SourceInfo::outermost(body.span);
1077 let poison_block = body.basic_blocks_mut().push(BasicBlockData {
1078 statements: vec![transform.set_discr(VariantIdx::new(POISONED), source_info)],
1079 terminator: Some(Terminator { source_info, kind: TerminatorKind::Resume }),
1080 is_cleanup: true,
1081 });
1082
1083 for (idx, block) in body.basic_blocks_mut().iter_enumerated_mut() {
1084 let source_info = block.terminator().source_info;
1085
1086 if let TerminatorKind::Resume = block.terminator().kind {
1087 // An existing `Resume` terminator is redirected to jump to our dedicated
1088 // "poisoning block" above.
1089 if idx != poison_block {
1090 *block.terminator_mut() = Terminator {
1091 source_info,
1092 kind: TerminatorKind::Goto { target: poison_block },
1093 };
1094 }
1095 } else if !block.is_cleanup {
1096 // Any terminators that *can* unwind but don't have an unwind target set are also
1097 // pointed at our poisoning block (unless they're part of the cleanup path).
1098 if let Some(unwind @ None) = block.terminator_mut().unwind_mut() {
1099 *unwind = Some(poison_block);
1100 }
1101 }
1102 }
1103 }
1104
1105 let mut cases = create_cases(body, &transform, Operation::Resume);
1106
1107 use rustc_middle::mir::AssertKind::{ResumedAfterPanic, ResumedAfterReturn};
1108
1109 // Jump to the entry point on the unresumed
1110 cases.insert(0, (UNRESUMED, BasicBlock::new(0)));
1111
1112 // Panic when resumed on the returned or poisoned state
1113 let generator_kind = body.generator_kind().unwrap();
1114
1115 if can_unwind {
1116 cases.insert(
1117 1,
1118 (POISONED, insert_panic_block(tcx, body, ResumedAfterPanic(generator_kind))),
1119 );
1120 }
1121
1122 if can_return {
1123 cases.insert(
1124 1,
1125 (RETURNED, insert_panic_block(tcx, body, ResumedAfterReturn(generator_kind))),
1126 );
1127 }
1128
1129 insert_switch(body, cases, &transform, TerminatorKind::Unreachable);
1130
1131 make_generator_state_argument_indirect(tcx, body);
1132 make_generator_state_argument_pinned(tcx, body);
1133
1134 // Make sure we remove dead blocks to remove
1135 // unrelated code from the drop part of the function
1136 simplify::remove_dead_blocks(tcx, body);
1137
1138 dump_mir(tcx, None, "generator_resume", &0, body, |_, _| Ok(()));
1139 }
1140
insert_clean_drop(body: &mut Body<'_>) -> BasicBlock1141 fn insert_clean_drop(body: &mut Body<'_>) -> BasicBlock {
1142 let return_block = insert_term_block(body, TerminatorKind::Return);
1143
1144 let term =
1145 TerminatorKind::Drop { place: Place::from(SELF_ARG), target: return_block, unwind: None };
1146 let source_info = SourceInfo::outermost(body.span);
1147
1148 // Create a block to destroy an unresumed generators. This can only destroy upvars.
1149 body.basic_blocks_mut().push(BasicBlockData {
1150 statements: Vec::new(),
1151 terminator: Some(Terminator { source_info, kind: term }),
1152 is_cleanup: false,
1153 })
1154 }
1155
1156 /// An operation that can be performed on a generator.
1157 #[derive(PartialEq, Copy, Clone)]
1158 enum Operation {
1159 Resume,
1160 Drop,
1161 }
1162
1163 impl Operation {
target_block(self, point: &SuspensionPoint<'_>) -> Option<BasicBlock>1164 fn target_block(self, point: &SuspensionPoint<'_>) -> Option<BasicBlock> {
1165 match self {
1166 Operation::Resume => Some(point.resume),
1167 Operation::Drop => point.drop,
1168 }
1169 }
1170 }
1171
create_cases<'tcx>( body: &mut Body<'tcx>, transform: &TransformVisitor<'tcx>, operation: Operation, ) -> Vec<(usize, BasicBlock)>1172 fn create_cases<'tcx>(
1173 body: &mut Body<'tcx>,
1174 transform: &TransformVisitor<'tcx>,
1175 operation: Operation,
1176 ) -> Vec<(usize, BasicBlock)> {
1177 let source_info = SourceInfo::outermost(body.span);
1178
1179 transform
1180 .suspension_points
1181 .iter()
1182 .filter_map(|point| {
1183 // Find the target for this suspension point, if applicable
1184 operation.target_block(point).map(|target| {
1185 let mut statements = Vec::new();
1186
1187 // Create StorageLive instructions for locals with live storage
1188 for i in 0..(body.local_decls.len()) {
1189 if i == 2 {
1190 // The resume argument is live on function entry. Don't insert a
1191 // `StorageLive`, or the following `Assign` will read from uninitialized
1192 // memory.
1193 continue;
1194 }
1195
1196 let l = Local::new(i);
1197 let needs_storage_live = point.storage_liveness.contains(l)
1198 && !transform.remap.contains_key(&l)
1199 && !transform.always_live_locals.contains(l);
1200 if needs_storage_live {
1201 statements
1202 .push(Statement { source_info, kind: StatementKind::StorageLive(l) });
1203 }
1204 }
1205
1206 if operation == Operation::Resume {
1207 // Move the resume argument to the destination place of the `Yield` terminator
1208 let resume_arg = Local::new(2); // 0 = return, 1 = self
1209 statements.push(Statement {
1210 source_info,
1211 kind: StatementKind::Assign(Box::new((
1212 point.resume_arg,
1213 Rvalue::Use(Operand::Move(resume_arg.into())),
1214 ))),
1215 });
1216 }
1217
1218 // Then jump to the real target
1219 let block = body.basic_blocks_mut().push(BasicBlockData {
1220 statements,
1221 terminator: Some(Terminator {
1222 source_info,
1223 kind: TerminatorKind::Goto { target },
1224 }),
1225 is_cleanup: false,
1226 });
1227
1228 (point.state, block)
1229 })
1230 })
1231 .collect()
1232 }
1233
1234 impl<'tcx> MirPass<'tcx> for StateTransform {
run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>)1235 fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
1236 let yield_ty = if let Some(yield_ty) = body.yield_ty() {
1237 yield_ty
1238 } else {
1239 // This only applies to generators
1240 return;
1241 };
1242
1243 assert!(body.generator_drop().is_none());
1244
1245 // The first argument is the generator type passed by value
1246 let gen_ty = body.local_decls.raw[1].ty;
1247
1248 // Get the interior types and substs which typeck computed
1249 let (upvars, interior, discr_ty, movable) = match *gen_ty.kind() {
1250 ty::Generator(_, substs, movability) => {
1251 let substs = substs.as_generator();
1252 (
1253 substs.upvar_tys().collect(),
1254 substs.witness(),
1255 substs.discr_ty(tcx),
1256 movability == hir::Movability::Movable,
1257 )
1258 }
1259 _ => {
1260 tcx.sess
1261 .delay_span_bug(body.span, &format!("unexpected generator type {}", gen_ty));
1262 return;
1263 }
1264 };
1265
1266 // Compute GeneratorState<yield_ty, return_ty>
1267 let state_did = tcx.require_lang_item(LangItem::GeneratorState, None);
1268 let state_adt_ref = tcx.adt_def(state_did);
1269 let state_substs = tcx.intern_substs(&[yield_ty.into(), body.return_ty().into()]);
1270 let ret_ty = tcx.mk_adt(state_adt_ref, state_substs);
1271
1272 // We rename RETURN_PLACE which has type mir.return_ty to new_ret_local
1273 // RETURN_PLACE then is a fresh unused local with type ret_ty.
1274 let new_ret_local = replace_local(RETURN_PLACE, ret_ty, body, tcx);
1275
1276 // We also replace the resume argument and insert an `Assign`.
1277 // This is needed because the resume argument `_2` might be live across a `yield`, in which
1278 // case there is no `Assign` to it that the transform can turn into a store to the generator
1279 // state. After the yield the slot in the generator state would then be uninitialized.
1280 let resume_local = Local::new(2);
1281 let new_resume_local =
1282 replace_local(resume_local, body.local_decls[resume_local].ty, body, tcx);
1283
1284 // When first entering the generator, move the resume argument into its new local.
1285 let source_info = SourceInfo::outermost(body.span);
1286 let stmts = &mut body.basic_blocks_mut()[BasicBlock::new(0)].statements;
1287 stmts.insert(
1288 0,
1289 Statement {
1290 source_info,
1291 kind: StatementKind::Assign(Box::new((
1292 new_resume_local.into(),
1293 Rvalue::Use(Operand::Move(resume_local.into())),
1294 ))),
1295 },
1296 );
1297
1298 let always_live_locals = storage::AlwaysLiveLocals::new(&body);
1299
1300 let liveness_info =
1301 locals_live_across_suspend_points(tcx, body, &always_live_locals, movable);
1302
1303 sanitize_witness(tcx, body, interior, upvars, &liveness_info.saved_locals);
1304
1305 if tcx.sess.opts.debugging_opts.validate_mir {
1306 let mut vis = EnsureGeneratorFieldAssignmentsNeverAlias {
1307 assigned_local: None,
1308 saved_locals: &liveness_info.saved_locals,
1309 storage_conflicts: &liveness_info.storage_conflicts,
1310 };
1311
1312 vis.visit_body(body);
1313 }
1314
1315 // Extract locals which are live across suspension point into `layout`
1316 // `remap` gives a mapping from local indices onto generator struct indices
1317 // `storage_liveness` tells us which locals have live storage at suspension points
1318 let (remap, layout, storage_liveness) = compute_layout(liveness_info, body);
1319
1320 let can_return = can_return(tcx, body, tcx.param_env(body.source.def_id()));
1321
1322 // Run the transformation which converts Places from Local to generator struct
1323 // accesses for locals in `remap`.
1324 // It also rewrites `return x` and `yield y` as writing a new generator state and returning
1325 // GeneratorState::Complete(x) and GeneratorState::Yielded(y) respectively.
1326 let mut transform = TransformVisitor {
1327 tcx,
1328 state_adt_ref,
1329 state_substs,
1330 remap,
1331 storage_liveness,
1332 always_live_locals,
1333 suspension_points: Vec::new(),
1334 new_ret_local,
1335 discr_ty,
1336 };
1337 transform.visit_body(body);
1338
1339 // Update our MIR struct to reflect the changes we've made
1340 body.arg_count = 2; // self, resume arg
1341 body.spread_arg = None;
1342
1343 body.generator.as_mut().unwrap().yield_ty = None;
1344 body.generator.as_mut().unwrap().generator_layout = Some(layout);
1345
1346 // Insert `drop(generator_struct)` which is used to drop upvars for generators in
1347 // the unresumed state.
1348 // This is expanded to a drop ladder in `elaborate_generator_drops`.
1349 let drop_clean = insert_clean_drop(body);
1350
1351 dump_mir(tcx, None, "generator_pre-elab", &0, body, |_, _| Ok(()));
1352
1353 // Expand `drop(generator_struct)` to a drop ladder which destroys upvars.
1354 // If any upvars are moved out of, drop elaboration will handle upvar destruction.
1355 // However we need to also elaborate the code generated by `insert_clean_drop`.
1356 elaborate_generator_drops(tcx, body);
1357
1358 dump_mir(tcx, None, "generator_post-transform", &0, body, |_, _| Ok(()));
1359
1360 // Create a copy of our MIR and use it to create the drop shim for the generator
1361 let drop_shim = create_generator_drop_shim(tcx, &transform, gen_ty, body, drop_clean);
1362
1363 body.generator.as_mut().unwrap().generator_drop = Some(drop_shim);
1364
1365 // Create the Generator::resume function
1366 create_generator_resume_function(tcx, transform, body, can_return);
1367 }
1368 }
1369
1370 /// Looks for any assignments between locals (e.g., `_4 = _5`) that will both be converted to fields
1371 /// in the generator state machine but whose storage is not marked as conflicting
1372 ///
1373 /// Validation needs to happen immediately *before* `TransformVisitor` is invoked, not after.
1374 ///
1375 /// This condition would arise when the assignment is the last use of `_5` but the initial
1376 /// definition of `_4` if we weren't extra careful to mark all locals used inside a statement as
1377 /// conflicting. Non-conflicting generator saved locals may be stored at the same location within
1378 /// the generator state machine, which would result in ill-formed MIR: the left-hand and right-hand
1379 /// sides of an assignment may not alias. This caused a miscompilation in [#73137].
1380 ///
1381 /// [#73137]: https://github.com/rust-lang/rust/issues/73137
1382 struct EnsureGeneratorFieldAssignmentsNeverAlias<'a> {
1383 saved_locals: &'a GeneratorSavedLocals,
1384 storage_conflicts: &'a BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal>,
1385 assigned_local: Option<GeneratorSavedLocal>,
1386 }
1387
1388 impl EnsureGeneratorFieldAssignmentsNeverAlias<'_> {
saved_local_for_direct_place(&self, place: Place<'_>) -> Option<GeneratorSavedLocal>1389 fn saved_local_for_direct_place(&self, place: Place<'_>) -> Option<GeneratorSavedLocal> {
1390 if place.is_indirect() {
1391 return None;
1392 }
1393
1394 self.saved_locals.get(place.local)
1395 }
1396
check_assigned_place(&mut self, place: Place<'tcx>, f: impl FnOnce(&mut Self))1397 fn check_assigned_place(&mut self, place: Place<'tcx>, f: impl FnOnce(&mut Self)) {
1398 if let Some(assigned_local) = self.saved_local_for_direct_place(place) {
1399 assert!(self.assigned_local.is_none(), "`check_assigned_place` must not recurse");
1400
1401 self.assigned_local = Some(assigned_local);
1402 f(self);
1403 self.assigned_local = None;
1404 }
1405 }
1406 }
1407
1408 impl Visitor<'tcx> for EnsureGeneratorFieldAssignmentsNeverAlias<'_> {
visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, location: Location)1409 fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, location: Location) {
1410 let lhs = match self.assigned_local {
1411 Some(l) => l,
1412 None => {
1413 // This visitor only invokes `visit_place` for the right-hand side of an assignment
1414 // and only after setting `self.assigned_local`. However, the default impl of
1415 // `Visitor::super_body` may call `visit_place` with a `NonUseContext` for places
1416 // with debuginfo. Ignore them here.
1417 assert!(!context.is_use());
1418 return;
1419 }
1420 };
1421
1422 let rhs = match self.saved_local_for_direct_place(*place) {
1423 Some(l) => l,
1424 None => return,
1425 };
1426
1427 if !self.storage_conflicts.contains(lhs, rhs) {
1428 bug!(
1429 "Assignment between generator saved locals whose storage is not \
1430 marked as conflicting: {:?}: {:?} = {:?}",
1431 location,
1432 lhs,
1433 rhs,
1434 );
1435 }
1436 }
1437
visit_statement(&mut self, statement: &Statement<'tcx>, location: Location)1438 fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
1439 match &statement.kind {
1440 StatementKind::Assign(box (lhs, rhs)) => {
1441 self.check_assigned_place(*lhs, |this| this.visit_rvalue(rhs, location));
1442 }
1443
1444 // FIXME: Does `llvm_asm!` have any aliasing requirements?
1445 StatementKind::LlvmInlineAsm(_) => {}
1446
1447 StatementKind::FakeRead(..)
1448 | StatementKind::SetDiscriminant { .. }
1449 | StatementKind::StorageLive(_)
1450 | StatementKind::StorageDead(_)
1451 | StatementKind::Retag(..)
1452 | StatementKind::AscribeUserType(..)
1453 | StatementKind::Coverage(..)
1454 | StatementKind::CopyNonOverlapping(..)
1455 | StatementKind::Nop => {}
1456 }
1457 }
1458
visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location)1459 fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
1460 // Checking for aliasing in terminators is probably overkill, but until we have actual
1461 // semantics, we should be conservative here.
1462 match &terminator.kind {
1463 TerminatorKind::Call {
1464 func,
1465 args,
1466 destination: Some((dest, _)),
1467 cleanup: _,
1468 from_hir_call: _,
1469 fn_span: _,
1470 } => {
1471 self.check_assigned_place(*dest, |this| {
1472 this.visit_operand(func, location);
1473 for arg in args {
1474 this.visit_operand(arg, location);
1475 }
1476 });
1477 }
1478
1479 TerminatorKind::Yield { value, resume: _, resume_arg, drop: _ } => {
1480 self.check_assigned_place(*resume_arg, |this| this.visit_operand(value, location));
1481 }
1482
1483 // FIXME: Does `asm!` have any aliasing requirements?
1484 TerminatorKind::InlineAsm { .. } => {}
1485
1486 TerminatorKind::Call { .. }
1487 | TerminatorKind::Goto { .. }
1488 | TerminatorKind::SwitchInt { .. }
1489 | TerminatorKind::Resume
1490 | TerminatorKind::Abort
1491 | TerminatorKind::Return
1492 | TerminatorKind::Unreachable
1493 | TerminatorKind::Drop { .. }
1494 | TerminatorKind::DropAndReplace { .. }
1495 | TerminatorKind::Assert { .. }
1496 | TerminatorKind::GeneratorDrop
1497 | TerminatorKind::FalseEdge { .. }
1498 | TerminatorKind::FalseUnwind { .. } => {}
1499 }
1500 }
1501 }
1502