Simplify transformed body.

Author: cjgillot

compiler/rustc_mir_transform/src/coroutine/mod.rs

@@ -868,12 +868,6 @@ fn create_coroutine_resume_function<'tcx>(
         }
     }
 
-    // Make sure we remove dead blocks to remove
-    // unrelated code from the drop part of the function
-    simplify::remove_dead_blocks(body);
-
-    pm::run_passes_no_validate(tcx, body, &[&abort_unwinding_calls::AbortUnwindingCalls], None);
-
     // Run derefer to fix Derefs that are not in the first place
     deref_finder(tcx, body, false);
 
@@ -884,6 +878,17 @@ fn create_coroutine_resume_function<'tcx>(
     if let Some(dumper) = MirDumper::new(tcx, "coroutine_resume", body) {
         dumper.dump_mir(body);
     }
+
+    pm::run_passes_no_validate(
+        tcx,
+        body,
+        &[
+            &crate::abort_unwinding_calls::AbortUnwindingCalls,
+            &crate::simplify::SimplifyCfg::PostStateTransform,
+            &crate::simplify::SimplifyLocals::PostStateTransform,
+        ],
+        None,
+    );
 }
 
 /// An operation that can be performed on a coroutine.

compiler/rustc_mir_transform/src/lib.rs

@@ -79,7 +79,7 @@ pub use shim::build_drop_shim;
 macro_rules! declare_passes {
     (
         $(
-            $vis:vis mod $mod_name:ident : $($pass_name:ident $( { $($ident:ident),* } )?),+ $(,)?;
+            $vis:vis mod $mod_name:ident : $($pass_name:ident $( { $($ident:ident),* $(,)? } )?),+ $(,)?;
         )*
     ) => {
         $(
@@ -186,12 +186,14 @@ declare_passes! {
             PreOptimizations,
             Final,
             MakeShim,
-            AfterUnreachableEnumBranching
+            AfterUnreachableEnumBranching,
+            PostStateTransform,
         },
         SimplifyLocals {
             BeforeConstProp,
             AfterGVN,
-            Final
+            Final,
+            PostStateTransform,
         };
     mod simplify_branches : SimplifyConstCondition {
         AfterInstSimplify,

compiler/rustc_mir_transform/src/simplify.rs

@@ -57,6 +57,8 @@ pub(super) enum SimplifyCfg {
     Final,
     MakeShim,
     AfterUnreachableEnumBranching,
+    /// Extra run introduced by `StateTransform`.
+    PostStateTransform,
 }
 
 impl SimplifyCfg {
@@ -72,6 +74,7 @@ impl SimplifyCfg {
             SimplifyCfg::AfterUnreachableEnumBranching => {
                 "SimplifyCfg-after-unreachable-enum-branching"
             }
+            SimplifyCfg::PostStateTransform => "SimplifyCfg-post-StateTransform",
         }
     }
 }
@@ -416,6 +419,8 @@ pub(super) enum SimplifyLocals {
     BeforeConstProp,
     AfterGVN,
     Final,
+    /// Extra run introduced by `StateTransform`.
+    PostStateTransform,
 }
 
 impl<'tcx> crate::MirPass<'tcx> for SimplifyLocals {
@@ -424,6 +429,7 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyLocals {
             SimplifyLocals::BeforeConstProp => "SimplifyLocals-before-const-prop",
             SimplifyLocals::AfterGVN => "SimplifyLocals-after-value-numbering",
             SimplifyLocals::Final => "SimplifyLocals-final",
+            SimplifyLocals::PostStateTransform => "SimplifyLocals-post-StateTransform",
         }
     }
 

tests/mir-opt/coroutine/async_await.a-{closure#0}.StateTransform.diff

@@ -30,32 +30,25 @@
 +         _6 = std::future::ResumeTy(move _7);
 +         _5 = copy (_1.0: &mut {async fn body of a()});
 +         _4 = discriminant((*_5));
-+         switchInt(move _4) -> [0: bb5, 1: bb3, otherwise: bb4];
++         switchInt(move _4) -> [0: bb3, 1: bb1, otherwise: bb2];
       }
 
       bb1: {
-+         _0 = Poll::<()>::Ready(move _3);
-+         discriminant((*_5)) = 1;
-          return;
+-         return;
++         assert(const false, "`async fn` resumed after completion") -> [success: bb1, unwind continue];
       }
 
 -     bb2 (cleanup): {
 -         resume;
 +     bb2: {
-+         goto -> bb1;
-+     }
-+ 
-+     bb3: {
-+         assert(const false, "`async fn` resumed after completion") -> [success: bb3, unwind continue];
-+     }
-+ 
-+     bb4: {
 +         unreachable;
 +     }
 + 
-+     bb5: {
++     bb3: {
 +         _3 = const ();
-+         goto -> bb2;
++         _0 = Poll::<()>::Ready(move _3);
++         discriminant((*_5)) = 1;
++         return;
       }
   }