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