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 _analysis: &MaybeRequiresStorage,
305 state: &DenseBitSet<Local>,
306 _statement: &Statement<'tcx>,
307 _loc: Location,
308 ) {
309 self.apply_state(state);
310 }
311
312 fn visit_after_early_terminator_effect(
313 &mut self,
314 _analysis: &MaybeRequiresStorage,
315 state: &DenseBitSet<Local>,
316 _terminator: &Terminator<'tcx>,
317 _loc: Location,
318 ) {
319 self.apply_state(state);
320 }
321}
322
323impl StorageConflictVisitor<'_> {
324 fn apply_state(&mut self, state: &DenseBitSet<Local>) {
325 self.eligible_storage_live.clone_from(state);
326 self.eligible_storage_live.intersect(&**self.saved_locals);
327
328 for local in self.eligible_storage_live.iter() {
329 self.local_conflicts.union_row_with(&self.eligible_storage_live, local);
330 }
331 }
332}
333
334#[tracing::instrument(level = "trace", skip(liveness, body))]
335pub(super) fn compute_layout<'tcx>(
336 liveness: LivenessInfo,
337 body: &Body<'tcx>,
338) -> (
339 IndexVec<Local, Option<(Ty<'tcx>, VariantIdx, FieldIdx)>>,
340 CoroutineLayout<'tcx>,
341 IndexVec<BasicBlock, Option<DenseBitSet<Local>>>,
342) {
343 let LivenessInfo {
344 saved_locals,
345 live_locals_at_suspension_points,
346 source_info_at_suspension_points,
347 storage_conflicts,
348 storage_liveness,
349 } = liveness;
350
351 let mut tys: IndexVec<CoroutineSavedLocal, CoroutineSavedTy<'_>> = saved_locals
353 .iter_enumerated()
354 .map(|(saved_local, local)| {
355 debug!("coroutine saved local {:?} => {:?}", saved_local, local);
356
357 let decl = &body.local_decls[local];
358
359 let ignore_for_traits = match decl.local_info {
364 ClearCrossCrate::Set(LocalInfo::StaticRef { is_thread_local, .. }) => {
367 !is_thread_local
368 }
369 ClearCrossCrate::Set(LocalInfo::FakeBorrow) => true,
372 _ => false,
373 };
374
375 CoroutineSavedTy {
376 ty: decl.ty,
377 source_info: decl.source_info,
378 ignore_for_traits,
379 debuginfo_name: None,
381 }
382 })
383 .collect();
384
385 let body_span = body.source_scopes[OUTERMOST_SOURCE_SCOPE].span;
389 let mut variant_source_info: IndexVec<VariantIdx, SourceInfo> = IndexVec::with_capacity(
390 CoroutineArgs::RESERVED_VARIANTS + live_locals_at_suspension_points.len(),
391 );
392 variant_source_info.extend([
393 SourceInfo::outermost(body_span.shrink_to_lo()),
394 SourceInfo::outermost(body_span.shrink_to_hi()),
395 SourceInfo::outermost(body_span.shrink_to_hi()),
396 ]);
397
398 let reverse_local_map: IndexVec<CoroutineSavedLocal, Local> = saved_locals.iter().collect();
400
401 let mut variant_fields: IndexVec<VariantIdx, _> = IndexVec::from_elem_n(
404 IndexVec::new(),
405 CoroutineArgs::RESERVED_VARIANTS + live_locals_at_suspension_points.len(),
406 );
407 let mut remap = IndexVec::from_elem_n(None, saved_locals.domain_size());
408 for (live_locals, &source_info_at_suspension_point, (variant_index, fields)) in izip!(
409 &live_locals_at_suspension_points,
410 &source_info_at_suspension_points,
411 variant_fields.iter_enumerated_mut().skip(CoroutineArgs::RESERVED_VARIANTS)
412 ) {
413 *fields = live_locals.iter().collect();
414 for (idx, &saved_local) in fields.iter_enumerated() {
415 remap[reverse_local_map[saved_local]] = Some((tys[saved_local].ty, variant_index, idx));
420 }
421 variant_source_info.push(source_info_at_suspension_point);
422 }
423 debug!(?variant_fields);
424 debug!(?storage_conflicts);
425
426 for var in &body.var_debug_info {
427 let VarDebugInfoContents::Place(place) = &var.value else { continue };
428 let Some(local) = place.as_local() else { continue };
429 let Some(&Some((_, variant, field))) = remap.get(local) else {
430 continue;
431 };
432
433 let saved_local: CoroutineSavedLocal = variant_fields[variant][field];
434 tys[saved_local].debuginfo_name.get_or_insert(var.name);
435 }
436
437 let layout =
438 CoroutineLayout { field_tys: tys, variant_fields, variant_source_info, storage_conflicts };
439 debug!(?remap);
440 debug!(?layout);
441 debug!(?storage_liveness);
442
443 (remap, layout, storage_liveness)
444}
445
446#[instrument(level = "debug", skip(tcx), ret)]
447pub(crate) fn mir_coroutine_witnesses<'tcx>(
448 tcx: TyCtxt<'tcx>,
449 def_id: LocalDefId,
450) -> Option<CoroutineLayout<'tcx>> {
451 let (body, _) = tcx.mir_promoted(def_id);
452 let body = body.borrow();
453 let body = &*body;
454
455 let coroutine_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
457
458 let movable = match *coroutine_ty.kind() {
459 ty::Coroutine(def_id, _) => tcx.coroutine_movability(def_id) == hir::Movability::Movable,
460 ty::Error(_) => return None,
461 _ => span_bug!(body.span, "unexpected coroutine type {}", coroutine_ty),
462 };
463
464 let always_live_locals = always_storage_live_locals(body);
467 let liveness_info = locals_live_across_suspend_points(tcx, body, &always_live_locals, movable);
468
469 let (_, coroutine_layout, _) = compute_layout(liveness_info, body);
473
474 check_suspend_tys(tcx, &coroutine_layout, body);
475 check_field_tys_sized(tcx, &coroutine_layout, def_id);
476
477 Some(coroutine_layout)
478}
479
480fn check_field_tys_sized<'tcx>(
481 tcx: TyCtxt<'tcx>,
482 coroutine_layout: &CoroutineLayout<'tcx>,
483 def_id: LocalDefId,
484) {
485 if !tcx.features().unsized_fn_params() {
488 return;
489 }
490
491 let infcx = tcx.infer_ctxt().ignoring_regions().build(TypingMode::non_body_analysis());
496 let param_env = tcx.param_env(def_id);
497
498 let ocx = ObligationCtxt::new_with_diagnostics(&infcx);
499 for field_ty in &coroutine_layout.field_tys {
500 ocx.register_bound(
501 ObligationCause::new(
502 field_ty.source_info.span,
503 def_id,
504 ObligationCauseCode::SizedCoroutineInterior(def_id),
505 ),
506 param_env,
507 field_ty.ty,
508 tcx.require_lang_item(LangItem::Sized, field_ty.source_info.span),
509 );
510 }
511
512 let errors = ocx.evaluate_obligations_error_on_ambiguity();
513 debug!(?errors);
514 if let TraitErrors::HasErrors(errors) = errors {
515 infcx.err_ctxt().report_fulfillment_errors(errors);
516 }
517}
518
519fn check_suspend_tys<'tcx>(tcx: TyCtxt<'tcx>, layout: &CoroutineLayout<'tcx>, body: &Body<'tcx>) {
520 let mut linted_tys = FxHashSet::default();
521
522 for (variant, yield_source_info) in
523 layout.variant_fields.iter().zip(&layout.variant_source_info)
524 {
525 debug!(?variant);
526 for &local in variant {
527 let decl = &layout.field_tys[local];
528 debug!(?decl);
529
530 if !decl.ignore_for_traits && linted_tys.insert(decl.ty) {
531 let Some(hir_id) = decl.source_info.scope.lint_root(&body.source_scopes) else {
532 continue;
533 };
534
535 check_must_not_suspend_ty(
536 tcx,
537 decl.ty,
538 hir_id,
539 SuspendCheckData {
540 source_span: decl.source_info.span,
541 yield_span: yield_source_info.span,
542 plural_len: 1,
543 ..Default::default()
544 },
545 );
546 }
547 }
548 }
549}
550
551#[derive(Default)]
552struct SuspendCheckData<'a> {
553 source_span: Span,
554 yield_span: Span,
555 descr_pre: &'a str,
556 descr_post: &'a str,
557 plural_len: usize,
558}
559
560fn check_must_not_suspend_ty<'tcx>(
567 tcx: TyCtxt<'tcx>,
568 ty: Ty<'tcx>,
569 hir_id: hir::HirId,
570 data: SuspendCheckData<'_>,
571) -> bool {
572 if ty.is_unit() {
573 return false;
574 }
575
576 let plural_suffix = pluralize!(data.plural_len);
577
578 debug!("Checking must_not_suspend for {}", ty);
579
580 match *ty.kind() {
581 ty::Adt(_, args) if ty.is_box() => {
582 let boxed_ty = args.type_at(0);
583 let allocator_ty = args.type_at(1);
584 check_must_not_suspend_ty(
585 tcx,
586 boxed_ty,
587 hir_id,
588 SuspendCheckData { descr_pre: &format!("{}boxed ", data.descr_pre), ..data },
589 ) || check_must_not_suspend_ty(
590 tcx,
591 allocator_ty,
592 hir_id,
593 SuspendCheckData { descr_pre: &format!("{}allocator ", data.descr_pre), ..data },
594 )
595 }
596 ty::Adt(def, _) if def.repr().scalable() => {
602 tcx.dcx()
603 .span_err(data.source_span, "scalable vectors cannot be held over await points");
604 true
605 }
606 ty::Adt(def, _) => check_must_not_suspend_def(tcx, def.did(), hir_id, data),
607 ty::Alias(_, ty::AliasTy { kind: ty::Opaque { def_id: def }, .. }) => {
609 let mut has_emitted = false;
610 for &(predicate, _) in tcx.explicit_item_bounds(def).skip_binder() {
611 if let ty::ClauseKind::Trait(ref poly_trait_predicate) =
613 predicate.kind().skip_binder()
614 {
615 let def_id = poly_trait_predicate.trait_ref.def_id;
616 let descr_pre = &format!("{}implementer{} of ", data.descr_pre, plural_suffix);
617 if check_must_not_suspend_def(
618 tcx,
619 def_id,
620 hir_id,
621 SuspendCheckData { descr_pre, ..data },
622 ) {
623 has_emitted = true;
624 break;
625 }
626 }
627 }
628 has_emitted
629 }
630 ty::Dynamic(binder, _) => {
631 let mut has_emitted = false;
632 for predicate in binder.iter() {
633 if let ty::ExistentialPredicate::Trait(ref trait_ref) = predicate.skip_binder() {
634 let def_id = trait_ref.def_id;
635 let descr_post = &format!(" trait object{}{}", plural_suffix, data.descr_post);
636 if check_must_not_suspend_def(
637 tcx,
638 def_id,
639 hir_id,
640 SuspendCheckData { descr_post, ..data },
641 ) {
642 has_emitted = true;
643 break;
644 }
645 }
646 }
647 has_emitted
648 }
649 ty::Tuple(fields) => {
650 let mut has_emitted = false;
651 for (i, ty) in fields.iter().enumerate() {
652 let descr_post = &format!(" in tuple element {i}");
653 if check_must_not_suspend_ty(
654 tcx,
655 ty,
656 hir_id,
657 SuspendCheckData { descr_post, ..data },
658 ) {
659 has_emitted = true;
660 }
661 }
662 has_emitted
663 }
664 ty::Array(ty, len) => {
665 let descr_pre = &format!("{}array{} of ", data.descr_pre, plural_suffix);
666 check_must_not_suspend_ty(
667 tcx,
668 ty,
669 hir_id,
670 SuspendCheckData {
671 descr_pre,
672 plural_len: len.try_to_target_usize(tcx).unwrap_or(0) as usize + 1,
674 ..data
675 },
676 )
677 }
678 ty::Ref(_region, ty, _mutability) => {
681 let descr_pre = &format!("{}reference{} to ", data.descr_pre, plural_suffix);
682 check_must_not_suspend_ty(tcx, ty, hir_id, SuspendCheckData { descr_pre, ..data })
683 }
684 _ => false,
685 }
686}
687
688fn check_must_not_suspend_def(
689 tcx: TyCtxt<'_>,
690 def_id: DefId,
691 hir_id: hir::HirId,
692 data: SuspendCheckData<'_>,
693) -> bool {
694 if let Some(reason_str) = find_attr!(tcx, def_id, MustNotSupend {reason} => reason) {
695 let reason = reason_str.map(|s| MustNotSuspendReason { span: data.source_span, reason: s });
696 tcx.emit_node_span_lint(
697 rustc_session::lint::builtin::MUST_NOT_SUSPEND,
698 hir_id,
699 data.source_span,
700 MustNotSupend {
701 tcx,
702 yield_sp: data.yield_span,
703 reason,
704 src_sp: data.source_span,
705 pre: data.descr_pre,
706 def_id,
707 post: data.descr_post,
708 },
709 );
710
711 true
712 } else {
713 false
714 }
715}