Fix validation on substituted callee bodies in MIR inliner

compiler/rustc_const_eval/src/transform/validate.rs

@@ -85,7 +85,7 @@ impl<'tcx> MirPass<'tcx> for Validator {
         cfg_checker.check_cleanup_control_flow();
 
         // Also run the TypeChecker.
-        for (location, msg) in validate_types(tcx, self.mir_phase, param_env, body) {
+        for (location, msg) in validate_types(tcx, self.mir_phase, param_env, body, body) {
             cfg_checker.fail(location, msg);
         }
 
@@ -541,19 +541,25 @@ impl<'a, 'tcx> Visitor<'tcx> for CfgChecker<'a, 'tcx> {
 
 /// A faster version of the validation pass that only checks those things which may break when
 /// instantiating any generic parameters.
+///
+/// `caller_body` is used to detect cycles in MIR inlining and MIR validation before
+/// `optimized_mir` is available.
 pub fn validate_types<'tcx>(
     tcx: TyCtxt<'tcx>,
     mir_phase: MirPhase,
     param_env: ty::ParamEnv<'tcx>,
     body: &Body<'tcx>,
+    caller_body: &Body<'tcx>,
 ) -> Vec<(Location, String)> {
-    let mut type_checker = TypeChecker { body, tcx, param_env, mir_phase, failures: Vec::new() };
+    let mut type_checker =
+        TypeChecker { body, caller_body, tcx, param_env, mir_phase, failures: Vec::new() };
     type_checker.visit_body(body);
     type_checker.failures
 }
 
 struct TypeChecker<'a, 'tcx> {
     body: &'a Body<'tcx>,
+    caller_body: &'a Body<'tcx>,
     tcx: TyCtxt<'tcx>,
     param_env: ParamEnv<'tcx>,
     mir_phase: MirPhase,
@@ -705,8 +711,13 @@ impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
                     }
                     &ty::Coroutine(def_id, args) => {
                         let f_ty = if let Some(var) = parent_ty.variant_index {
-                            let gen_body = if def_id == self.body.source.def_id() {
-                                self.body
+                            // If we're currently validating an inlined copy of this body,
+                            // then it will no longer be parameterized over the original
+                            // args of the coroutine. Otherwise, we prefer to use this body
+                            // since we may be in the process of computing this MIR in the
+                            // first place.
+                            let gen_body = if def_id == self.caller_body.source.def_id() {
+                                self.caller_body
                             } else {
                                 self.tcx.optimized_mir(def_id)
                             };

compiler/rustc_mir_transform/src/coroutine/by_move_body.rs

@@ -64,12 +64,9 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
         let mut by_move_body = body.clone();
         MakeByMoveBody { tcx, by_ref_fields, by_move_coroutine_ty }.visit_body(&mut by_move_body);
         dump_mir(tcx, false, "coroutine_by_move", &0, &by_move_body, |_, _| Ok(()));
-        by_move_body.source = mir::MirSource {
-            instance: InstanceDef::CoroutineKindShim {
-                coroutine_def_id: coroutine_def_id.to_def_id(),
-            },
-            promoted: None,
-        };
+        by_move_body.source = mir::MirSource::from_instance(InstanceDef::CoroutineKindShim {
+            coroutine_def_id: coroutine_def_id.to_def_id(),
+        });
         body.coroutine.as_mut().unwrap().by_move_body = Some(by_move_body);
     }
 }

compiler/rustc_mir_transform/src/inline.rs

@@ -213,6 +213,7 @@ impl<'tcx> Inliner<'tcx> {
             MirPhase::Runtime(RuntimePhase::Optimized),
             self.param_env,
             &callee_body,
+            &caller_body,
         )
         .is_empty()
         {

tests/mir-opt/inline_coroutine_body.rs

@@ -0,0 +1,28 @@
+// EMIT_MIR_FOR_EACH_PANIC_STRATEGY
+// skip-filecheck
+//@ unit-test: Inline
+//@ edition: 2021
+//@ compile-flags: -Zinline-mir-hint-threshold=10000 -Zinline-mir-threshold=10000 --crate-type=lib
+
+pub async fn run(permit: ActionPermit<'_, ()>, ctx: &mut core::task::Context<'_>) {
+    run2(permit, ctx);
+}
+
+// EMIT_MIR inline_coroutine_body.run2-{closure#0}.Inline.diff
+fn run2<T>(permit: ActionPermit<'_, T>, ctx: &mut core::task::Context) {
+    _ = || {
+        let mut fut = ActionPermit::perform(permit);
+        let fut = unsafe { core::pin::Pin::new_unchecked(&mut fut) };
+        _ = core::future::Future::poll(fut, ctx);
+    };
+}
+
+pub struct ActionPermit<'a, T> {
+    _guard: core::cell::Ref<'a, T>,
+}
+
+impl<'a, T> ActionPermit<'a, T> {
+    async fn perform(self) {
+        core::future::ready(()).await
+    }
+}

tests/mir-opt/inline_coroutine_body.run2-{closure#0}.Inline.panic-abort.diff

@@ -0,0 +1,281 @@
+- // MIR for `run2::{closure#0}` before Inline
++ // MIR for `run2::{closure#0}` after Inline
+
+  fn run2::{closure#0}(_1: {closure@$DIR/inline_coroutine_body.rs:13:9: 13:11}) -> () {
+      debug permit => (_1.0: ActionPermit<'_, T>);
+      debug ctx => (*(_1.1: &mut std::task::Context<'_>));
+      let mut _0: ();
+      let mut _2: {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6};
+      let mut _3: ActionPermit<'_, T>;
+      let mut _5: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6};
+      let _6: ();
+      let mut _7: std::task::Poll<()>;
+      let mut _8: std::pin::Pin<&mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6}>;
+      let mut _9: &mut std::task::Context<'_>;
+      let mut _10: &mut std::task::Context<'_>;
+      scope 1 {
+          debug fut => _2;
+          let _4: std::pin::Pin<&mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6}>;
+          scope 2 {
+              debug fut => _4;
+              scope 4 {
+              }
++             scope 7 (inlined ActionPermit::<'_, T>::perform::{closure#0}) {
++                 debug _task_context => _31;
++                 debug self => ((*(_8.0: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6})).0: ActionPermit<'_, T>);
++                 let _11: ActionPermit<'_, T>;
++                 let mut _12: std::future::Ready<()>;
++                 let mut _13: std::future::Ready<()>;
++                 let mut _14: ();
++                 let mut _16: ();
++                 let _17: ();
++                 let mut _18: std::task::Poll<()>;
++                 let mut _19: std::pin::Pin<&mut std::future::Ready<()>>;
++                 let mut _20: &mut std::future::Ready<()>;
++                 let mut _21: &mut std::future::Ready<()>;
++                 let mut _22: &mut std::task::Context<'_>;
++                 let mut _23: &mut std::task::Context<'_>;
++                 let mut _24: &mut std::task::Context<'_>;
++                 let mut _25: isize;
++                 let mut _27: !;
++                 let mut _28: &mut std::task::Context<'_>;
++                 let mut _29: ();
++                 let mut _30: ();
++                 let mut _31: &mut std::task::Context<'_>;
++                 let mut _32: u32;
++                 let mut _33: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6};
++                 let mut _34: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6};
++                 let mut _35: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6};
++                 let mut _36: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6};
++                 let mut _37: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6};
++                 let mut _38: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6};
++                 let mut _39: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6};
++                 let mut _40: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6};
++                 scope 8 {
++                     debug self => (((*(_8.0: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6})) as variant#3).0: ActionPermit<'_, T>);
++                     let mut _15: std::future::Ready<()>;
++                     scope 9 {
++                         debug __awaitee => (((*(_8.0: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6})) as variant#3).1: std::future::Ready<()>);
++                         let _26: ();
++                         scope 10 {
++                         }
++                         scope 11 {
++                             debug result => _26;
++                         }
++                     }
++                     scope 12 (inlined ready::<()>) {
++                         debug t => _14;
++                         let mut _41: std::option::Option<()>;
++                     }
++                 }
++             }
+          }
+          scope 3 {
++             scope 6 (inlined Pin::<&mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6}>::new_unchecked) {
++                 debug pointer => _5;
++             }
+          }
+      }
++     scope 5 (inlined ActionPermit::<'_, T>::perform) {
++         debug self => _3;
++     }
+
+      bb0: {
+          StorageLive(_2);
+          StorageLive(_3);
+          _3 = move (_1.0: ActionPermit<'_, T>);
+-         _2 = ActionPermit::<'_, T>::perform(move _3) -> [return: bb1, unwind unreachable];
+-     }
+- 
+-     bb1: {
++         _2 = {coroutine@$DIR/inline_coroutine_body.rs:25:28: 27:6 (#0)} { self: move _3 };
+          StorageDead(_3);
+          StorageLive(_4);
+          StorageLive(_5);
+          _5 = &mut _2;
+-         _4 = Pin::<&mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6}>::new_unchecked(move _5) -> [return: bb2, unwind unreachable];
+-     }
+- 
+-     bb2: {
++         _4 = Pin::<&mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6}> { __pointer: _5 };
+          StorageDead(_5);
+          StorageLive(_6);
+          StorageLive(_7);
+          StorageLive(_8);
+          _8 = move _4;
+          StorageLive(_9);
+          _10 = deref_copy (_1.1: &mut std::task::Context<'_>);
+          _9 = &mut (*_10);
+-         _7 = <{async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6} as Future>::poll(move _8, move _9) -> [return: bb3, unwind unreachable];
++         StorageLive(_11);
++         StorageLive(_15);
++         StorageLive(_16);
++         StorageLive(_25);
++         StorageLive(_27);
++         StorageLive(_30);
++         StorageLive(_31);
++         StorageLive(_32);
++         StorageLive(_33);
++         StorageLive(_34);
++         StorageLive(_35);
++         StorageLive(_36);
++         StorageLive(_37);
++         StorageLive(_38);
++         StorageLive(_39);
++         StorageLive(_40);
++         _33 = deref_copy (_8.0: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6});
++         _32 = discriminant((*_33));
++         switchInt(move _32) -> [0: bb3, 1: bb13, 3: bb12, otherwise: bb8];
+      }
+
+-     bb3: {
++     bb1: {
++         StorageDead(_2);
++         return;
++     }
++ 
++     bb2: {
++         StorageDead(_40);
++         StorageDead(_39);
++         StorageDead(_38);
++         StorageDead(_37);
++         StorageDead(_36);
++         StorageDead(_35);
++         StorageDead(_34);
++         StorageDead(_33);
++         StorageDead(_32);
++         StorageDead(_31);
++         StorageDead(_30);
++         StorageDead(_27);
++         StorageDead(_25);
++         StorageDead(_16);
++         StorageDead(_15);
++         StorageDead(_11);
+          StorageDead(_9);
+          StorageDead(_8);
+          StorageDead(_7);
+          _6 = const ();
+          StorageDead(_6);
+          _0 = const ();
+          StorageDead(_4);
+-         drop(_2) -> [return: bb4, unwind unreachable];
++         drop(_2) -> [return: bb1, unwind unreachable];
+      }
+
++     bb3: {
++         _31 = move _9;
++         _34 = deref_copy (_8.0: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6});
++         _35 = deref_copy (_8.0: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6});
++         (((*_34) as variant#3).0: ActionPermit<'_, T>) = move ((*_35).0: ActionPermit<'_, T>);
++         StorageLive(_12);
++         StorageLive(_13);
++         StorageLive(_14);
++         _14 = ();
++         StorageLive(_41);
++         _41 = Option::<()>::Some(_14);
++         _13 = std::future::Ready::<()>(move _41);
++         StorageDead(_41);
++         StorageDead(_14);
++         _12 = <std::future::Ready<()> as IntoFuture>::into_future(move _13) -> [return: bb4, unwind unreachable];
++     }
++ 
+      bb4: {
+-         StorageDead(_2);
+-         return;
++         StorageDead(_13);
++         _36 = deref_copy (_8.0: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6});
++         (((*_36) as variant#3).1: std::future::Ready<()>) = move _12;
++         goto -> bb5;
++     }
++ 
++     bb5: {
++         StorageLive(_17);
++         StorageLive(_18);
++         StorageLive(_19);
++         StorageLive(_20);
++         StorageLive(_21);
++         _37 = deref_copy (_8.0: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6});
++         _21 = &mut (((*_37) as variant#3).1: std::future::Ready<()>);
++         _20 = &mut (*_21);
++         _19 = Pin::<&mut std::future::Ready<()>>::new_unchecked(move _20) -> [return: bb6, unwind unreachable];
++     }
++ 
++     bb6: {
++         StorageDead(_20);
++         StorageLive(_22);
++         StorageLive(_23);
++         StorageLive(_24);
++         _24 = _31;
++         _23 = move _24;
++         _22 = &mut (*_23);
++         StorageDead(_24);
++         _18 = <std::future::Ready<()> as Future>::poll(move _19, move _22) -> [return: bb7, unwind unreachable];
++     }
++ 
++     bb7: {
++         StorageDead(_22);
++         StorageDead(_19);
++         _25 = discriminant(_18);
++         switchInt(move _25) -> [0: bb10, 1: bb9, otherwise: bb8];
++     }
++ 
++     bb8: {
++         unreachable;
++     }
++ 
++     bb9: {
++         _17 = const ();
++         StorageDead(_23);
++         StorageDead(_21);
++         StorageDead(_18);
++         StorageDead(_17);
++         StorageLive(_28);
++         StorageLive(_29);
++         _29 = ();
++         _7 = Poll::<()>::Pending;
++         StorageDead(_12);
++         StorageDead(_28);
++         StorageDead(_29);
++         _38 = deref_copy (_8.0: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6});
++         discriminant((*_38)) = 3;
++         goto -> bb2;
++     }
++ 
++     bb10: {
++         StorageLive(_26);
++         _26 = ((_18 as Ready).0: ());
++         _30 = _26;
++         StorageDead(_26);
++         StorageDead(_23);
++         StorageDead(_21);
++         StorageDead(_18);
++         StorageDead(_17);
++         StorageDead(_12);
++         _39 = deref_copy (_8.0: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6});
++         drop((((*_39) as variant#3).0: ActionPermit<'_, T>)) -> [return: bb11, unwind unreachable];
++     }
++ 
++     bb11: {
++         _7 = Poll::<()>::Ready(move _30);
++         _40 = deref_copy (_8.0: &mut {async fn body@$DIR/inline_coroutine_body.rs:25:28: 27:6});
++         discriminant((*_40)) = 1;
++         goto -> bb2;
++     }
++ 
++     bb12: {
++         StorageLive(_12);
++         StorageLive(_28);
++         StorageLive(_29);
++         _28 = move _9;
++         StorageDead(_29);
++         _31 = move _28;
++         StorageDead(_28);
++         _16 = const ();
++         goto -> bb5;
++     }
++ 
++     bb13: {
++         assert(const false, "`async fn` resumed after completion") -> [success: bb13, unwind unreachable];
+      }
+  }
+