1mod by_move_body;
54mod drop;
55use std::{iter, ops};
56
57pub(super) use by_move_body::coroutine_by_move_body_def_id;
58use drop::{
59 cleanup_async_drops, create_coroutine_drop_shim, create_coroutine_drop_shim_async,
60 create_coroutine_drop_shim_proxy_async, elaborate_coroutine_drops, expand_async_drops,
61 has_expandable_async_drops, insert_clean_drop,
62};
63use rustc_abi::{FieldIdx, VariantIdx};
64use rustc_data_structures::fx::FxHashSet;
65use rustc_errors::pluralize;
66use rustc_hir as hir;
67use rustc_hir::lang_items::LangItem;
68use rustc_hir::{CoroutineDesugaring, CoroutineKind};
69use rustc_index::bit_set::{BitMatrix, DenseBitSet, GrowableBitSet};
70use rustc_index::{Idx, IndexVec};
71use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
72use rustc_middle::mir::*;
73use rustc_middle::ty::util::Discr;
74use rustc_middle::ty::{
75 self, CoroutineArgs, CoroutineArgsExt, GenericArgsRef, InstanceKind, Ty, TyCtxt, TypingMode,
76};
77use rustc_middle::{bug, span_bug};
78use rustc_mir_dataflow::impls::{
79 MaybeBorrowedLocals, MaybeLiveLocals, MaybeRequiresStorage, MaybeStorageLive,
80 always_storage_live_locals,
81};
82use rustc_mir_dataflow::{
83 Analysis, Results, ResultsCursor, ResultsVisitor, visit_reachable_results,
84};
85use rustc_span::def_id::{DefId, LocalDefId};
86use rustc_span::source_map::dummy_spanned;
87use rustc_span::symbol::sym;
88use rustc_span::{DUMMY_SP, Span};
89use rustc_target::spec::PanicStrategy;
90use rustc_trait_selection::error_reporting::InferCtxtErrorExt;
91use rustc_trait_selection::infer::TyCtxtInferExt as _;
92use rustc_trait_selection::traits::{ObligationCause, ObligationCauseCode, ObligationCtxt};
93use tracing::{debug, instrument, trace};
94
95use crate::deref_separator::deref_finder;
96use crate::{abort_unwinding_calls, errors, pass_manager as pm, simplify};
97
98pub(super) struct StateTransform;
99
100struct RenameLocalVisitor<'tcx> {
101 from: Local,
102 to: Local,
103 tcx: TyCtxt<'tcx>,
104}
105
106impl<'tcx> MutVisitor<'tcx> for RenameLocalVisitor<'tcx> {
107 fn tcx(&self) -> TyCtxt<'tcx> {
108 self.tcx
109 }
110
111 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
112 if *local == self.from {
113 *local = self.to;
114 }
115 }
116
117 fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) {
118 match terminator.kind {
119 TerminatorKind::Return => {
120 }
123 _ => self.super_terminator(terminator, location),
124 }
125 }
126}
127
128struct SelfArgVisitor<'tcx> {
129 tcx: TyCtxt<'tcx>,
130 new_base: Place<'tcx>,
131}
132
133impl<'tcx> SelfArgVisitor<'tcx> {
134 fn new(tcx: TyCtxt<'tcx>, elem: ProjectionElem<Local, Ty<'tcx>>) -> Self {
135 Self { tcx, new_base: Place { local: SELF_ARG, projection: tcx.mk_place_elems(&[elem]) } }
136 }
137}
138
139impl<'tcx> MutVisitor<'tcx> for SelfArgVisitor<'tcx> {
140 fn tcx(&self) -> TyCtxt<'tcx> {
141 self.tcx
142 }
143
144 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
145 assert_ne!(*local, SELF_ARG);
146 }
147
148 fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) {
149 if place.local == SELF_ARG {
150 replace_base(place, self.new_base, self.tcx);
151 } else {
152 self.visit_local(&mut place.local, context, location);
153
154 for elem in place.projection.iter() {
155 if let PlaceElem::Index(local) = elem {
156 assert_ne!(local, SELF_ARG);
157 }
158 }
159 }
160 }
161}
162
163fn replace_base<'tcx>(place: &mut Place<'tcx>, new_base: Place<'tcx>, tcx: TyCtxt<'tcx>) {
164 place.local = new_base.local;
165
166 let mut new_projection = new_base.projection.to_vec();
167 new_projection.append(&mut place.projection.to_vec());
168
169 place.projection = tcx.mk_place_elems(&new_projection);
170}
171
172const SELF_ARG: Local = Local::from_u32(1);
173const CTX_ARG: Local = Local::from_u32(2);
174
175struct SuspensionPoint<'tcx> {
177 state: usize,
179 resume: BasicBlock,
181 resume_arg: Place<'tcx>,
183 drop: Option<BasicBlock>,
185 storage_liveness: GrowableBitSet<Local>,
187}
188
189struct TransformVisitor<'tcx> {
190 tcx: TyCtxt<'tcx>,
191 coroutine_kind: hir::CoroutineKind,
192
193 discr_ty: Ty<'tcx>,
195
196 remap: IndexVec<Local, Option<(Ty<'tcx>, VariantIdx, FieldIdx)>>,
198
199 storage_liveness: IndexVec<BasicBlock, Option<DenseBitSet<Local>>>,
201
202 suspension_points: Vec<SuspensionPoint<'tcx>>,
204
205 always_live_locals: DenseBitSet<Local>,
207
208 old_ret_local: Local,
210
211 old_yield_ty: Ty<'tcx>,
212
213 old_ret_ty: Ty<'tcx>,
214}
215
216impl<'tcx> TransformVisitor<'tcx> {
217 fn insert_none_ret_block(&self, body: &mut Body<'tcx>) -> BasicBlock {
218 let block = body.basic_blocks.next_index();
219 let source_info = SourceInfo::outermost(body.span);
220
221 let none_value = match self.coroutine_kind {
222 CoroutineKind::Desugared(CoroutineDesugaring::Async, _) => {
223 span_bug!(body.span, "`Future`s are not fused inherently")
224 }
225 CoroutineKind::Coroutine(_) => span_bug!(body.span, "`Coroutine`s cannot be fused"),
226 CoroutineKind::Desugared(CoroutineDesugaring::Gen, _) => {
228 let option_def_id = self.tcx.require_lang_item(LangItem::Option, body.span);
229 make_aggregate_adt(
230 option_def_id,
231 VariantIdx::ZERO,
232 self.tcx.mk_args(&[self.old_yield_ty.into()]),
233 IndexVec::new(),
234 )
235 }
236 CoroutineKind::Desugared(CoroutineDesugaring::AsyncGen, _) => {
238 let ty::Adt(_poll_adt, args) = *self.old_yield_ty.kind() else { bug!() };
239 let ty::Adt(_option_adt, args) = *args.type_at(0).kind() else { bug!() };
240 let yield_ty = args.type_at(0);
241 Rvalue::Use(Operand::Constant(Box::new(ConstOperand {
242 span: source_info.span,
243 const_: Const::Unevaluated(
244 UnevaluatedConst::new(
245 self.tcx.require_lang_item(LangItem::AsyncGenFinished, body.span),
246 self.tcx.mk_args(&[yield_ty.into()]),
247 ),
248 self.old_yield_ty,
249 ),
250 user_ty: None,
251 })))
252 }
253 };
254
255 let statements = vec![Statement::new(
256 source_info,
257 StatementKind::Assign(Box::new((Place::return_place(), none_value))),
258 )];
259
260 body.basic_blocks_mut().push(BasicBlockData::new_stmts(
261 statements,
262 Some(Terminator { source_info, kind: TerminatorKind::Return }),
263 false,
264 ));
265
266 block
267 }
268
269 fn make_state(
275 &self,
276 val: Operand<'tcx>,
277 source_info: SourceInfo,
278 is_return: bool,
279 statements: &mut Vec<Statement<'tcx>>,
280 ) {
281 const ZERO: VariantIdx = VariantIdx::ZERO;
282 const ONE: VariantIdx = VariantIdx::from_usize(1);
283 let rvalue = match self.coroutine_kind {
284 CoroutineKind::Desugared(CoroutineDesugaring::Async, _) => {
285 let poll_def_id = self.tcx.require_lang_item(LangItem::Poll, source_info.span);
286 let args = self.tcx.mk_args(&[self.old_ret_ty.into()]);
287 let (variant_idx, operands) = if is_return {
288 (ZERO, IndexVec::from_raw(vec![val])) } else {
290 (ONE, IndexVec::new()) };
292 make_aggregate_adt(poll_def_id, variant_idx, args, operands)
293 }
294 CoroutineKind::Desugared(CoroutineDesugaring::Gen, _) => {
295 let option_def_id = self.tcx.require_lang_item(LangItem::Option, source_info.span);
296 let args = self.tcx.mk_args(&[self.old_yield_ty.into()]);
297 let (variant_idx, operands) = if is_return {
298 (ZERO, IndexVec::new()) } else {
300 (ONE, IndexVec::from_raw(vec![val])) };
302 make_aggregate_adt(option_def_id, variant_idx, args, operands)
303 }
304 CoroutineKind::Desugared(CoroutineDesugaring::AsyncGen, _) => {
305 if is_return {
306 let ty::Adt(_poll_adt, args) = *self.old_yield_ty.kind() else { bug!() };
307 let ty::Adt(_option_adt, args) = *args.type_at(0).kind() else { bug!() };
308 let yield_ty = args.type_at(0);
309 Rvalue::Use(Operand::Constant(Box::new(ConstOperand {
310 span: source_info.span,
311 const_: Const::Unevaluated(
312 UnevaluatedConst::new(
313 self.tcx.require_lang_item(
314 LangItem::AsyncGenFinished,
315 source_info.span,
316 ),
317 self.tcx.mk_args(&[yield_ty.into()]),
318 ),
319 self.old_yield_ty,
320 ),
321 user_ty: None,
322 })))
323 } else {
324 Rvalue::Use(val)
325 }
326 }
327 CoroutineKind::Coroutine(_) => {
328 let coroutine_state_def_id =
329 self.tcx.require_lang_item(LangItem::CoroutineState, source_info.span);
330 let args = self.tcx.mk_args(&[self.old_yield_ty.into(), self.old_ret_ty.into()]);
331 let variant_idx = if is_return {
332 ONE } else {
334 ZERO };
336 make_aggregate_adt(
337 coroutine_state_def_id,
338 variant_idx,
339 args,
340 IndexVec::from_raw(vec![val]),
341 )
342 }
343 };
344
345 statements.push(Statement::new(
346 source_info,
347 StatementKind::Assign(Box::new((Place::return_place(), rvalue))),
348 ));
349 }
350
351 fn make_field(&self, variant_index: VariantIdx, idx: FieldIdx, ty: Ty<'tcx>) -> Place<'tcx> {
353 let self_place = Place::from(SELF_ARG);
354 let base = self.tcx.mk_place_downcast_unnamed(self_place, variant_index);
355 let mut projection = base.projection.to_vec();
356 projection.push(ProjectionElem::Field(idx, ty));
357
358 Place { local: base.local, projection: self.tcx.mk_place_elems(&projection) }
359 }
360
361 fn set_discr(&self, state_disc: VariantIdx, source_info: SourceInfo) -> Statement<'tcx> {
363 let self_place = Place::from(SELF_ARG);
364 Statement::new(
365 source_info,
366 StatementKind::SetDiscriminant {
367 place: Box::new(self_place),
368 variant_index: state_disc,
369 },
370 )
371 }
372
373 fn get_discr(&self, body: &mut Body<'tcx>) -> (Statement<'tcx>, Place<'tcx>) {
375 let temp_decl = LocalDecl::new(self.discr_ty, body.span);
376 let local_decls_len = body.local_decls.push(temp_decl);
377 let temp = Place::from(local_decls_len);
378
379 let self_place = Place::from(SELF_ARG);
380 let assign = Statement::new(
381 SourceInfo::outermost(body.span),
382 StatementKind::Assign(Box::new((temp, Rvalue::Discriminant(self_place)))),
383 );
384 (assign, temp)
385 }
386}
387
388impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> {
389 fn tcx(&self) -> TyCtxt<'tcx> {
390 self.tcx
391 }
392
393 fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) {
394 assert!(!self.remap.contains(*local));
395 }
396
397 fn visit_place(
398 &mut self,
399 place: &mut Place<'tcx>,
400 _context: PlaceContext,
401 _location: Location,
402 ) {
403 if let Some(&Some((ty, variant_index, idx))) = self.remap.get(place.local) {
405 replace_base(place, self.make_field(variant_index, idx, ty), self.tcx);
406 }
407 }
408
409 fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) {
410 for s in &mut data.statements {
412 if let StatementKind::StorageLive(l) | StatementKind::StorageDead(l) = s.kind
413 && self.remap.contains(l)
414 {
415 s.make_nop();
416 }
417 }
418
419 let ret_val = match data.terminator().kind {
420 TerminatorKind::Return => {
421 Some((true, None, Operand::Move(Place::from(self.old_ret_local)), None))
422 }
423 TerminatorKind::Yield { ref value, resume, resume_arg, drop } => {
424 Some((false, Some((resume, resume_arg)), value.clone(), drop))
425 }
426 _ => None,
427 };
428
429 if let Some((is_return, resume, v, drop)) = ret_val {
430 let source_info = data.terminator().source_info;
431 self.make_state(v, source_info, is_return, &mut data.statements);
433 let state = if let Some((resume, mut resume_arg)) = resume {
434 let state = CoroutineArgs::RESERVED_VARIANTS + self.suspension_points.len();
436
437 if let Some(&Some((ty, variant, idx))) = self.remap.get(resume_arg.local) {
440 replace_base(&mut resume_arg, self.make_field(variant, idx, ty), self.tcx);
441 }
442
443 let storage_liveness: GrowableBitSet<Local> =
444 self.storage_liveness[block].clone().unwrap().into();
445
446 for i in 0..self.always_live_locals.domain_size() {
447 let l = Local::new(i);
448 let needs_storage_dead = storage_liveness.contains(l)
449 && !self.remap.contains(l)
450 && !self.always_live_locals.contains(l);
451 if needs_storage_dead {
452 data.statements
453 .push(Statement::new(source_info, StatementKind::StorageDead(l)));
454 }
455 }
456
457 self.suspension_points.push(SuspensionPoint {
458 state,
459 resume,
460 resume_arg,
461 drop,
462 storage_liveness,
463 });
464
465 VariantIdx::new(state)
466 } else {
467 VariantIdx::new(CoroutineArgs::RETURNED) };
470 data.statements.push(self.set_discr(state, source_info));
471 data.terminator_mut().kind = TerminatorKind::Return;
472 }
473
474 self.super_basic_block_data(block, data);
475 }
476}
477
478fn make_aggregate_adt<'tcx>(
479 def_id: DefId,
480 variant_idx: VariantIdx,
481 args: GenericArgsRef<'tcx>,
482 operands: IndexVec<FieldIdx, Operand<'tcx>>,
483) -> Rvalue<'tcx> {
484 Rvalue::Aggregate(Box::new(AggregateKind::Adt(def_id, variant_idx, args, None, None)), operands)
485}
486
487fn make_coroutine_state_argument_indirect<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
488 let coroutine_ty = body.local_decls.raw[1].ty;
489
490 let ref_coroutine_ty = Ty::new_mut_ref(tcx, tcx.lifetimes.re_erased, coroutine_ty);
491
492 body.local_decls.raw[1].ty = ref_coroutine_ty;
494
495 SelfArgVisitor::new(tcx, ProjectionElem::Deref).visit_body(body);
497}
498
499fn make_coroutine_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
500 let ref_coroutine_ty = body.local_decls.raw[1].ty;
501
502 let pin_did = tcx.require_lang_item(LangItem::Pin, body.span);
503 let pin_adt_ref = tcx.adt_def(pin_did);
504 let args = tcx.mk_args(&[ref_coroutine_ty.into()]);
505 let pin_ref_coroutine_ty = Ty::new_adt(tcx, pin_adt_ref, args);
506
507 body.local_decls.raw[1].ty = pin_ref_coroutine_ty;
509
510 SelfArgVisitor::new(tcx, ProjectionElem::Field(FieldIdx::ZERO, ref_coroutine_ty))
512 .visit_body(body);
513}
514
515fn replace_local<'tcx>(
522 local: Local,
523 ty: Ty<'tcx>,
524 body: &mut Body<'tcx>,
525 tcx: TyCtxt<'tcx>,
526) -> Local {
527 let new_decl = LocalDecl::new(ty, body.span);
528 let new_local = body.local_decls.push(new_decl);
529 body.local_decls.swap(local, new_local);
530
531 RenameLocalVisitor { from: local, to: new_local, tcx }.visit_body(body);
532
533 new_local
534}
535
536fn transform_async_context<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> Ty<'tcx> {
558 let context_mut_ref = Ty::new_task_context(tcx);
559
560 replace_resume_ty_local(tcx, body, CTX_ARG, context_mut_ref);
562
563 let get_context_def_id = tcx.require_lang_item(LangItem::GetContext, body.span);
564
565 for bb in body.basic_blocks.indices() {
566 let bb_data = &body[bb];
567 if bb_data.is_cleanup {
568 continue;
569 }
570
571 match &bb_data.terminator().kind {
572 TerminatorKind::Call { func, .. } => {
573 let func_ty = func.ty(body, tcx);
574 if let ty::FnDef(def_id, _) = *func_ty.kind()
575 && def_id == get_context_def_id
576 {
577 let local = eliminate_get_context_call(&mut body[bb]);
578 replace_resume_ty_local(tcx, body, local, context_mut_ref);
579 }
580 }
581 TerminatorKind::Yield { resume_arg, .. } => {
582 replace_resume_ty_local(tcx, body, resume_arg.local, context_mut_ref);
583 }
584 _ => {}
585 }
586 }
587 context_mut_ref
588}
589
590fn eliminate_get_context_call<'tcx>(bb_data: &mut BasicBlockData<'tcx>) -> Local {
591 let terminator = bb_data.terminator.take().unwrap();
592 let TerminatorKind::Call { args, destination, target, .. } = terminator.kind else {
593 bug!();
594 };
595 let [arg] = *Box::try_from(args).unwrap();
596 let local = arg.node.place().unwrap().local;
597
598 let arg = Rvalue::Use(arg.node);
599 let assign =
600 Statement::new(terminator.source_info, StatementKind::Assign(Box::new((destination, arg))));
601 bb_data.statements.push(assign);
602 bb_data.terminator = Some(Terminator {
603 source_info: terminator.source_info,
604 kind: TerminatorKind::Goto { target: target.unwrap() },
605 });
606 local
607}
608
609#[cfg_attr(not(debug_assertions), allow(unused))]
610fn replace_resume_ty_local<'tcx>(
611 tcx: TyCtxt<'tcx>,
612 body: &mut Body<'tcx>,
613 local: Local,
614 context_mut_ref: Ty<'tcx>,
615) {
616 let local_ty = std::mem::replace(&mut body.local_decls[local].ty, context_mut_ref);
617 #[cfg(debug_assertions)]
620 {
621 if let ty::Adt(resume_ty_adt, _) = local_ty.kind() {
622 let expected_adt = tcx.adt_def(tcx.require_lang_item(LangItem::ResumeTy, body.span));
623 assert_eq!(*resume_ty_adt, expected_adt);
624 } else {
625 panic!("expected `ResumeTy`, found `{:?}`", local_ty);
626 };
627 }
628}
629
630fn transform_gen_context<'tcx>(body: &mut Body<'tcx>) {
640 body.arg_count = 1;
644}
645
646struct LivenessInfo {
647 saved_locals: CoroutineSavedLocals,
649
650 live_locals_at_suspension_points: Vec<DenseBitSet<CoroutineSavedLocal>>,
652
653 source_info_at_suspension_points: Vec<SourceInfo>,
655
656 storage_conflicts: BitMatrix<CoroutineSavedLocal, CoroutineSavedLocal>,
660
661 storage_liveness: IndexVec<BasicBlock, Option<DenseBitSet<Local>>>,
664}
665
666fn locals_live_across_suspend_points<'tcx>(
675 tcx: TyCtxt<'tcx>,
676 body: &Body<'tcx>,
677 always_live_locals: &DenseBitSet<Local>,
678 movable: bool,
679) -> LivenessInfo {
680 let mut storage_live = MaybeStorageLive::new(std::borrow::Cow::Borrowed(always_live_locals))
683 .iterate_to_fixpoint(tcx, body, None)
684 .into_results_cursor(body);
685
686 let borrowed_locals = MaybeBorrowedLocals.iterate_to_fixpoint(tcx, body, Some("coroutine"));
688 let mut borrowed_locals_analysis1 = borrowed_locals.analysis;
689 let mut borrowed_locals_analysis2 = borrowed_locals_analysis1.clone(); let borrowed_locals_cursor1 = ResultsCursor::new_borrowing(
691 body,
692 &mut borrowed_locals_analysis1,
693 &borrowed_locals.results,
694 );
695 let mut borrowed_locals_cursor2 = ResultsCursor::new_borrowing(
696 body,
697 &mut borrowed_locals_analysis2,
698 &borrowed_locals.results,
699 );
700
701 let mut requires_storage =
703 MaybeRequiresStorage::new(borrowed_locals_cursor1).iterate_to_fixpoint(tcx, body, None);
704 let mut requires_storage_cursor = ResultsCursor::new_borrowing(
705 body,
706 &mut requires_storage.analysis,
707 &requires_storage.results,
708 );
709
710 let mut liveness =
712 MaybeLiveLocals.iterate_to_fixpoint(tcx, body, Some("coroutine")).into_results_cursor(body);
713
714 let mut storage_liveness_map = IndexVec::from_elem(None, &body.basic_blocks);
715 let mut live_locals_at_suspension_points = Vec::new();
716 let mut source_info_at_suspension_points = Vec::new();
717 let mut live_locals_at_any_suspension_point = DenseBitSet::new_empty(body.local_decls.len());
718
719 for (block, data) in body.basic_blocks.iter_enumerated() {
720 if let TerminatorKind::Yield { .. } = data.terminator().kind {
721 let loc = Location { block, statement_index: data.statements.len() };
722
723 liveness.seek_to_block_end(block);
724 let mut live_locals = liveness.get().clone();
725
726 if !movable {
727 borrowed_locals_cursor2.seek_before_primary_effect(loc);
738 live_locals.union(borrowed_locals_cursor2.get());
739 }
740
741 storage_live.seek_before_primary_effect(loc);
744 storage_liveness_map[block] = Some(storage_live.get().clone());
745
746 requires_storage_cursor.seek_before_primary_effect(loc);
750 live_locals.intersect(requires_storage_cursor.get());
751
752 live_locals.remove(SELF_ARG);
754
755 debug!("loc = {:?}, live_locals = {:?}", loc, live_locals);
756
757 live_locals_at_any_suspension_point.union(&live_locals);
760
761 live_locals_at_suspension_points.push(live_locals);
762 source_info_at_suspension_points.push(data.terminator().source_info);
763 }
764 }
765
766 debug!("live_locals_anywhere = {:?}", live_locals_at_any_suspension_point);
767 let saved_locals = CoroutineSavedLocals(live_locals_at_any_suspension_point);
768
769 let live_locals_at_suspension_points = live_locals_at_suspension_points
772 .iter()
773 .map(|live_here| saved_locals.renumber_bitset(live_here))
774 .collect();
775
776 let storage_conflicts = compute_storage_conflicts(
777 body,
778 &saved_locals,
779 always_live_locals.clone(),
780 &mut requires_storage.analysis,
781 &requires_storage.results,
782 );
783
784 LivenessInfo {
785 saved_locals,
786 live_locals_at_suspension_points,
787 source_info_at_suspension_points,
788 storage_conflicts,
789 storage_liveness: storage_liveness_map,
790 }
791}
792
793struct CoroutineSavedLocals(DenseBitSet<Local>);
799
800impl CoroutineSavedLocals {
801 fn iter_enumerated(&self) -> impl '_ + Iterator<Item = (CoroutineSavedLocal, Local)> {
804 self.iter().enumerate().map(|(i, l)| (CoroutineSavedLocal::from(i), l))
805 }
806
807 fn renumber_bitset(&self, input: &DenseBitSet<Local>) -> DenseBitSet<CoroutineSavedLocal> {
810 assert!(self.superset(input), "{:?} not a superset of {:?}", self.0, input);
811 let mut out = DenseBitSet::new_empty(self.count());
812 for (saved_local, local) in self.iter_enumerated() {
813 if input.contains(local) {
814 out.insert(saved_local);
815 }
816 }
817 out
818 }
819
820 fn get(&self, local: Local) -> Option<CoroutineSavedLocal> {
821 if !self.contains(local) {
822 return None;
823 }
824
825 let idx = self.iter().take_while(|&l| l < local).count();
826 Some(CoroutineSavedLocal::new(idx))
827 }
828}
829
830impl ops::Deref for CoroutineSavedLocals {
831 type Target = DenseBitSet<Local>;
832
833 fn deref(&self) -> &Self::Target {
834 &self.0
835 }
836}
837
838fn compute_storage_conflicts<'mir, 'tcx>(
843 body: &'mir Body<'tcx>,
844 saved_locals: &'mir CoroutineSavedLocals,
845 always_live_locals: DenseBitSet<Local>,
846 analysis: &mut MaybeRequiresStorage<'mir, 'tcx>,
847 results: &Results<DenseBitSet<Local>>,
848) -> BitMatrix<CoroutineSavedLocal, CoroutineSavedLocal> {
849 assert_eq!(body.local_decls.len(), saved_locals.domain_size());
850
851 debug!("compute_storage_conflicts({:?})", body.span);
852 debug!("always_live = {:?}", always_live_locals);
853
854 let mut ineligible_locals = always_live_locals;
857 ineligible_locals.intersect(&**saved_locals);
858
859 let mut visitor = StorageConflictVisitor {
861 body,
862 saved_locals,
863 local_conflicts: BitMatrix::from_row_n(&ineligible_locals, body.local_decls.len()),
864 eligible_storage_live: DenseBitSet::new_empty(body.local_decls.len()),
865 };
866
867 visit_reachable_results(body, analysis, results, &mut visitor);
868
869 let local_conflicts = visitor.local_conflicts;
870
871 let mut storage_conflicts = BitMatrix::new(saved_locals.count(), saved_locals.count());
879 for (saved_local_a, local_a) in saved_locals.iter_enumerated() {
880 if ineligible_locals.contains(local_a) {
881 storage_conflicts.insert_all_into_row(saved_local_a);
883 } else {
884 for (saved_local_b, local_b) in saved_locals.iter_enumerated() {
886 if local_conflicts.contains(local_a, local_b) {
887 storage_conflicts.insert(saved_local_a, saved_local_b);
888 }
889 }
890 }
891 }
892 storage_conflicts
893}
894
895struct StorageConflictVisitor<'a, 'tcx> {
896 body: &'a Body<'tcx>,
897 saved_locals: &'a CoroutineSavedLocals,
898 local_conflicts: BitMatrix<Local, Local>,
901 eligible_storage_live: DenseBitSet<Local>,
903}
904
905impl<'a, 'tcx> ResultsVisitor<'tcx, MaybeRequiresStorage<'a, 'tcx>>
906 for StorageConflictVisitor<'a, 'tcx>
907{
908 fn visit_after_early_statement_effect(
909 &mut self,
910 _analysis: &mut MaybeRequiresStorage<'a, 'tcx>,
911 state: &DenseBitSet<Local>,
912 _statement: &Statement<'tcx>,
913 loc: Location,
914 ) {
915 self.apply_state(state, loc);
916 }
917
918 fn visit_after_early_terminator_effect(
919 &mut self,
920 _analysis: &mut MaybeRequiresStorage<'a, 'tcx>,
921 state: &DenseBitSet<Local>,
922 _terminator: &Terminator<'tcx>,
923 loc: Location,
924 ) {
925 self.apply_state(state, loc);
926 }
927}
928
929impl StorageConflictVisitor<'_, '_> {
930 fn apply_state(&mut self, state: &DenseBitSet<Local>, loc: Location) {
931 if let TerminatorKind::Unreachable = self.body.basic_blocks[loc.block].terminator().kind {
933 return;
934 }
935
936 self.eligible_storage_live.clone_from(state);
937 self.eligible_storage_live.intersect(&**self.saved_locals);
938
939 for local in self.eligible_storage_live.iter() {
940 self.local_conflicts.union_row_with(&self.eligible_storage_live, local);
941 }
942
943 if self.eligible_storage_live.count() > 1 {
944 trace!("at {:?}, eligible_storage_live={:?}", loc, self.eligible_storage_live);
945 }
946 }
947}
948
949fn compute_layout<'tcx>(
950 liveness: LivenessInfo,
951 body: &Body<'tcx>,
952) -> (
953 IndexVec<Local, Option<(Ty<'tcx>, VariantIdx, FieldIdx)>>,
954 CoroutineLayout<'tcx>,
955 IndexVec<BasicBlock, Option<DenseBitSet<Local>>>,
956) {
957 let LivenessInfo {
958 saved_locals,
959 live_locals_at_suspension_points,
960 source_info_at_suspension_points,
961 storage_conflicts,
962 storage_liveness,
963 } = liveness;
964
965 let mut locals = IndexVec::<CoroutineSavedLocal, _>::new();
967 let mut tys = IndexVec::<CoroutineSavedLocal, _>::new();
968 for (saved_local, local) in saved_locals.iter_enumerated() {
969 debug!("coroutine saved local {:?} => {:?}", saved_local, local);
970
971 locals.push(local);
972 let decl = &body.local_decls[local];
973 debug!(?decl);
974
975 let ignore_for_traits = match decl.local_info {
980 ClearCrossCrate::Set(box LocalInfo::StaticRef { is_thread_local, .. }) => {
983 !is_thread_local
984 }
985 ClearCrossCrate::Set(box LocalInfo::FakeBorrow) => true,
988 _ => false,
989 };
990 let decl =
991 CoroutineSavedTy { ty: decl.ty, source_info: decl.source_info, ignore_for_traits };
992 debug!(?decl);
993
994 tys.push(decl);
995 }
996
997 let body_span = body.source_scopes[OUTERMOST_SOURCE_SCOPE].span;
1001 let mut variant_source_info: IndexVec<VariantIdx, SourceInfo> = [
1002 SourceInfo::outermost(body_span.shrink_to_lo()),
1003 SourceInfo::outermost(body_span.shrink_to_hi()),
1004 SourceInfo::outermost(body_span.shrink_to_hi()),
1005 ]
1006 .iter()
1007 .copied()
1008 .collect();
1009
1010 let mut variant_fields: IndexVec<VariantIdx, IndexVec<FieldIdx, CoroutineSavedLocal>> =
1013 iter::repeat(IndexVec::new()).take(CoroutineArgs::RESERVED_VARIANTS).collect();
1014 let mut remap = IndexVec::from_elem_n(None, saved_locals.domain_size());
1015 for (suspension_point_idx, live_locals) in live_locals_at_suspension_points.iter().enumerate() {
1016 let variant_index =
1017 VariantIdx::from(CoroutineArgs::RESERVED_VARIANTS + suspension_point_idx);
1018 let mut fields = IndexVec::new();
1019 for (idx, saved_local) in live_locals.iter().enumerate() {
1020 fields.push(saved_local);
1021 let idx = FieldIdx::from_usize(idx);
1026 remap[locals[saved_local]] = Some((tys[saved_local].ty, variant_index, idx));
1027 }
1028 variant_fields.push(fields);
1029 variant_source_info.push(source_info_at_suspension_points[suspension_point_idx]);
1030 }
1031 debug!("coroutine variant_fields = {:?}", variant_fields);
1032 debug!("coroutine storage_conflicts = {:#?}", storage_conflicts);
1033
1034 let mut field_names = IndexVec::from_elem(None, &tys);
1035 for var in &body.var_debug_info {
1036 let VarDebugInfoContents::Place(place) = &var.value else { continue };
1037 let Some(local) = place.as_local() else { continue };
1038 let Some(&Some((_, variant, field))) = remap.get(local) else {
1039 continue;
1040 };
1041
1042 let saved_local = variant_fields[variant][field];
1043 field_names.get_or_insert_with(saved_local, || var.name);
1044 }
1045
1046 let layout = CoroutineLayout {
1047 field_tys: tys,
1048 field_names,
1049 variant_fields,
1050 variant_source_info,
1051 storage_conflicts,
1052 };
1053 debug!(?layout);
1054
1055 (remap, layout, storage_liveness)
1056}
1057
1058fn insert_switch<'tcx>(
1063 body: &mut Body<'tcx>,
1064 cases: Vec<(usize, BasicBlock)>,
1065 transform: &TransformVisitor<'tcx>,
1066 default_block: BasicBlock,
1067) {
1068 let (assign, discr) = transform.get_discr(body);
1069 let switch_targets =
1070 SwitchTargets::new(cases.iter().map(|(i, bb)| ((*i) as u128, *bb)), default_block);
1071 let switch = TerminatorKind::SwitchInt { discr: Operand::Move(discr), targets: switch_targets };
1072
1073 let source_info = SourceInfo::outermost(body.span);
1074 body.basic_blocks_mut().raw.insert(
1075 0,
1076 BasicBlockData::new_stmts(
1077 vec![assign],
1078 Some(Terminator { source_info, kind: switch }),
1079 false,
1080 ),
1081 );
1082
1083 for b in body.basic_blocks_mut().iter_mut() {
1084 b.terminator_mut().successors_mut(|target| *target += 1);
1085 }
1086}
1087
1088fn insert_term_block<'tcx>(body: &mut Body<'tcx>, kind: TerminatorKind<'tcx>) -> BasicBlock {
1089 let source_info = SourceInfo::outermost(body.span);
1090 body.basic_blocks_mut().push(BasicBlockData::new(Some(Terminator { source_info, kind }), false))
1091}
1092
1093fn return_poll_ready_assign<'tcx>(tcx: TyCtxt<'tcx>, source_info: SourceInfo) -> Statement<'tcx> {
1094 let poll_def_id = tcx.require_lang_item(LangItem::Poll, source_info.span);
1096 let args = tcx.mk_args(&[tcx.types.unit.into()]);
1097 let val = Operand::Constant(Box::new(ConstOperand {
1098 span: source_info.span,
1099 user_ty: None,
1100 const_: Const::zero_sized(tcx.types.unit),
1101 }));
1102 let ready_val = Rvalue::Aggregate(
1103 Box::new(AggregateKind::Adt(poll_def_id, VariantIdx::from_usize(0), args, None, None)),
1104 IndexVec::from_raw(vec![val]),
1105 );
1106 Statement::new(source_info, StatementKind::Assign(Box::new((Place::return_place(), ready_val))))
1107}
1108
1109fn insert_poll_ready_block<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> BasicBlock {
1110 let source_info = SourceInfo::outermost(body.span);
1111 body.basic_blocks_mut().push(BasicBlockData::new_stmts(
1112 [return_poll_ready_assign(tcx, source_info)].to_vec(),
1113 Some(Terminator { source_info, kind: TerminatorKind::Return }),
1114 false,
1115 ))
1116}
1117
1118fn insert_panic_block<'tcx>(
1119 tcx: TyCtxt<'tcx>,
1120 body: &mut Body<'tcx>,
1121 message: AssertMessage<'tcx>,
1122) -> BasicBlock {
1123 let assert_block = body.basic_blocks.next_index();
1124 let kind = TerminatorKind::Assert {
1125 cond: Operand::Constant(Box::new(ConstOperand {
1126 span: body.span,
1127 user_ty: None,
1128 const_: Const::from_bool(tcx, false),
1129 })),
1130 expected: true,
1131 msg: Box::new(message),
1132 target: assert_block,
1133 unwind: UnwindAction::Continue,
1134 };
1135
1136 insert_term_block(body, kind)
1137}
1138
1139fn can_return<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, typing_env: ty::TypingEnv<'tcx>) -> bool {
1140 if body.return_ty().is_privately_uninhabited(tcx, typing_env) {
1142 return false;
1143 }
1144
1145 body.basic_blocks.iter().any(|block| matches!(block.terminator().kind, TerminatorKind::Return))
1147 }
1149
1150fn can_unwind<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) -> bool {
1151 if tcx.sess.panic_strategy() == PanicStrategy::Abort {
1153 return false;
1154 }
1155
1156 for block in body.basic_blocks.iter() {
1158 match block.terminator().kind {
1159 TerminatorKind::Goto { .. }
1161 | TerminatorKind::SwitchInt { .. }
1162 | TerminatorKind::UnwindTerminate(_)
1163 | TerminatorKind::Return
1164 | TerminatorKind::Unreachable
1165 | TerminatorKind::CoroutineDrop
1166 | TerminatorKind::FalseEdge { .. }
1167 | TerminatorKind::FalseUnwind { .. } => {}
1168
1169 TerminatorKind::UnwindResume => {}
1172
1173 TerminatorKind::Yield { .. } => {
1174 unreachable!("`can_unwind` called before coroutine transform")
1175 }
1176
1177 TerminatorKind::Drop { .. }
1179 | TerminatorKind::Call { .. }
1180 | TerminatorKind::InlineAsm { .. }
1181 | TerminatorKind::Assert { .. } => return true,
1182
1183 TerminatorKind::TailCall { .. } => {
1184 unreachable!("tail calls can't be present in generators")
1185 }
1186 }
1187 }
1188
1189 false
1191}
1192
1193fn generate_poison_block_and_redirect_unwinds_there<'tcx>(
1195 transform: &TransformVisitor<'tcx>,
1196 body: &mut Body<'tcx>,
1197) {
1198 let source_info = SourceInfo::outermost(body.span);
1199 let poison_block = body.basic_blocks_mut().push(BasicBlockData::new_stmts(
1200 vec![transform.set_discr(VariantIdx::new(CoroutineArgs::POISONED), source_info)],
1201 Some(Terminator { source_info, kind: TerminatorKind::UnwindResume }),
1202 true,
1203 ));
1204
1205 for (idx, block) in body.basic_blocks_mut().iter_enumerated_mut() {
1206 let source_info = block.terminator().source_info;
1207
1208 if let TerminatorKind::UnwindResume = block.terminator().kind {
1209 if idx != poison_block {
1212 *block.terminator_mut() =
1213 Terminator { source_info, kind: TerminatorKind::Goto { target: poison_block } };
1214 }
1215 } else if !block.is_cleanup
1216 && let Some(unwind @ UnwindAction::Continue) = block.terminator_mut().unwind_mut()
1219 {
1220 *unwind = UnwindAction::Cleanup(poison_block);
1221 }
1222 }
1223}
1224
1225fn create_coroutine_resume_function<'tcx>(
1226 tcx: TyCtxt<'tcx>,
1227 transform: TransformVisitor<'tcx>,
1228 body: &mut Body<'tcx>,
1229 can_return: bool,
1230 can_unwind: bool,
1231) {
1232 if can_unwind {
1234 generate_poison_block_and_redirect_unwinds_there(&transform, body);
1235 }
1236
1237 let mut cases = create_cases(body, &transform, Operation::Resume);
1238
1239 use rustc_middle::mir::AssertKind::{ResumedAfterPanic, ResumedAfterReturn};
1240
1241 cases.insert(0, (CoroutineArgs::UNRESUMED, START_BLOCK));
1243
1244 if can_unwind {
1246 cases.insert(
1247 1,
1248 (
1249 CoroutineArgs::POISONED,
1250 insert_panic_block(tcx, body, ResumedAfterPanic(transform.coroutine_kind)),
1251 ),
1252 );
1253 }
1254
1255 if can_return {
1256 let block = match transform.coroutine_kind {
1257 CoroutineKind::Desugared(CoroutineDesugaring::Async, _)
1258 | CoroutineKind::Coroutine(_) => {
1259 if tcx.is_async_drop_in_place_coroutine(body.source.def_id()) {
1262 insert_poll_ready_block(tcx, body)
1263 } else {
1264 insert_panic_block(tcx, body, ResumedAfterReturn(transform.coroutine_kind))
1265 }
1266 }
1267 CoroutineKind::Desugared(CoroutineDesugaring::AsyncGen, _)
1268 | CoroutineKind::Desugared(CoroutineDesugaring::Gen, _) => {
1269 transform.insert_none_ret_block(body)
1270 }
1271 };
1272 cases.insert(1, (CoroutineArgs::RETURNED, block));
1273 }
1274
1275 let default_block = insert_term_block(body, TerminatorKind::Unreachable);
1276 insert_switch(body, cases, &transform, default_block);
1277
1278 make_coroutine_state_argument_indirect(tcx, body);
1279
1280 match transform.coroutine_kind {
1281 CoroutineKind::Coroutine(_)
1282 | CoroutineKind::Desugared(CoroutineDesugaring::Async | CoroutineDesugaring::AsyncGen, _) =>
1283 {
1284 make_coroutine_state_argument_pinned(tcx, body);
1285 }
1286 CoroutineKind::Desugared(CoroutineDesugaring::Gen, _) => {}
1289 }
1290
1291 simplify::remove_dead_blocks(body);
1294
1295 pm::run_passes_no_validate(tcx, body, &[&abort_unwinding_calls::AbortUnwindingCalls], None);
1296
1297 dump_mir(tcx, false, "coroutine_resume", &0, body, |_, _| Ok(()));
1298}
1299
1300#[derive(PartialEq, Copy, Clone)]
1302enum Operation {
1303 Resume,
1304 Drop,
1305}
1306
1307impl Operation {
1308 fn target_block(self, point: &SuspensionPoint<'_>) -> Option<BasicBlock> {
1309 match self {
1310 Operation::Resume => Some(point.resume),
1311 Operation::Drop => point.drop,
1312 }
1313 }
1314}
1315
1316fn create_cases<'tcx>(
1317 body: &mut Body<'tcx>,
1318 transform: &TransformVisitor<'tcx>,
1319 operation: Operation,
1320) -> Vec<(usize, BasicBlock)> {
1321 let source_info = SourceInfo::outermost(body.span);
1322
1323 transform
1324 .suspension_points
1325 .iter()
1326 .filter_map(|point| {
1327 operation.target_block(point).map(|target| {
1329 let mut statements = Vec::new();
1330
1331 for l in body.local_decls.indices() {
1333 let needs_storage_live = point.storage_liveness.contains(l)
1334 && !transform.remap.contains(l)
1335 && !transform.always_live_locals.contains(l);
1336 if needs_storage_live {
1337 statements.push(Statement::new(source_info, StatementKind::StorageLive(l)));
1338 }
1339 }
1340
1341 if operation == Operation::Resume {
1342 let resume_arg = CTX_ARG;
1344 statements.push(Statement::new(
1345 source_info,
1346 StatementKind::Assign(Box::new((
1347 point.resume_arg,
1348 Rvalue::Use(Operand::Move(resume_arg.into())),
1349 ))),
1350 ));
1351 }
1352
1353 let block = body.basic_blocks_mut().push(BasicBlockData::new_stmts(
1355 statements,
1356 Some(Terminator { source_info, kind: TerminatorKind::Goto { target } }),
1357 false,
1358 ));
1359
1360 (point.state, block)
1361 })
1362 })
1363 .collect()
1364}
1365
1366#[instrument(level = "debug", skip(tcx), ret)]
1367pub(crate) fn mir_coroutine_witnesses<'tcx>(
1368 tcx: TyCtxt<'tcx>,
1369 def_id: LocalDefId,
1370) -> Option<CoroutineLayout<'tcx>> {
1371 let (body, _) = tcx.mir_promoted(def_id);
1372 let body = body.borrow();
1373 let body = &*body;
1374
1375 let coroutine_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
1377
1378 let movable = match *coroutine_ty.kind() {
1379 ty::Coroutine(def_id, _) => tcx.coroutine_movability(def_id) == hir::Movability::Movable,
1380 ty::Error(_) => return None,
1381 _ => span_bug!(body.span, "unexpected coroutine type {}", coroutine_ty),
1382 };
1383
1384 let always_live_locals = always_storage_live_locals(body);
1387 let liveness_info = locals_live_across_suspend_points(tcx, body, &always_live_locals, movable);
1388
1389 let (_, coroutine_layout, _) = compute_layout(liveness_info, body);
1393
1394 check_suspend_tys(tcx, &coroutine_layout, body);
1395 check_field_tys_sized(tcx, &coroutine_layout, def_id);
1396
1397 Some(coroutine_layout)
1398}
1399
1400fn check_field_tys_sized<'tcx>(
1401 tcx: TyCtxt<'tcx>,
1402 coroutine_layout: &CoroutineLayout<'tcx>,
1403 def_id: LocalDefId,
1404) {
1405 if !tcx.features().unsized_fn_params() {
1408 return;
1409 }
1410
1411 let infcx = tcx.infer_ctxt().ignoring_regions().build(TypingMode::non_body_analysis());
1416 let param_env = tcx.param_env(def_id);
1417
1418 let ocx = ObligationCtxt::new_with_diagnostics(&infcx);
1419 for field_ty in &coroutine_layout.field_tys {
1420 ocx.register_bound(
1421 ObligationCause::new(
1422 field_ty.source_info.span,
1423 def_id,
1424 ObligationCauseCode::SizedCoroutineInterior(def_id),
1425 ),
1426 param_env,
1427 field_ty.ty,
1428 tcx.require_lang_item(hir::LangItem::Sized, field_ty.source_info.span),
1429 );
1430 }
1431
1432 let errors = ocx.select_all_or_error();
1433 debug!(?errors);
1434 if !errors.is_empty() {
1435 infcx.err_ctxt().report_fulfillment_errors(errors);
1436 }
1437}
1438
1439impl<'tcx> crate::MirPass<'tcx> for StateTransform {
1440 fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
1441 let Some(old_yield_ty) = body.yield_ty() else {
1442 return;
1444 };
1445 let old_ret_ty = body.return_ty();
1446
1447 assert!(body.coroutine_drop().is_none() && body.coroutine_drop_async().is_none());
1448
1449 dump_mir(tcx, false, "coroutine_before", &0, body, |_, _| Ok(()));
1450
1451 let coroutine_ty = body.local_decls.raw[1].ty;
1453 let coroutine_kind = body.coroutine_kind().unwrap();
1454
1455 let ty::Coroutine(_, args) = coroutine_ty.kind() else {
1457 tcx.dcx().span_bug(body.span, format!("unexpected coroutine type {coroutine_ty}"));
1458 };
1459 let discr_ty = args.as_coroutine().discr_ty(tcx);
1460
1461 let new_ret_ty = match coroutine_kind {
1462 CoroutineKind::Desugared(CoroutineDesugaring::Async, _) => {
1463 let poll_did = tcx.require_lang_item(LangItem::Poll, body.span);
1465 let poll_adt_ref = tcx.adt_def(poll_did);
1466 let poll_args = tcx.mk_args(&[old_ret_ty.into()]);
1467 Ty::new_adt(tcx, poll_adt_ref, poll_args)
1468 }
1469 CoroutineKind::Desugared(CoroutineDesugaring::Gen, _) => {
1470 let option_did = tcx.require_lang_item(LangItem::Option, body.span);
1472 let option_adt_ref = tcx.adt_def(option_did);
1473 let option_args = tcx.mk_args(&[old_yield_ty.into()]);
1474 Ty::new_adt(tcx, option_adt_ref, option_args)
1475 }
1476 CoroutineKind::Desugared(CoroutineDesugaring::AsyncGen, _) => {
1477 old_yield_ty
1479 }
1480 CoroutineKind::Coroutine(_) => {
1481 let state_did = tcx.require_lang_item(LangItem::CoroutineState, body.span);
1483 let state_adt_ref = tcx.adt_def(state_did);
1484 let state_args = tcx.mk_args(&[old_yield_ty.into(), old_ret_ty.into()]);
1485 Ty::new_adt(tcx, state_adt_ref, state_args)
1486 }
1487 };
1488
1489 let old_ret_local = replace_local(RETURN_PLACE, new_ret_ty, body, tcx);
1492
1493 let has_async_drops = matches!(
1498 coroutine_kind,
1499 CoroutineKind::Desugared(CoroutineDesugaring::Async | CoroutineDesugaring::AsyncGen, _)
1500 ) && has_expandable_async_drops(tcx, body, coroutine_ty);
1501
1502 if matches!(
1504 coroutine_kind,
1505 CoroutineKind::Desugared(CoroutineDesugaring::Async | CoroutineDesugaring::AsyncGen, _)
1506 ) {
1507 let context_mut_ref = transform_async_context(tcx, body);
1508 expand_async_drops(tcx, body, context_mut_ref, coroutine_kind, coroutine_ty);
1509 dump_mir(tcx, false, "coroutine_async_drop_expand", &0, body, |_, _| Ok(()));
1510 } else {
1511 cleanup_async_drops(body);
1512 }
1513
1514 let resume_local = CTX_ARG;
1519 let resume_ty = body.local_decls[resume_local].ty;
1520 let old_resume_local = replace_local(resume_local, resume_ty, body, tcx);
1521
1522 let source_info = SourceInfo::outermost(body.span);
1525 let stmts = &mut body.basic_blocks_mut()[START_BLOCK].statements;
1526 stmts.insert(
1527 0,
1528 Statement::new(
1529 source_info,
1530 StatementKind::Assign(Box::new((
1531 old_resume_local.into(),
1532 Rvalue::Use(Operand::Move(resume_local.into())),
1533 ))),
1534 ),
1535 );
1536
1537 let always_live_locals = always_storage_live_locals(body);
1538
1539 let movable = coroutine_kind.movability() == hir::Movability::Movable;
1540 let liveness_info =
1541 locals_live_across_suspend_points(tcx, body, &always_live_locals, movable);
1542
1543 if tcx.sess.opts.unstable_opts.validate_mir {
1544 let mut vis = EnsureCoroutineFieldAssignmentsNeverAlias {
1545 assigned_local: None,
1546 saved_locals: &liveness_info.saved_locals,
1547 storage_conflicts: &liveness_info.storage_conflicts,
1548 };
1549
1550 vis.visit_body(body);
1551 }
1552
1553 let (remap, layout, storage_liveness) = compute_layout(liveness_info, body);
1557
1558 let can_return = can_return(tcx, body, body.typing_env(tcx));
1559
1560 let mut transform = TransformVisitor {
1566 tcx,
1567 coroutine_kind,
1568 remap,
1569 storage_liveness,
1570 always_live_locals,
1571 suspension_points: Vec::new(),
1572 old_ret_local,
1573 discr_ty,
1574 old_ret_ty,
1575 old_yield_ty,
1576 };
1577 transform.visit_body(body);
1578
1579 body.arg_count = 2; body.spread_arg = None;
1582
1583 if matches!(coroutine_kind, CoroutineKind::Desugared(CoroutineDesugaring::Gen, _)) {
1585 transform_gen_context(body);
1586 }
1587
1588 for var in &mut body.var_debug_info {
1592 var.argument_index = None;
1593 }
1594
1595 body.coroutine.as_mut().unwrap().yield_ty = None;
1596 body.coroutine.as_mut().unwrap().resume_ty = None;
1597 body.coroutine.as_mut().unwrap().coroutine_layout = Some(layout);
1598
1599 let drop_clean = insert_clean_drop(tcx, body, has_async_drops);
1607
1608 dump_mir(tcx, false, "coroutine_pre-elab", &0, body, |_, _| Ok(()));
1609
1610 elaborate_coroutine_drops(tcx, body);
1614
1615 dump_mir(tcx, false, "coroutine_post-transform", &0, body, |_, _| Ok(()));
1616
1617 let can_unwind = can_unwind(tcx, body);
1618
1619 if has_async_drops {
1621 let mut drop_shim =
1623 create_coroutine_drop_shim_async(tcx, &transform, body, drop_clean, can_unwind);
1624 deref_finder(tcx, &mut drop_shim);
1626 body.coroutine.as_mut().unwrap().coroutine_drop_async = Some(drop_shim);
1627 } else {
1628 let mut drop_shim =
1630 create_coroutine_drop_shim(tcx, &transform, coroutine_ty, body, drop_clean);
1631 deref_finder(tcx, &mut drop_shim);
1633 body.coroutine.as_mut().unwrap().coroutine_drop = Some(drop_shim);
1634
1635 let mut proxy_shim = create_coroutine_drop_shim_proxy_async(tcx, body);
1637 deref_finder(tcx, &mut proxy_shim);
1638 body.coroutine.as_mut().unwrap().coroutine_drop_proxy_async = Some(proxy_shim);
1639 }
1640
1641 create_coroutine_resume_function(tcx, transform, body, can_return, can_unwind);
1643
1644 deref_finder(tcx, body);
1646 }
1647
1648 fn is_required(&self) -> bool {
1649 true
1650 }
1651}
1652
1653struct EnsureCoroutineFieldAssignmentsNeverAlias<'a> {
1666 saved_locals: &'a CoroutineSavedLocals,
1667 storage_conflicts: &'a BitMatrix<CoroutineSavedLocal, CoroutineSavedLocal>,
1668 assigned_local: Option<CoroutineSavedLocal>,
1669}
1670
1671impl EnsureCoroutineFieldAssignmentsNeverAlias<'_> {
1672 fn saved_local_for_direct_place(&self, place: Place<'_>) -> Option<CoroutineSavedLocal> {
1673 if place.is_indirect() {
1674 return None;
1675 }
1676
1677 self.saved_locals.get(place.local)
1678 }
1679
1680 fn check_assigned_place(&mut self, place: Place<'_>, f: impl FnOnce(&mut Self)) {
1681 if let Some(assigned_local) = self.saved_local_for_direct_place(place) {
1682 assert!(self.assigned_local.is_none(), "`check_assigned_place` must not recurse");
1683
1684 self.assigned_local = Some(assigned_local);
1685 f(self);
1686 self.assigned_local = None;
1687 }
1688 }
1689}
1690
1691impl<'tcx> Visitor<'tcx> for EnsureCoroutineFieldAssignmentsNeverAlias<'_> {
1692 fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, location: Location) {
1693 let Some(lhs) = self.assigned_local else {
1694 assert!(!context.is_use());
1699 return;
1700 };
1701
1702 let Some(rhs) = self.saved_local_for_direct_place(*place) else { return };
1703
1704 if !self.storage_conflicts.contains(lhs, rhs) {
1705 bug!(
1706 "Assignment between coroutine saved locals whose storage is not \
1707 marked as conflicting: {:?}: {:?} = {:?}",
1708 location,
1709 lhs,
1710 rhs,
1711 );
1712 }
1713 }
1714
1715 fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
1716 match &statement.kind {
1717 StatementKind::Assign(box (lhs, rhs)) => {
1718 self.check_assigned_place(*lhs, |this| this.visit_rvalue(rhs, location));
1719 }
1720
1721 StatementKind::FakeRead(..)
1722 | StatementKind::SetDiscriminant { .. }
1723 | StatementKind::Deinit(..)
1724 | StatementKind::StorageLive(_)
1725 | StatementKind::StorageDead(_)
1726 | StatementKind::Retag(..)
1727 | StatementKind::AscribeUserType(..)
1728 | StatementKind::PlaceMention(..)
1729 | StatementKind::Coverage(..)
1730 | StatementKind::Intrinsic(..)
1731 | StatementKind::ConstEvalCounter
1732 | StatementKind::BackwardIncompatibleDropHint { .. }
1733 | StatementKind::Nop => {}
1734 }
1735 }
1736
1737 fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
1738 match &terminator.kind {
1741 TerminatorKind::Call {
1742 func,
1743 args,
1744 destination,
1745 target: Some(_),
1746 unwind: _,
1747 call_source: _,
1748 fn_span: _,
1749 } => {
1750 self.check_assigned_place(*destination, |this| {
1751 this.visit_operand(func, location);
1752 for arg in args {
1753 this.visit_operand(&arg.node, location);
1754 }
1755 });
1756 }
1757
1758 TerminatorKind::Yield { value, resume: _, resume_arg, drop: _ } => {
1759 self.check_assigned_place(*resume_arg, |this| this.visit_operand(value, location));
1760 }
1761
1762 TerminatorKind::InlineAsm { .. } => {}
1764
1765 TerminatorKind::Call { .. }
1766 | TerminatorKind::Goto { .. }
1767 | TerminatorKind::SwitchInt { .. }
1768 | TerminatorKind::UnwindResume
1769 | TerminatorKind::UnwindTerminate(_)
1770 | TerminatorKind::Return
1771 | TerminatorKind::TailCall { .. }
1772 | TerminatorKind::Unreachable
1773 | TerminatorKind::Drop { .. }
1774 | TerminatorKind::Assert { .. }
1775 | TerminatorKind::CoroutineDrop
1776 | TerminatorKind::FalseEdge { .. }
1777 | TerminatorKind::FalseUnwind { .. } => {}
1778 }
1779 }
1780}
1781
1782fn check_suspend_tys<'tcx>(tcx: TyCtxt<'tcx>, layout: &CoroutineLayout<'tcx>, body: &Body<'tcx>) {
1783 let mut linted_tys = FxHashSet::default();
1784
1785 for (variant, yield_source_info) in
1786 layout.variant_fields.iter().zip(&layout.variant_source_info)
1787 {
1788 debug!(?variant);
1789 for &local in variant {
1790 let decl = &layout.field_tys[local];
1791 debug!(?decl);
1792
1793 if !decl.ignore_for_traits && linted_tys.insert(decl.ty) {
1794 let Some(hir_id) = decl.source_info.scope.lint_root(&body.source_scopes) else {
1795 continue;
1796 };
1797
1798 check_must_not_suspend_ty(
1799 tcx,
1800 decl.ty,
1801 hir_id,
1802 SuspendCheckData {
1803 source_span: decl.source_info.span,
1804 yield_span: yield_source_info.span,
1805 plural_len: 1,
1806 ..Default::default()
1807 },
1808 );
1809 }
1810 }
1811 }
1812}
1813
1814#[derive(Default)]
1815struct SuspendCheckData<'a> {
1816 source_span: Span,
1817 yield_span: Span,
1818 descr_pre: &'a str,
1819 descr_post: &'a str,
1820 plural_len: usize,
1821}
1822
1823fn check_must_not_suspend_ty<'tcx>(
1830 tcx: TyCtxt<'tcx>,
1831 ty: Ty<'tcx>,
1832 hir_id: hir::HirId,
1833 data: SuspendCheckData<'_>,
1834) -> bool {
1835 if ty.is_unit() {
1836 return false;
1837 }
1838
1839 let plural_suffix = pluralize!(data.plural_len);
1840
1841 debug!("Checking must_not_suspend for {}", ty);
1842
1843 match *ty.kind() {
1844 ty::Adt(_, args) if ty.is_box() => {
1845 let boxed_ty = args.type_at(0);
1846 let allocator_ty = args.type_at(1);
1847 check_must_not_suspend_ty(
1848 tcx,
1849 boxed_ty,
1850 hir_id,
1851 SuspendCheckData { descr_pre: &format!("{}boxed ", data.descr_pre), ..data },
1852 ) || check_must_not_suspend_ty(
1853 tcx,
1854 allocator_ty,
1855 hir_id,
1856 SuspendCheckData { descr_pre: &format!("{}allocator ", data.descr_pre), ..data },
1857 )
1858 }
1859 ty::Adt(def, _) => check_must_not_suspend_def(tcx, def.did(), hir_id, data),
1860 ty::Alias(ty::Opaque, ty::AliasTy { def_id: def, .. }) => {
1862 let mut has_emitted = false;
1863 for &(predicate, _) in tcx.explicit_item_bounds(def).skip_binder() {
1864 if let ty::ClauseKind::Trait(ref poly_trait_predicate) =
1866 predicate.kind().skip_binder()
1867 {
1868 let def_id = poly_trait_predicate.trait_ref.def_id;
1869 let descr_pre = &format!("{}implementer{} of ", data.descr_pre, plural_suffix);
1870 if check_must_not_suspend_def(
1871 tcx,
1872 def_id,
1873 hir_id,
1874 SuspendCheckData { descr_pre, ..data },
1875 ) {
1876 has_emitted = true;
1877 break;
1878 }
1879 }
1880 }
1881 has_emitted
1882 }
1883 ty::Dynamic(binder, _, _) => {
1884 let mut has_emitted = false;
1885 for predicate in binder.iter() {
1886 if let ty::ExistentialPredicate::Trait(ref trait_ref) = predicate.skip_binder() {
1887 let def_id = trait_ref.def_id;
1888 let descr_post = &format!(" trait object{}{}", plural_suffix, data.descr_post);
1889 if check_must_not_suspend_def(
1890 tcx,
1891 def_id,
1892 hir_id,
1893 SuspendCheckData { descr_post, ..data },
1894 ) {
1895 has_emitted = true;
1896 break;
1897 }
1898 }
1899 }
1900 has_emitted
1901 }
1902 ty::Tuple(fields) => {
1903 let mut has_emitted = false;
1904 for (i, ty) in fields.iter().enumerate() {
1905 let descr_post = &format!(" in tuple element {i}");
1906 if check_must_not_suspend_ty(
1907 tcx,
1908 ty,
1909 hir_id,
1910 SuspendCheckData { descr_post, ..data },
1911 ) {
1912 has_emitted = true;
1913 }
1914 }
1915 has_emitted
1916 }
1917 ty::Array(ty, len) => {
1918 let descr_pre = &format!("{}array{} of ", data.descr_pre, plural_suffix);
1919 check_must_not_suspend_ty(
1920 tcx,
1921 ty,
1922 hir_id,
1923 SuspendCheckData {
1924 descr_pre,
1925 plural_len: len.try_to_target_usize(tcx).unwrap_or(0) as usize + 1,
1927 ..data
1928 },
1929 )
1930 }
1931 ty::Ref(_region, ty, _mutability) => {
1934 let descr_pre = &format!("{}reference{} to ", data.descr_pre, plural_suffix);
1935 check_must_not_suspend_ty(tcx, ty, hir_id, SuspendCheckData { descr_pre, ..data })
1936 }
1937 _ => false,
1938 }
1939}
1940
1941fn check_must_not_suspend_def(
1942 tcx: TyCtxt<'_>,
1943 def_id: DefId,
1944 hir_id: hir::HirId,
1945 data: SuspendCheckData<'_>,
1946) -> bool {
1947 if let Some(attr) = tcx.get_attr(def_id, sym::must_not_suspend) {
1948 let reason = attr.value_str().map(|s| errors::MustNotSuspendReason {
1949 span: data.source_span,
1950 reason: s.as_str().to_string(),
1951 });
1952 tcx.emit_node_span_lint(
1953 rustc_session::lint::builtin::MUST_NOT_SUSPEND,
1954 hir_id,
1955 data.source_span,
1956 errors::MustNotSupend {
1957 tcx,
1958 yield_sp: data.yield_span,
1959 reason,
1960 src_sp: data.source_span,
1961 pre: data.descr_pre,
1962 def_id,
1963 post: data.descr_post,
1964 },
1965 );
1966
1967 true
1968 } else {
1969 false
1970 }
1971}