Separate witness type computation from the generator transform.

compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs

@@ -202,6 +202,7 @@ provide! { tcx, def_id, other, cdata,
     thir_abstract_const => { table }
     optimized_mir => { table }
     mir_for_ctfe => { table }
+    mir_generator_witnesses => { table }
     promoted_mir => { table }
     def_span => { table }
     def_ident_span => { table }

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -1414,6 +1414,10 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
             debug!("EntryBuilder::encode_mir({:?})", def_id);
             if encode_opt {
                 record!(self.tables.optimized_mir[def_id.to_def_id()] <- tcx.optimized_mir(def_id));
+
+                if let DefKind::Generator = self.tcx.def_kind(def_id) {
+                    record!(self.tables.mir_generator_witnesses[def_id.to_def_id()] <- tcx.mir_generator_witnesses(def_id));
+                }
             }
             if encode_const {
                 record!(self.tables.mir_for_ctfe[def_id.to_def_id()] <- tcx.mir_for_ctfe(def_id));

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -376,6 +376,7 @@ define_tables! {
     object_lifetime_default: Table<DefIndex, LazyValue<ObjectLifetimeDefault>>,
     optimized_mir: Table<DefIndex, LazyValue<mir::Body<'static>>>,
     mir_for_ctfe: Table<DefIndex, LazyValue<mir::Body<'static>>>,
+    mir_generator_witnesses: Table<DefIndex, LazyValue<mir::GeneratorLayout<'static>>>,
     promoted_mir: Table<DefIndex, LazyValue<IndexVec<mir::Promoted, mir::Body<'static>>>>,
     // FIXME(compiler-errors): Why isn't this a LazyArray?
     thir_abstract_const: Table<DefIndex, LazyValue<ty::Const<'static>>>,

compiler/rustc_middle/src/query/mod.rs

@@ -471,6 +471,13 @@ rustc_queries! {
         }
     }
 
+    query mir_generator_witnesses(key: DefId) -> mir::GeneratorLayout<'tcx> {
+        arena_cache
+        desc { |tcx| "generator witness types for `{}`", tcx.def_path_str(key) }
+        cache_on_disk_if { key.is_local() }
+        separate_provide_extern
+    }
+
     /// MIR after our optimization passes have run. This is MIR that is ready
     /// for codegen. This is also the only query that can fetch non-local MIR, at present.
     query optimized_mir(key: DefId) -> &'tcx mir::Body<'tcx> {

compiler/rustc_middle/src/ty/parameterized.rs

@@ -117,6 +117,7 @@ macro_rules! parameterized_over_tcx {
 parameterized_over_tcx! {
     crate::middle::exported_symbols::ExportedSymbol,
     crate::mir::Body,
+    crate::mir::GeneratorLayout,
     ty::Ty,
     ty::FnSig,
     ty::GenericPredicates,

compiler/rustc_mir_build/src/build/matches/mod.rs

@@ -1747,8 +1747,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 };
                 let fake_borrow_deref_ty = matched_place.ty(&self.local_decls, tcx).ty;
                 let fake_borrow_ty = tcx.mk_imm_ref(tcx.lifetimes.re_erased, fake_borrow_deref_ty);
-                let fake_borrow_temp =
-                    self.local_decls.push(LocalDecl::new(fake_borrow_ty, temp_span));
+                let mut fake_borrow_temp = LocalDecl::new(fake_borrow_ty, temp_span);
+                fake_borrow_temp.internal = self.local_decls[matched_place.local].internal;
+                let fake_borrow_temp = self.local_decls.push(fake_borrow_temp);
 
                 (matched_place, fake_borrow_temp)
             })

compiler/rustc_mir_transform/src/generator.rs

@@ -54,7 +54,8 @@ use crate::deref_separator::deref_finder;
 use crate::simplify;
 use crate::util::expand_aggregate;
 use crate::MirPass;
-use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::fx::{FxHashMap, FxHashSet};
+use rustc_errors::pluralize;
 use rustc_hir as hir;
 use rustc_hir::lang_items::LangItem;
 use rustc_hir::GeneratorKind;
@@ -70,6 +71,9 @@ use rustc_mir_dataflow::impls::{
 };
 use rustc_mir_dataflow::storage::always_storage_live_locals;
 use rustc_mir_dataflow::{self, Analysis};
+use rustc_span::def_id::DefId;
+use rustc_span::symbol::sym;
+use rustc_span::Span;
 use rustc_target::abi::VariantIdx;
 use rustc_target::spec::PanicStrategy;
 use std::{iter, ops};
@@ -854,7 +858,7 @@ fn sanitize_witness<'tcx>(
     body: &Body<'tcx>,
     witness: Ty<'tcx>,
     upvars: Vec<Ty<'tcx>>,
-    saved_locals: &GeneratorSavedLocals,
+    layout: &GeneratorLayout<'tcx>,
 ) {
     let did = body.source.def_id();
     let param_env = tcx.param_env(did);
@@ -873,31 +877,35 @@ fn sanitize_witness<'tcx>(
         }
     };
 
