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::{self as hir, find_attr};
30use rustc_index::bit_set::{BitMatrix, DenseBitSet};
31use rustc_index::{Idx, IndexVec};
32use rustc_middle::mir::*;
33use rustc_middle::span_bug;
34use rustc_middle::ty::{self, CoroutineArgs, CoroutineArgsExt, Ty, TyCtxt, TypingMode};
35use rustc_mir_dataflow::impls::{
36 MaybeBorrowedLocals, MaybeLiveLocals, MaybeRequiresStorage, MaybeStorageLive,
37 always_storage_live_locals,
38};
39use rustc_mir_dataflow::{
40 Analysis, Results, ResultsCursor, ResultsVisitor, visit_reachable_results,
41};
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, trace};
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 body,
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 visit_reachable_results(body, results, &mut visitor);
259
260 let local_conflicts = visitor.local_conflicts;
261
262 let mut storage_conflicts = BitMatrix::new(saved_locals.count(), saved_locals.count());
270 for (saved_local_a, local_a) in saved_locals.iter_enumerated() {
271 if ineligible_locals.contains(local_a) {
272 storage_conflicts.insert_all_into_row(saved_local_a);
274 } else {
275 for (saved_local_b, local_b) in saved_locals.iter_enumerated() {
277 if local_conflicts.contains(local_a, local_b) {
278 storage_conflicts.insert(saved_local_a, saved_local_b);
279 }
280 }
281 }
282 }
283 storage_conflicts
284}
285
286struct StorageConflictVisitor<'a, 'tcx> {
287 body: &'a Body<'tcx>,
288 saved_locals: &'a CoroutineSavedLocals,
289 local_conflicts: BitMatrix<Local, Local>,
292 eligible_storage_live: DenseBitSet<Local>,
294}
295
296impl<'a, 'tcx> ResultsVisitor<'tcx, MaybeRequiresStorage> for StorageConflictVisitor<'a, 'tcx> {
297 fn visit_after_early_statement_effect(
298 &mut self,
299 _analysis: &MaybeRequiresStorage,
300 state: &DenseBitSet<Local>,
301 _statement: &Statement<'tcx>,
302 loc: Location,
303 ) {
304 self.apply_state(state, loc);
305 }
306
307 fn visit_after_early_terminator_effect(
308 &mut self,
309 _analysis: &MaybeRequiresStorage,
310 state: &DenseBitSet<Local>,
311 _terminator: &Terminator<'tcx>,
312 loc: Location,
313 ) {
314 self.apply_state(state, loc);
315 }
316}
317
318impl StorageConflictVisitor<'_, '_> {
319 fn apply_state(&mut self, state: &DenseBitSet<Local>, loc: Location) {
320 if let TerminatorKind::Unreachable = self.body.basic_blocks[loc.block].terminator().kind {
322 return;
323 }
324
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 if self.eligible_storage_live.count() > 1 {
333 trace!("at {:?}, eligible_storage_live={:?}", loc, self.eligible_storage_live);
334 }
335 }
336}
337
338#[tracing::instrument(level = "trace", skip(liveness, body))]
339pub(super) fn compute_layout<'tcx>(
340 liveness: LivenessInfo,
341 body: &Body<'tcx>,
342) -> (
343 IndexVec<Local, Option<(Ty<'tcx>, VariantIdx, FieldIdx)>>,
344 CoroutineLayout<'tcx>,
345 IndexVec<BasicBlock, Option<DenseBitSet<Local>>>,
346) {
347 let LivenessInfo {
348 saved_locals,
349 live_locals_at_suspension_points,
350 source_info_at_suspension_points,
351 storage_conflicts,
352 storage_liveness,
353 } = liveness;
354
355 let mut tys: IndexVec<CoroutineSavedLocal, CoroutineSavedTy<'_>> = saved_locals
357 .iter_enumerated()
358 .map(|(saved_local, local)| {
359 debug!("coroutine saved local {:?} => {:?}", saved_local, local);
360
361 let decl = &body.local_decls[local];
362
363 let ignore_for_traits = match decl.local_info {
368 ClearCrossCrate::Set(LocalInfo::StaticRef { is_thread_local, .. }) => {
371 !is_thread_local
372 }
373 ClearCrossCrate::Set(LocalInfo::FakeBorrow) => true,
376 _ => false,
377 };
378
379 CoroutineSavedTy {
380 ty: decl.ty,
381 source_info: decl.source_info,
382 ignore_for_traits,
383 debuginfo_name: None,
385 }
386 })
387 .collect();
388
389 let body_span = body.source_scopes[OUTERMOST_SOURCE_SCOPE].span;
393 let mut variant_source_info: IndexVec<VariantIdx, SourceInfo> = IndexVec::with_capacity(
394 CoroutineArgs::RESERVED_VARIANTS + live_locals_at_suspension_points.len(),
395 );
396 variant_source_info.extend([
397 SourceInfo::outermost(body_span.shrink_to_lo()),
398 SourceInfo::outermost(body_span.shrink_to_hi()),
399 SourceInfo::outermost(body_span.shrink_to_hi()),
400 ]);
401
402 let reverse_local_map: IndexVec<CoroutineSavedLocal, Local> = saved_locals.iter().collect();
404
405 let mut variant_fields: IndexVec<VariantIdx, _> = IndexVec::from_elem_n(
408 IndexVec::new(),
409 CoroutineArgs::RESERVED_VARIANTS + live_locals_at_suspension_points.len(),
410 );
411 let mut remap = IndexVec::from_elem_n(None, saved_locals.domain_size());
412 for (live_locals, &source_info_at_suspension_point, (variant_index, fields)) in izip!(
413 &live_locals_at_suspension_points,
414 &source_info_at_suspension_points,
415 variant_fields.iter_enumerated_mut().skip(CoroutineArgs::RESERVED_VARIANTS)
416 ) {
417 *fields = live_locals.iter().collect();
418 for (idx, &saved_local) in fields.iter_enumerated() {
419 remap[reverse_local_map[saved_local]] = Some((tys[saved_local].ty, variant_index, idx));
424 }
425 variant_source_info.push(source_info_at_suspension_point);
426 }
427 debug!(?variant_fields);
428 debug!(?storage_conflicts);
429
430 for var in &body.var_debug_info {
431 let VarDebugInfoContents::Place(place) = &var.value else { continue };
432 let Some(local) = place.as_local() else { continue };
433 let Some(&Some((_, variant, field))) = remap.get(local) else {
434 continue;
435 };
436
437 let saved_local: CoroutineSavedLocal = variant_fields[variant][field];
438 tys[saved_local].debuginfo_name.get_or_insert(var.name);
439 }
440
441 let layout =
442 CoroutineLayout { field_tys: tys, variant_fields, variant_source_info, storage_conflicts };
443 debug!(?remap);
444 debug!(?layout);
445 debug!(?storage_liveness);
446
447 (remap, layout, storage_liveness)
448}
449
450#[instrument(level = "debug", skip(tcx), ret)]
451pub(crate) fn mir_coroutine_witnesses<'tcx>(
452 tcx: TyCtxt<'tcx>,
453 def_id: LocalDefId,
454) -> Option<CoroutineLayout<'tcx>> {
455 let (body, _) = tcx.mir_promoted(def_id);
456 let body = body.borrow();
457 let body = &*body;
458
459 let coroutine_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
461
462 let movable = match *coroutine_ty.kind() {
463 ty::Coroutine(def_id, _) => tcx.coroutine_movability(def_id) == hir::Movability::Movable,
464 ty::Error(_) => return None,
465 _ => span_bug!(body.span, "unexpected coroutine type {}", coroutine_ty),
466 };
467
468 let always_live_locals = always_storage_live_locals(body);
471 let liveness_info = locals_live_across_suspend_points(tcx, body, &always_live_locals, movable);
472
473 let (_, coroutine_layout, _) = compute_layout(liveness_info, body);
477
478 check_suspend_tys(tcx, &coroutine_layout, body);
479 check_field_tys_sized(tcx, &coroutine_layout, def_id);
480
481 Some(coroutine_layout)
482}
483
484fn check_field_tys_sized<'tcx>(
485 tcx: TyCtxt<'tcx>,
486 coroutine_layout: &CoroutineLayout<'tcx>,
487 def_id: LocalDefId,
488) {
489 if !tcx.features().unsized_fn_params() {
492 return;
493 }
494
495 let infcx = tcx.infer_ctxt().ignoring_regions().build(TypingMode::non_body_analysis());
500 let param_env = tcx.param_env(def_id);
501
502 let ocx = ObligationCtxt::new_with_diagnostics(&infcx);
503 for field_ty in &coroutine_layout.field_tys {
504 ocx.register_bound(
505 ObligationCause::new(
506 field_ty.source_info.span,
507 def_id,
508 ObligationCauseCode::SizedCoroutineInterior(def_id),
509 ),
510 param_env,
511 field_ty.ty,
512 tcx.require_lang_item(hir::LangItem::Sized, field_ty.source_info.span),
513 );
514 }
515
516 let errors = ocx.evaluate_obligations_error_on_ambiguity();
517 debug!(?errors);
518 if !errors.is_empty() {
519 infcx.err_ctxt().report_fulfillment_errors(errors);
520 }
521}
522
523fn check_suspend_tys<'tcx>(tcx: TyCtxt<'tcx>, layout: &CoroutineLayout<'tcx>, body: &Body<'tcx>) {
524 let mut linted_tys = FxHashSet::default();
525
526 for (variant, yield_source_info) in
527 layout.variant_fields.iter().zip(&layout.variant_source_info)
528 {
529 debug!(?variant);
530 for &local in variant {
531 let decl = &layout.field_tys[local];
532 debug!(?decl);
533
534 if !decl.ignore_for_traits && linted_tys.insert(decl.ty) {
535 let Some(hir_id) = decl.source_info.scope.lint_root(&body.source_scopes) else {
536 continue;
537 };
538
539 check_must_not_suspend_ty(
540 tcx,
541 decl.ty,
542 hir_id,
543 SuspendCheckData {
544 source_span: decl.source_info.span,
545 yield_span: yield_source_info.span,
546 plural_len: 1,
547 ..Default::default()
548 },
549 );
550 }
551 }
552 }
553}
554
555#[derive(Default)]
556struct SuspendCheckData<'a> {
557 source_span: Span,
558 yield_span: Span,
559 descr_pre: &'a str,
560 descr_post: &'a str,
561 plural_len: usize,
562}
563
564fn check_must_not_suspend_ty<'tcx>(
571 tcx: TyCtxt<'tcx>,
572 ty: Ty<'tcx>,
573 hir_id: hir::HirId,
574 data: SuspendCheckData<'_>,
575) -> bool {
576 if ty.is_unit() {
577 return false;
578 }
579
580 let plural_suffix = pluralize!(data.plural_len);
581
582 debug!("Checking must_not_suspend for {}", ty);
583
584 match *ty.kind() {
585 ty::Adt(_, args) if ty.is_box() => {
586 let boxed_ty = args.type_at(0);
587 let allocator_ty = args.type_at(1);
588 check_must_not_suspend_ty(
589 tcx,
590 boxed_ty,
591 hir_id,
592 SuspendCheckData { descr_pre: &format!("{}boxed ", data.descr_pre), ..data },
593 ) || check_must_not_suspend_ty(
594 tcx,
595 allocator_ty,
596 hir_id,
597 SuspendCheckData { descr_pre: &format!("{}allocator ", data.descr_pre), ..data },
598 )
599 }
600 ty::Adt(def, _) if def.repr().scalable() => {
606 tcx.dcx()
607 .span_err(data.source_span, "scalable vectors cannot be held over await points");
608 true
609 }
610 ty::Adt(def, _) => check_must_not_suspend_def(tcx, def.did(), hir_id, data),
611 ty::Alias(_, ty::AliasTy { kind: ty::Opaque { def_id: def }, .. }) => {
613 let mut has_emitted = false;
614 for &(predicate, _) in tcx.explicit_item_bounds(def).skip_binder() {
615 if let ty::ClauseKind::Trait(ref poly_trait_predicate) =
617 predicate.kind().skip_binder()
618 {
619 let def_id = poly_trait_predicate.trait_ref.def_id;
620 let descr_pre = &format!("{}implementer{} of ", data.descr_pre, plural_suffix);
621 if check_must_not_suspend_def(
622 tcx,
623 def_id,
624 hir_id,
625 SuspendCheckData { descr_pre, ..data },
626 ) {
627 has_emitted = true;
628 break;
629 }
630 }
631 }
632 has_emitted
633 }
634 ty::Dynamic(binder, _) => {
635 let mut has_emitted = false;
636 for predicate in binder.iter() {
637 if let ty::ExistentialPredicate::Trait(ref trait_ref) = predicate.skip_binder() {
638 let def_id = trait_ref.def_id;
639 let descr_post = &format!(" trait object{}{}", plural_suffix, data.descr_post);
640 if check_must_not_suspend_def(
641 tcx,
642 def_id,
643 hir_id,
644 SuspendCheckData { descr_post, ..data },
645 ) {
646 has_emitted = true;
647 break;
648 }
649 }
650 }
651 has_emitted
652 }
653 ty::Tuple(fields) => {
654 let mut has_emitted = false;
655 for (i, ty) in fields.iter().enumerate() {
656 let descr_post = &format!(" in tuple element {i}");
657 if check_must_not_suspend_ty(
658 tcx,
659 ty,
660 hir_id,
661 SuspendCheckData { descr_post, ..data },
662 ) {
663 has_emitted = true;
664 }
665 }
666 has_emitted
667 }
668 ty::Array(ty, len) => {
669 let descr_pre = &format!("{}array{} of ", data.descr_pre, plural_suffix);
670 check_must_not_suspend_ty(
671 tcx,
672 ty,
673 hir_id,
674 SuspendCheckData {
675 descr_pre,
676 plural_len: len.try_to_target_usize(tcx).unwrap_or(0) as usize + 1,
678 ..data
679 },
680 )
681 }
682 ty::Ref(_region, ty, _mutability) => {
685 let descr_pre = &format!("{}reference{} to ", data.descr_pre, plural_suffix);
686 check_must_not_suspend_ty(tcx, ty, hir_id, SuspendCheckData { descr_pre, ..data })
687 }
688 _ => false,
689 }
690}
691
692fn check_must_not_suspend_def(
693 tcx: TyCtxt<'_>,
694 def_id: DefId,
695 hir_id: hir::HirId,
696 data: SuspendCheckData<'_>,
697) -> bool {
698 if let Some(reason_str) = find_attr!(tcx, def_id, MustNotSupend {reason} => reason) {
699 let reason = reason_str.map(|s| MustNotSuspendReason { span: data.source_span, reason: s });
700 tcx.emit_node_span_lint(
701 rustc_session::lint::builtin::MUST_NOT_SUSPEND,
702 hir_id,
703 data.source_span,
704 MustNotSupend {
705 tcx,
706 yield_sp: data.yield_span,
707 reason,
708 src_sp: data.source_span,
709 pre: data.descr_pre,
710 def_id,
711 post: data.descr_post,
712 },
713 );
714
715 true
716 } else {
717 false
718 }
719}