1use std::ops;
24
25use itertools::izip;
26use rustc_abi::{FieldIdx, VariantIdx};
27use rustc_data_structures::fx::FxHashSet;
28use rustc_errors::pluralize;
29use rustc_hir::attrs::lang_items::LangItem;
30use rustc_hir::{self as hir, find_attr};
31use rustc_index::bit_set::{BitMatrix, DenseBitSet};
32use rustc_index::{Idx, IndexVec};
33use rustc_infer::traits::TraitErrors;
34use rustc_middle::mir::*;
35use rustc_middle::span_bug;
36use rustc_middle::ty::{self, CoroutineArgs, CoroutineArgsExt, Ty, TyCtxt, TypingMode};
37use rustc_mir_dataflow::impls::{
38 MaybeBorrowedLocals, MaybeLiveLocals, MaybeRequiresStorage, MaybeStorageLive,
39 always_storage_live_locals,
40};
41use rustc_mir_dataflow::{Analysis, Results, ResultsCursor, ResultsVisitor, visit_results};
42use rustc_span::Span;
43use rustc_span::def_id::{DefId, LocalDefId};
44use rustc_trait_selection::error_reporting::InferCtxtErrorExt;
45use rustc_trait_selection::infer::TyCtxtInferExt as _;
46use rustc_trait_selection::traits::{ObligationCause, ObligationCauseCode, ObligationCtxt};
47use tracing::{debug, instrument};
48
49use crate::diagnostics::{MustNotSupend, MustNotSuspendReason};
50
51const SELF_ARG: Local = Local::arg(0);
52
53pub(super) struct LivenessInfo {
54 pub(super) saved_locals: CoroutineSavedLocals,
56
57 live_locals_at_suspension_points: Vec<DenseBitSet<CoroutineSavedLocal>>,
59
60 source_info_at_suspension_points: Vec<SourceInfo>,
62
63 pub(super) storage_conflicts: BitMatrix<CoroutineSavedLocal, CoroutineSavedLocal>,
67
68 storage_liveness: IndexVec<BasicBlock, Option<DenseBitSet<Local>>>,
71}
72
73#[tracing::instrument(level = "trace", skip(tcx, body))]
82pub(super) fn locals_live_across_suspend_points<'tcx>(
83 tcx: TyCtxt<'tcx>,
84 body: &Body<'tcx>,
85 always_live_locals: &DenseBitSet<Local>,
86 movable: bool,
87) -> LivenessInfo {
88 let mut storage_live = MaybeStorageLive::new(std::borrow::Cow::Borrowed(always_live_locals))
91 .iterate_to_fixpoint(tcx, body, None)
92 .into_results_cursor(body);
93
94 let borrowed_locals = MaybeBorrowedLocals.iterate_to_fixpoint(tcx, body, Some("coroutine"));
96
97 let requires_storage =
99 MaybeRequiresStorage::new(body, &borrowed_locals).iterate_to_fixpoint(tcx, body, None);
100 let mut requires_storage_cursor = ResultsCursor::new_borrowing(body, &requires_storage);
101
102 let mut liveness =
104 MaybeLiveLocals.iterate_to_fixpoint(tcx, body, Some("coroutine")).into_results_cursor(body);
105
106 let mut storage_liveness_map = IndexVec::from_elem(None, &body.basic_blocks);
107 let mut live_locals_at_suspension_points = Vec::new();
108 let mut source_info_at_suspension_points = Vec::new();
109 let mut live_locals_at_any_suspension_point = DenseBitSet::new_empty(body.local_decls.len());
110 let mut borrowed_locals_cursor = ResultsCursor::new_owning(body, borrowed_locals);
111
112 for (block, data) in body.basic_blocks.iter_enumerated() {
113 let TerminatorKind::Yield { .. } = data.terminator().kind else { continue };
114
115 let loc = Location { block, statement_index: data.statements.len() };
116
117 liveness.seek_to_block_end(block);
118 let mut live_locals = liveness.get().clone();
119
120 if !movable {
121 borrowed_locals_cursor.seek_before_primary_effect(loc);
132 live_locals.union(borrowed_locals_cursor.get());
133 }
134
135 storage_live.seek_before_primary_effect(loc);
138 storage_liveness_map[block] = Some(storage_live.get().clone());
139
140 requires_storage_cursor.seek_before_primary_effect(loc);
144 live_locals.intersect(requires_storage_cursor.get());
145
146 live_locals.remove(SELF_ARG);
148
149 debug!(?loc, ?live_locals);
150
151 live_locals_at_any_suspension_point.union(&live_locals);
154
155 live_locals_at_suspension_points.push(live_locals);
156 source_info_at_suspension_points.push(data.terminator().source_info);
157 }
158
159 debug!(?live_locals_at_any_suspension_point);
160 let saved_locals = CoroutineSavedLocals(live_locals_at_any_suspension_point);
161
162 let live_locals_at_suspension_points = live_locals_at_suspension_points
165 .iter()
166 .map(|live_here| saved_locals.renumber_bitset(live_here))
167 .collect();
168
169 let storage_conflicts = compute_storage_conflicts(
170 body,
171 &saved_locals,
172 always_live_locals.clone(),
173 &requires_storage,
174 );
175
176 LivenessInfo {
177 saved_locals,
178 live_locals_at_suspension_points,
179 source_info_at_suspension_points,
180 storage_conflicts,
181 storage_liveness: storage_liveness_map,
182 }
183}
184
185pub(super) struct CoroutineSavedLocals(DenseBitSet<Local>);
191
192impl CoroutineSavedLocals {
193 fn iter_enumerated(&self) -> impl '_ + Iterator<Item = (CoroutineSavedLocal, Local)> {
196 self.iter().enumerate().map(|(i, l)| (CoroutineSavedLocal::from(i), l))
197 }
198
199 fn renumber_bitset(&self, input: &DenseBitSet<Local>) -> DenseBitSet<CoroutineSavedLocal> {
202 assert!(self.superset(input), "{:?} not a superset of {:?}", self.0, input);
203 let mut out = DenseBitSet::new_empty(self.count());
204 for (saved_local, local) in self.iter_enumerated() {
205 if input.contains(local) {
206 out.insert(saved_local);
207 }
208 }
209 out
210 }
211
212 pub(super) fn get(&self, local: Local) -> Option<CoroutineSavedLocal> {
213 if !self.contains(local) {
214 return None;
215 }
216
217 let idx = self.iter().take_while(|&l| l < local).count();
218 Some(CoroutineSavedLocal::new(idx))
219 }
220}
221
222impl ops::Deref for CoroutineSavedLocals {
223 type Target = DenseBitSet<Local>;
224
225 fn deref(&self) -> &Self::Target {
226 &self.0
227 }
228}
229
230fn compute_storage_conflicts<'mir, 'tcx>(
235 body: &'mir Body<'tcx>,
236 saved_locals: &'mir CoroutineSavedLocals,
237 always_live_locals: DenseBitSet<Local>,
238 results: &Results<'tcx, MaybeRequiresStorage>,
239) -> BitMatrix<CoroutineSavedLocal, CoroutineSavedLocal> {
240 assert_eq!(body.local_decls.len(), saved_locals.domain_size());
241
242 debug!("compute_storage_conflicts({:?})", body.span);
243 debug!("always_live = {:?}", always_live_locals);
244
245 let mut ineligible_locals = always_live_locals;
248 ineligible_locals.intersect(&**saved_locals);
249
250 let mut visitor = StorageConflictVisitor {
252 saved_locals,
253 local_conflicts: BitMatrix::from_row_n(&ineligible_locals, body.local_decls.len()),
254 eligible_storage_live: DenseBitSet::new_empty(body.local_decls.len()),
255 };
256
257 let blocks = traversal::reachable(body).filter_map(|(bb, data)| {
262 (!matches!(data.terminator().kind, TerminatorKind::Unreachable)).then_some(bb)
263 });
264 visit_results(body, blocks, results, &mut visitor);
265
266 let local_conflicts = visitor.local_conflicts;
267
268 let mut storage_conflicts = BitMatrix::new(saved_locals.count(), saved_locals.count());
276 for (saved_local_a, local_a) in saved_locals.iter_enumerated() {
277 if ineligible_locals.contains(local_a) {
278 storage_conflicts.insert_all_into_row(saved_local_a);
280 } else {
281 for (saved_local_b, local_b) in saved_locals.iter_enumerated() {
283 if local_conflicts.contains(local_a, local_b) {
284 storage_conflicts.insert(saved_local_a, saved_local_b);
285 }
286 }
287 }
288 }
289 storage_conflicts
290}
291
292struct StorageConflictVisitor<'a> {
293 saved_locals: &'a CoroutineSavedLocals,
294 local_conflicts: BitMatrix<Local, Local>,
297 eligible_storage_live: DenseBitSet<Local>,
299}
300
301impl<'a, 'tcx> ResultsVisitor<'tcx, MaybeRequiresStorage> for StorageConflictVisitor<'a> {
302 fn visit_after_early_statement_effect(
303 &mut self,
304 state: &DenseBitSet<Local>,
305 _statement: &Statement<'tcx>,
306 _loc: Location,
307 ) {
308 self.apply_state(state);
309 }
310
311 fn visit_after_early_terminator_effect(
312 &mut self,
313 state: &DenseBitSet<Local>,
314 _terminator: &Terminator<'tcx>,
315 _loc: Location,
316 ) {
317 self.apply_state(state);
318 }
319}
320
321impl StorageConflictVisitor<'_> {
322 fn apply_state(&mut self, state: &DenseBitSet<Local>) {
323 self.eligible_storage_live.clone_from(state);
324 self.eligible_storage_live.intersect(&**self.saved_locals);
325
326 for local in self.eligible_storage_live.iter() {
327 self.local_conflicts.union_row_with(&self.eligible_storage_live, local);
328 }
329 }
330}
331
332#[tracing::instrument(level = "trace", skip(liveness, body))]
333pub(super) fn compute_layout<'tcx>(
334 liveness: LivenessInfo,
335 body: &Body<'tcx>,
336) -> (
337 IndexVec<Local, Option<(Ty<'tcx>, VariantIdx, FieldIdx)>>,
338 CoroutineLayout<'tcx>,
339 IndexVec<BasicBlock, Option<DenseBitSet<Local>>>,
340) {
341 let LivenessInfo {
342 saved_locals,
343 live_locals_at_suspension_points,
344 source_info_at_suspension_points,
345 storage_conflicts,
346 storage_liveness,
347 } = liveness;
348
349 let mut tys: IndexVec<CoroutineSavedLocal, CoroutineSavedTy<'_>> = saved_locals
351 .iter_enumerated()
352 .map(|(saved_local, local)| {
353 debug!("coroutine saved local {:?} => {:?}", saved_local, local);
354
355 let decl = &body.local_decls[local];
356
357 let ignore_for_traits = match decl.local_info {
362 ClearCrossCrate::Set(LocalInfo::StaticRef { is_thread_local, .. }) => {
365 !is_thread_local
366 }
367 ClearCrossCrate::Set(LocalInfo::FakeBorrow) => true,
370 _ => false,
371 };
372
373 CoroutineSavedTy {
374 ty: decl.ty,
375 source_info: decl.source_info,
376 ignore_for_traits,
377 debuginfo_name: None,
379 }
380 })
381 .collect();
382
383 let body_span = body.source_scopes[OUTERMOST_SOURCE_SCOPE].span;
387 let mut variant_source_info: IndexVec<VariantIdx, SourceInfo> = IndexVec::with_capacity(
388 CoroutineArgs::RESERVED_VARIANTS + live_locals_at_suspension_points.len(),
389 );
390 variant_source_info.extend([
391 SourceInfo::outermost(body_span.shrink_to_lo()),
392 SourceInfo::outermost(body_span.shrink_to_hi()),
393 SourceInfo::outermost(body_span.shrink_to_hi()),
394 ]);
395
396 let reverse_local_map: IndexVec<CoroutineSavedLocal, Local> = saved_locals.iter().collect();
398
399 let mut variant_fields: IndexVec<VariantIdx, _> = IndexVec::from_elem_n(
402 IndexVec::new(),
403 CoroutineArgs::RESERVED_VARIANTS + live_locals_at_suspension_points.len(),
404 );
405 let mut remap = IndexVec::from_elem_n(None, saved_locals.domain_size());
406 for (live_locals, &source_info_at_suspension_point, (variant_index, fields)) in izip!(
407 &live_locals_at_suspension_points,
408 &source_info_at_suspension_points,
409 variant_fields.iter_enumerated_mut().skip(CoroutineArgs::RESERVED_VARIANTS)
410 ) {
411 *fields = live_locals.iter().collect();
412 for (idx, &saved_local) in fields.iter_enumerated() {
413 remap[reverse_local_map[saved_local]] = Some((tys[saved_local].ty, variant_index, idx));
418 }
419 variant_source_info.push(source_info_at_suspension_point);
420 }
421 debug!(?variant_fields);
422 debug!(?storage_conflicts);
423
424 for var in &body.var_debug_info {
425 let VarDebugInfoContents::Place(place) = &var.value else { continue };
426 let Some(local) = place.as_local() else { continue };
427 let Some(&Some((_, variant, field))) = remap.get(local) else {
428 continue;
429 };
430
431 let saved_local: CoroutineSavedLocal = variant_fields[variant][field];
432 tys[saved_local].debuginfo_name.get_or_insert(var.name);
433 }
434
435 let layout =
436 CoroutineLayout { field_tys: tys, variant_fields, variant_source_info, storage_conflicts };
437 debug!(?remap);
438 debug!(?layout);
439 debug!(?storage_liveness);
440
441 (remap, layout, storage_liveness)
442}
443
444#[instrument(level = "debug", skip(tcx), ret)]
445pub(crate) fn mir_coroutine_witnesses<'tcx>(
446 tcx: TyCtxt<'tcx>,
447 def_id: LocalDefId,
448) -> Option<CoroutineLayout<'tcx>> {
449 let (body, _) = tcx.mir_promoted(def_id);
450 let body = body.borrow();
451 let body = &*body;
452
453 let coroutine_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
455
456 let movable = match *coroutine_ty.kind() {
457 ty::Coroutine(def_id, _) => tcx.coroutine_movability(def_id) == hir::Movability::Movable,
458 ty::Error(_) => return None,
459 _ => span_bug!(body.span, "unexpected coroutine type {}", coroutine_ty),
460 };
461
462 let always_live_locals = always_storage_live_locals(body);
465 let liveness_info = locals_live_across_suspend_points(tcx, body, &always_live_locals, movable);
466
467 let (_, coroutine_layout, _) = compute_layout(liveness_info, body);
471
472 check_suspend_tys(tcx, &coroutine_layout, body);
473 check_field_tys_sized(tcx, &coroutine_layout, def_id);
474
475 Some(coroutine_layout)
476}
477
478fn check_field_tys_sized<'tcx>(
479 tcx: TyCtxt<'tcx>,
480 coroutine_layout: &CoroutineLayout<'tcx>,
481 def_id: LocalDefId,
482) {
483 if !tcx.features().unsized_fn_params() {
486 return;
487 }
488
489 let infcx = tcx.infer_ctxt().ignoring_regions().build(TypingMode::non_body_analysis());
494 let param_env = tcx.param_env(def_id);
495
496 let ocx = ObligationCtxt::new_with_diagnostics(&infcx);
497 for field_ty in &coroutine_layout.field_tys {
498 ocx.register_bound(
499 ObligationCause::new(
500 field_ty.source_info.span,
501 def_id,
502 ObligationCauseCode::SizedCoroutineInterior(def_id),
503 ),
504 param_env,
505 field_ty.ty,
506 tcx.require_lang_item(LangItem::Sized, field_ty.source_info.span),
507 );
508 }
509
510 let errors = ocx.evaluate_obligations_error_on_ambiguity();
511 debug!(?errors);
512 if let TraitErrors::HasErrors(errors) = errors {
513 infcx.err_ctxt().report_fulfillment_errors(errors);
514 }
515}
516
517fn check_suspend_tys<'tcx>(tcx: TyCtxt<'tcx>, layout: &CoroutineLayout<'tcx>, body: &Body<'tcx>) {
518 let mut linted_tys = FxHashSet::default();
519
520 for (variant, yield_source_info) in
521 layout.variant_fields.iter().zip(&layout.variant_source_info)
522 {
523 debug!(?variant);
524 for &local in variant {
525 let decl = &layout.field_tys[local];
526 debug!(?decl);
527
528 if !decl.ignore_for_traits && linted_tys.insert(decl.ty) {
529 let Some(hir_id) = decl.source_info.scope.lint_root(&body.source_scopes) else {
530 continue;
531 };
532
533 check_must_not_suspend_ty(
534 tcx,
535 decl.ty,
536 hir_id,
537 SuspendCheckData {
538 source_span: decl.source_info.span,
539 yield_span: yield_source_info.span,
540 plural_len: 1,
541 ..Default::default()
542 },
543 );
544 }
545 }
546 }
547}
548
549#[derive(Default)]
550struct SuspendCheckData<'a> {
551 source_span: Span,
552 yield_span: Span,
553 descr_pre: &'a str,
554 descr_post: &'a str,
555 plural_len: usize,
556}
557
558fn check_must_not_suspend_ty<'tcx>(
565 tcx: TyCtxt<'tcx>,
566 ty: Ty<'tcx>,
567 hir_id: hir::HirId,
568 data: SuspendCheckData<'_>,
569) -> bool {
570 if ty.is_unit() {
571 return false;
572 }
573
574 let plural_suffix = pluralize!(data.plural_len);
575
576 debug!("Checking must_not_suspend for {}", ty);
577
578 match *ty.kind() {
579 ty::Adt(_, args) if ty.is_box() => {
580 let boxed_ty = args.type_at(0);
581 let allocator_ty = args.type_at(1);
582 check_must_not_suspend_ty(
583 tcx,
584 boxed_ty,
585 hir_id,
586 SuspendCheckData { descr_pre: &format!("{}boxed ", data.descr_pre), ..data },
587 ) || check_must_not_suspend_ty(
588 tcx,
589 allocator_ty,
590 hir_id,
591 SuspendCheckData { descr_pre: &format!("{}allocator ", data.descr_pre), ..data },
592 )
593 }
594 ty::Adt(def, _) if def.repr().scalable() => {
600 tcx.dcx()
601 .span_err(data.source_span, "scalable vectors cannot be held over await points");
602 true
603 }
604 ty::Adt(def, _) => check_must_not_suspend_def(tcx, def.did(), hir_id, data),
605 ty::Alias(_, ty::AliasTy { kind: ty::Opaque { def_id: def }, .. }) => {
607 let mut has_emitted = false;
608 for &(predicate, _) in tcx.explicit_item_bounds(def).skip_binder() {
609 if let ty::ClauseKind::Trait(ref poly_trait_predicate) =
611 predicate.kind().skip_binder()
612 {
613 let def_id = poly_trait_predicate.trait_ref.def_id;
614 let descr_pre = &format!("{}implementer{} of ", data.descr_pre, plural_suffix);
615 if check_must_not_suspend_def(
616 tcx,
617 def_id,
618 hir_id,
619 SuspendCheckData { descr_pre, ..data },
620 ) {
621 has_emitted = true;
622 break;
623 }
624 }
625 }
626 has_emitted
627 }
628 ty::Dynamic(binder, _) => {
629 let mut has_emitted = false;
630 for predicate in binder.iter() {
631 if let ty::ExistentialPredicate::Trait(ref trait_ref) = predicate.skip_binder() {
632 let def_id = trait_ref.def_id;
633 let descr_post = &format!(" trait object{}{}", plural_suffix, data.descr_post);
634 if check_must_not_suspend_def(
635 tcx,
636 def_id,
637 hir_id,
638 SuspendCheckData { descr_post, ..data },
639 ) {
640 has_emitted = true;
641 break;
642 }
643 }
644 }
645 has_emitted
646 }
647 ty::Tuple(fields) => {
648 let mut has_emitted = false;
649 for (i, ty) in fields.iter().enumerate() {
650 let descr_post = &format!(" in tuple element {i}");
651 if check_must_not_suspend_ty(
652 tcx,
653 ty,
654 hir_id,
655 SuspendCheckData { descr_post, ..data },
656 ) {
657 has_emitted = true;
658 }
659 }
660 has_emitted
661 }
662 ty::Array(ty, len) => {
663 let descr_pre = &format!("{}array{} of ", data.descr_pre, plural_suffix);
664 check_must_not_suspend_ty(
665 tcx,
666 ty,
667 hir_id,
668 SuspendCheckData {
669 descr_pre,
670 plural_len: len.try_to_target_usize(tcx).unwrap_or(0) as usize + 1,
672 ..data
673 },
674 )
675 }
676 ty::Ref(_region, ty, _mutability) => {
679 let descr_pre = &format!("{}reference{} to ", data.descr_pre, plural_suffix);
680 check_must_not_suspend_ty(tcx, ty, hir_id, SuspendCheckData { descr_pre, ..data })
681 }
682 _ => false,
683 }
684}
685
686fn check_must_not_suspend_def(
687 tcx: TyCtxt<'_>,
688 def_id: DefId,
689 hir_id: hir::HirId,
690 data: SuspendCheckData<'_>,
691) -> bool {
692 if let Some(reason_str) = find_attr!(tcx, def_id, MustNotSupend {reason} => reason) {
693 let reason = reason_str.map(|s| MustNotSuspendReason { span: data.source_span, reason: s });
694 tcx.emit_node_span_lint(
695 rustc_session::lint::builtin::MUST_NOT_SUSPEND,
696 hir_id,
697 data.source_span,
698 MustNotSupend {
699 tcx,
700 yield_sp: data.yield_span,
701 reason,
702 src_sp: data.source_span,
703 pre: data.descr_pre,
704 def_id,
705 post: data.descr_post,
706 },
707 );
708
709 true
710 } else {
711 false
712 }
713}