-    for (local, decl) in body.local_decls.iter_enumerated() {
-        // Ignore locals which are internal or not saved between yields.
-        if !saved_locals.contains(local) || decl.internal {
+    let mut mismatches = Vec::new();
+    for fty in &layout.field_tys {
+        if fty.is_static_ptr {
             continue;
         }
-        let decl_ty = tcx.normalize_erasing_regions(param_env, decl.ty);
+        let decl_ty = tcx.normalize_erasing_regions(param_env, fty.ty);
 
         // Sanity check that typeck knows about the type of locals which are
         // live across a suspension point
         if !allowed.contains(&decl_ty) && !allowed_upvars.contains(&decl_ty) {
-            span_bug!(
-                body.span,
-                "Broken MIR: generator contains type {} in MIR, \
-                       but typeck only knows about {} and {:?}",
-                decl_ty,
-                allowed,
-                allowed_upvars
-            );
+            mismatches.push(decl_ty);
         }
     }
+
+    if !mismatches.is_empty() {
+        span_bug!(
+            body.span,
+            "Broken MIR: generator contains type {:?} in MIR, \
+                       but typeck only knows about {} and {:?}",
+            mismatches,
+            allowed,
+            allowed_upvars
+        );
+    }
 }
 
 fn compute_layout<'tcx>(
     liveness: LivenessInfo,
-    body: &mut Body<'tcx>,
+    body: &Body<'tcx>,
 ) -> (
     FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>,
     GeneratorLayout<'tcx>,
@@ -920,9 +928,7 @@ fn compute_layout<'tcx>(
         let decl = GeneratorSavedTy {
             ty: decl.ty,
             source_info: decl.source_info,
-            is_static_ptr: decl.internal
-                && decl.ty.is_unsafe_ptr()
-                && matches!(decl.local_info.as_deref(), Some(LocalInfo::StaticRef { .. })),
+            is_static_ptr: decl.internal,
         };
         tys.push(decl);
         debug!("generator saved local {:?} => {:?}", saved_local, local);
@@ -1360,6 +1366,52 @@ fn create_cases<'tcx>(
         .collect()
 }
 
