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