2 files changed, 18 insertions, 17 deletions
diff --git a/compiler/rustc_mir_transform/src/generator.rs b/compiler/rustc_mir_transform/src/generator.rs
index e8871ff37f2..a9fd95f4541 100644
--- a/compiler/rustc_mir_transform/src/generator.rs
+++ b/compiler/rustc_mir_transform/src/generator.rs
@@ -1390,6 +1390,7 @@ pub(crate) fn mir_generator_witnesses<'tcx>(
     tcx: TyCtxt<'tcx>,
     def_id: DefId,
 ) -> GeneratorLayout<'tcx> {
+    assert!(tcx.sess.opts.unstable_opts.drop_tracking_mir);
     let def_id = def_id.expect_local();
 
     let (body, _) = tcx.mir_promoted(ty::WithOptConstParam::unknown(def_id));
@@ -1400,15 +1401,8 @@ pub(crate) fn mir_generator_witnesses<'tcx>(
     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,
-            )
-        }
+    let movable = match *gen_ty.kind() {
+        ty::Generator(_, _, movability) => movability == hir::Movability::Movable,
         _ => span_bug!(body.span, "unexpected generator type {}", gen_ty),
     };
 
@@ -1422,11 +1416,7 @@ pub(crate) fn mir_generator_witnesses<'tcx>(
     // `storage_liveness` tells us which locals have live storage at suspension points
     let (_, generator_layout, _) = compute_layout(tcx, 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);
-    }
+    check_suspend_tys(tcx, &generator_layout, &body);
 
     generator_layout
 }
@@ -1444,10 +1434,15 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
         let gen_ty = body.local_decls.raw[1].ty;
 
         // Get the discriminant type and substs which typeck computed
-        let (discr_ty, movable) = match *gen_ty.kind() {
+        let (discr_ty, upvars, interior, movable) = match *gen_ty.kind() {
             ty::Generator(_, substs, movability) => {
                 let substs = substs.as_generator();
-                (substs.discr_ty(tcx), movability == hir::Movability::Movable)
+                (
+                    substs.discr_ty(tcx),
+                    substs.upvar_tys().collect::<Vec<_>>(),
+                    substs.witness(),
+                    movability == hir::Movability::Movable,
+                )
             }
             _ => {
                 tcx.sess
@@ -1524,6 +1519,12 @@ impl<'tcx> MirPass<'tcx> for StateTransform {
         // `storage_liveness` tells us which locals have live storage at suspension points
         let (remap, layout, storage_liveness) = compute_layout(tcx, liveness_info, body);
 
+        if tcx.sess.opts.unstable_opts.validate_mir
+            && !tcx.sess.opts.unstable_opts.drop_tracking_mir
+        {
+            sanitize_witness(tcx, body, interior, upvars, &layout);
+        }
+
         let can_return = can_return(tcx, body, tcx.param_env(body.source.def_id()));
 
         // Run the transformation which converts Places from Local to generator struct
diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs
index fe3d5b1cce4..fcb09fa02dd 100644
--- a/compiler/rustc_mir_transform/src/lib.rs
+++ b/compiler/rustc_mir_transform/src/lib.rs
@@ -426,7 +426,7 @@ fn mir_drops_elaborated_and_const_checked(
         return tcx.mir_drops_elaborated_and_const_checked(def);
     }
 
-    if tcx.generator_kind(def.did).is_some() {
+    if tcx.generator_kind(def.did).is_some() && tcx.sess.opts.unstable_opts.drop_tracking_mir {
         tcx.ensure().mir_generator_witnesses(def.did);
     }
     let mir_borrowck = tcx.mir_borrowck_opt_const_arg(def);