+#[instrument(level = "debug", skip(tcx), ret)]
+pub(crate) fn mir_generator_witnesses<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def_id: DefId,
+) -> GeneratorLayout<'tcx> {
+    let def_id = def_id.expect_local();
+
+    let (body, _) = tcx.mir_promoted(ty::WithOptConstParam::unknown(def_id));
+    let body = body.borrow();
+    let body = &*body;
+
+    // The first argument is the generator type passed by value
+    let gen_ty = body.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
+
+    // Get the interior types and substs which typeck computed
+    let (upvars, interior, movable) = match *gen_ty.kind() {
+        ty::Generator(_, substs, movability) => {
+            let substs = substs.as_generator();
+            (
+                substs.upvar_tys().collect::<Vec<_>>(),
+                substs.witness(),
+                movability == hir::Movability::Movable,
+            )
+        }
+        _ => span_bug!(body.span, "unexpected generator type {}", gen_ty),
+    };
+
+    // When first entering the generator, move the resume argument into its new local.
+    let always_live_locals = always_storage_live_locals(&body);
+
+    let liveness_info = locals_live_across_suspend_points(tcx, body, &always_live_locals, movable);
+
+    // Extract locals which are live across suspension point into `layout`
+    // `remap` gives a mapping from local indices onto generator struct indices
+    // `storage_liveness` tells us which locals have live storage at suspension points
+    let (_, generator_layout, _) = compute_layout(liveness_info, body);
+
+    if tcx.sess.opts.unstable_opts.drop_tracking_mir {
+        check_suspend_tys(tcx, &generator_layout, &body);
+    } else {
+        sanitize_witness(tcx, body, interior, upvars, &generator_layout);
+    }
+
+    generator_layout
+}
+
 impl<'tcx> MirPass<'tcx> for StateTransform {
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         let Some(yield_ty) = body.yield_ty() else {
@@ -1372,16 +1424,11 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
         // The first argument is the generator type passed by value
         let gen_ty = body.local_decls.raw[1].ty;
 
-        // Get the interior types and substs which typeck computed
-        let (upvars, interior, discr_ty, movable) = match *gen_ty.kind() {
+        // Get the discriminant type and substs which typeck computed
+        let (discr_ty, movable) = match *gen_ty.kind() {
             ty::Generator(_, substs, movability) => {
                 let substs = substs.as_generator();
-                (
-                    substs.upvar_tys().collect(),
-                    substs.witness(),
-                    substs.discr_ty(tcx),
-                    movability == hir::Movability::Movable,
-                )
+                (substs.discr_ty(tcx), movability == hir::Movability::Movable)
             }
             _ => {
                 tcx.sess
@@ -1443,8 +1490,6 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
         let liveness_info =
             locals_live_across_suspend_points(tcx, body, &always_live_locals, movable);
 
-        sanitize_witness(tcx, body, interior, upvars, &liveness_info.saved_locals);
-
         if tcx.sess.opts.unstable_opts.validate_mir {
             let mut vis = EnsureGeneratorFieldAssignmentsNeverAlias {
                 assigned_local: None,
@@ -1640,3 +1685,212 @@ impl<'tcx> Visitor<'tcx> for EnsureGeneratorFieldAssignmentsNeverAlias<'_> {
         }
     }
 }
+
+fn check_suspend_tys<'tcx>(tcx: TyCtxt<'tcx>, layout: &GeneratorLayout<'tcx>, body: &Body<'tcx>) {
+    let mut linted_tys = FxHashSet::default();
+
+    // We want a user-facing param-env.
+    let param_env = tcx.param_env(body.source.def_id());
+
+    for (variant, yield_source_info) in
+        layout.variant_fields.iter().zip(&layout.variant_source_info)
+    {
+        debug!(?variant);
+        for &local in variant {
+            let decl = &layout.field_tys[local];
+            debug!(?decl);
+
+            if !decl.is_static_ptr && linted_tys.insert(decl.ty) {
+                let Some(hir_id) = decl.source_info.scope.lint_root(&body.source_scopes) else { continue };
+
+                check_must_not_suspend_ty(
+                    tcx,
+                    decl.ty,
+                    hir_id,
+                    param_env,
+                    SuspendCheckData {
+                        source_span: decl.source_info.span,
+                        yield_span: yield_source_info.span,
+                        plural_len: 1,
+                        ..Default::default()
+                    },
+                );
+            }
+        }
+    }
+}
+
+#[derive(Default)]
+struct SuspendCheckData<'a> {
+    source_span: Span,
+    yield_span: Span,
+    descr_pre: &'a str,
+    descr_post: &'a str,
+    plural_len: usize,
+}
+
+// Returns whether it emitted a diagnostic or not
+// Note that this fn and the proceeding one are based on the code
+// for creating must_use diagnostics
+//
+// Note that this technique was chosen over things like a `Suspend` marker trait
+// as it is simpler and has precedent in the compiler
+fn check_must_not_suspend_ty<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    ty: Ty<'tcx>,
+    hir_id: hir::HirId,
+    param_env: ty::ParamEnv<'tcx>,
+    data: SuspendCheckData<'_>,
+) -> bool {
+    if ty.is_unit() {
+        return false;
+    }
+
+    let plural_suffix = pluralize!(data.plural_len);
+
+    debug!("Checking must_not_suspend for {}", ty);
+
+    match *ty.kind() {
+        ty::Adt(..) if ty.is_box() => {
+            let boxed_ty = ty.boxed_ty();
+            let descr_pre = &format!("{}boxed ", data.descr_pre);
+            check_must_not_suspend_ty(
+                tcx,
+                boxed_ty,
+                hir_id,
+                param_env,
+                SuspendCheckData { descr_pre, ..data },
+            )
+        }
+        ty::Adt(def, _) => check_must_not_suspend_def(tcx, def.did(), hir_id, data),
+        // FIXME: support adding the attribute to TAITs
+        ty::Alias(ty::Opaque, ty::AliasTy { def_id: def, .. }) => {
+            let mut has_emitted = false;
+            for &(predicate, _) in tcx.explicit_item_bounds(def) {
+                // We only look at the `DefId`, so it is safe to skip the binder here.
+                if let ty::PredicateKind::Clause(ty::Clause::Trait(ref poly_trait_predicate)) =
+                    predicate.kind().skip_binder()
+                {
+                    let def_id = poly_trait_predicate.trait_ref.def_id;
+                    let descr_pre = &format!("{}implementer{} of ", data.descr_pre, plural_suffix);
+                    if check_must_not_suspend_def(
+                        tcx,
+                        def_id,
+                        hir_id,
+                        SuspendCheckData { descr_pre, ..data },
+                    ) {
+                        has_emitted = true;
+                        break;
+                    }
+                }
+            }
+            has_emitted
+        }
+        ty::Dynamic(binder, _, _) => {
+            let mut has_emitted = false;
+            for predicate in binder.iter() {
+                if let ty::ExistentialPredicate::Trait(ref trait_ref) = predicate.skip_binder() {
+                    let def_id = trait_ref.def_id;
+                    let descr_post = &format!(" trait object{}{}", plural_suffix, data.descr_post);
+                    if check_must_not_suspend_def(
+                        tcx,
+                        def_id,
+                        hir_id,
+                        SuspendCheckData { descr_post, ..data },
+                    ) {
+                        has_emitted = true;
+                        break;
+                    }
+                }
+            }
+            has_emitted
+        }
+        ty::Tuple(fields) => {
+            let mut has_emitted = false;
+            for (i, ty) in fields.iter().enumerate() {
+                let descr_post = &format!(" in tuple element {i}");
+                if check_must_not_suspend_ty(
+                    tcx,
+                    ty,
+                    hir_id,
+                    param_env,
+                    SuspendCheckData { descr_post, ..data },
+                ) {
+                    has_emitted = true;
+                }
+            }
+            has_emitted
+        }
+        ty::Array(ty, len) => {
+            let descr_pre = &format!("{}array{} of ", data.descr_pre, plural_suffix);
+            check_must_not_suspend_ty(
+                tcx,
+                ty,
+                hir_id,
+                param_env,
+                SuspendCheckData {
+                    descr_pre,
+                    plural_len: len.try_eval_usize(tcx, param_env).unwrap_or(0) as usize + 1,
+                    ..data
+                },
+            )
+        }
+        // If drop tracking is enabled, we want to look through references, since the referrent
+        // may not be considered live across the await point.
+        ty::Ref(_region, ty, _mutability) => {
+            let descr_pre = &format!("{}reference{} to ", data.descr_pre, plural_suffix);
+            check_must_not_suspend_ty(
+                tcx,
+                ty,
+                hir_id,
+                param_env,
+                SuspendCheckData { descr_pre, ..data },
+            )
+        }
+        _ => false,
+    }
+}
+
+fn check_must_not_suspend_def(
+    tcx: TyCtxt<'_>,
+    def_id: DefId,
+    hir_id: hir::HirId,
+    data: SuspendCheckData<'_>,
+) -> bool {
+    if let Some(attr) = tcx.get_attr(def_id, sym::must_not_suspend) {
+        let msg = format!(
+            "{}`{}`{} held across a suspend point, but should not be",
+            data.descr_pre,
+            tcx.def_path_str(def_id),
+            data.descr_post,
+        );
+        tcx.struct_span_lint_hir(
+            rustc_session::lint::builtin::MUST_NOT_SUSPEND,
+            hir_id,
+            data.source_span,
+            msg,
+            |lint| {
+                // add span pointing to the offending yield/await
+                lint.span_label(data.yield_span, "the value is held across this suspend point");
+
+                // Add optional reason note
+                if let Some(note) = attr.value_str() {
+                    // FIXME(guswynn): consider formatting this better
+                    lint.span_note(data.source_span, note.as_str());
+                }
+
+                // Add some quick suggestions on what to do
+                // FIXME: can `drop` work as a suggestion here as well?
+                lint.span_help(
+                    data.source_span,
+                    "consider using a block (`{ ... }`) \
+                    to shrink the value's scope, ending before the suspend point",
+                )
+            },
+        );
+
+        true
+    } else {
+        false
+    }
+}

compiler/rustc_mir_transform/src/lib.rs

@@ -123,6 +123,7 @@ pub fn provide(providers: &mut Providers) {
         mir_drops_elaborated_and_const_checked,
         mir_for_ctfe,
         mir_for_ctfe_of_const_arg,
+        mir_generator_witnesses: generator::mir_generator_witnesses,
         optimized_mir,
         is_mir_available,
         is_ctfe_mir_available: |tcx, did| is_mir_available(tcx, did),
@@ -425,6 +426,9 @@ fn mir_drops_elaborated_and_const_checked(
         return tcx.mir_drops_elaborated_and_const_checked(def);
     }
 
+    if tcx.generator_kind(def.did).is_some() {
+        tcx.ensure().mir_generator_witnesses(def.did);
+    }
     let mir_borrowck = tcx.mir_borrowck_opt_const_arg(def);
 
     let is_fn_like = tcx.def_kind(def.did).is_fn_like();

compiler/rustc_type_ir/src/sty.rs

@@ -170,7 +170,7 @@ pub enum TyKind<I: Interner> {
     ///
     /// This variant is only using when `drop_tracking_mir` is set.
     /// This contains the `DefId` and the `SubstRef` of the generator.
-    /// The actual witness types are computed on MIR by the `mir_generator_info` query.
+    /// The actual witness types are computed on MIR by the `mir_generator_witnesses` query.
     ///
     /// Looking at the following example, the witness for this generator
     /// may end up as something like `for<'a> [Vec<i32>, &'a Vec<i32>]`: