Rollup merge of #130201 - compiler-errors:foreign-synthetic-body, r=lcnr

Encode `coroutine_by_move_body_def_id` in crate metadata

We synthesize the MIR for a by-move body for the `FnOnce` implementation of async closures. It can be accessed with the `coroutine_by_move_body_def_id` query. We weren't encoding this query in the metadata though, nor were we properly recording that synthetic MIR in `mir_keys`, so the `optimized_mir` wasn't getting encoded either!

Stacked on top is a fix to consider `DefKind::SyntheticCoroutineBody` to return true in several places I missed. Specifically, we should consider the def-kind in `fn DefKind::is_fn_like()`, since that's what we were using to make sure we ensure `query mir_inliner_callees` before the MIR gets stolen for the body. This led to some CI failures that were caught by miri but which I added a test for.

9 files changed, 81 insertions, 20 deletions

diff --git a/compiler/rustc_hir/src/def.rs b/compiler/rustc_hir/src/def.rs
index bd55617d84e..9b4174013a6 100644
--- a/compiler/rustc_hir/src/def.rs
+++ b/compiler/rustc_hir/src/def.rs
@@ -287,7 +287,10 @@ impl DefKind {
 
     #[inline]
     pub fn is_fn_like(self) -> bool {
-        matches!(self, DefKind::Fn | DefKind::AssocFn | DefKind::Closure)
+        matches!(
+            self,
+            DefKind::Fn | DefKind::AssocFn | DefKind::Closure | DefKind::SyntheticCoroutineBody
+        )
     }
 
     /// Whether `query get_codegen_attrs` should be used with this definition.
diff --git a/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs b/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
index 53da07aeaa6..b9d287730fa 100644
--- a/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
+++ b/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
@@ -290,6 +290,7 @@ provide! { tcx, def_id, other, cdata,
     fn_arg_names => { table }
     coroutine_kind => { table_direct }
     coroutine_for_closure => { table }
+    coroutine_by_move_body_def_id => { table }
     eval_static_initializer => {
         Ok(cdata
             .root
diff --git a/compiler/rustc_metadata/src/rmeta/encoder.rs b/compiler/rustc_metadata/src/rmeta/encoder.rs
index b617d5236b9..46dc0696638 100644
--- a/compiler/rustc_metadata/src/rmeta/encoder.rs
+++ b/compiler/rustc_metadata/src/rmeta/encoder.rs
@@ -1488,9 +1488,18 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
             if def_kind == DefKind::Closure
                 && tcx.type_of(def_id).skip_binder().is_coroutine_closure()
             {
+                let coroutine_for_closure = self.tcx.coroutine_for_closure(def_id);
                 self.tables
                     .coroutine_for_closure
-                    .set_some(def_id.index, self.tcx.coroutine_for_closure(def_id).into());
+                    .set_some(def_id.index, coroutine_for_closure.into());
+
+                // If this async closure has a by-move body, record it too.
+                if tcx.needs_coroutine_by_move_body_def_id(coroutine_for_closure) {
+                    self.tables.coroutine_by_move_body_def_id.set_some(
+                        coroutine_for_closure.index,
+                        self.tcx.coroutine_by_move_body_def_id(coroutine_for_closure).into(),
+                    );
+                }
             }
             if let DefKind::Static { .. } = def_kind {
                 if !self.tcx.is_foreign_item(def_id) {
diff --git a/compiler/rustc_metadata/src/rmeta/mod.rs b/compiler/rustc_metadata/src/rmeta/mod.rs
index 8180a507a51..c791a1f541c 100644
--- a/compiler/rustc_metadata/src/rmeta/mod.rs
+++ b/compiler/rustc_metadata/src/rmeta/mod.rs
@@ -446,6 +446,7 @@ define_tables! {
     fn_arg_names: Table<DefIndex, LazyArray<Ident>>,
     coroutine_kind: Table<DefIndex, hir::CoroutineKind>,
     coroutine_for_closure: Table<DefIndex, RawDefId>,
+    coroutine_by_move_body_def_id: Table<DefIndex, RawDefId>,
     eval_static_initializer: Table<DefIndex, LazyValue<mir::interpret::ConstAllocation<'static>>>,
     trait_def: Table<DefIndex, LazyValue<ty::TraitDef>>,
     trait_item_def_id: Table<DefIndex, RawDefId>,
diff --git a/compiler/rustc_mir_transform/src/cross_crate_inline.rs b/compiler/rustc_mir_transform/src/cross_crate_inline.rs
index ce109ef7674..42cbece32d8 100644
--- a/compiler/rustc_mir_transform/src/cross_crate_inline.rs
+++ b/compiler/rustc_mir_transform/src/cross_crate_inline.rs
@@ -24,7 +24,7 @@ fn cross_crate_inlinable(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
 
     // This just reproduces the logic from Instance::requires_inline.
     match tcx.def_kind(def_id) {
-        DefKind::Ctor(..) | DefKind::Closure => return true,
+        DefKind::Ctor(..) | DefKind::Closure | DefKind::SyntheticCoroutineBody => return true,
         DefKind::Fn | DefKind::AssocFn => {}
         _ => return false,
     }
diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs
index 424e7008326..0868f4b3d88 100644
--- a/compiler/rustc_mir_transform/src/lib.rs
+++ b/compiler/rustc_mir_transform/src/lib.rs
@@ -21,9 +21,8 @@ use rustc_const_eval::util;
 use rustc_data_structures::fx::FxIndexSet;
 use rustc_data_structures::steal::Steal;
 use rustc_hir as hir;
-use rustc_hir::def::DefKind;
+use rustc_hir::def::{CtorKind, DefKind};
 use rustc_hir::def_id::LocalDefId;
-use rustc_hir::intravisit::{self, Visitor};
 use rustc_index::IndexVec;
 use rustc_middle::mir::{
     AnalysisPhase, Body, CallSource, ClearCrossCrate, ConstOperand, ConstQualifs, LocalDecl,
@@ -224,26 +223,31 @@ fn is_mir_available(tcx: TyCtxt<'_>, def_id: LocalDefId) -> bool {
 /// MIR associated with them.
 fn mir_keys(tcx: TyCtxt<'_>, (): ()) -> FxIndexSet<LocalDefId> {
     // All body-owners have MIR associated with them.
-    let set: FxIndexSet<_> = tcx.hir().body_owners().collect();
+    let mut set: FxIndexSet<_> = tcx.hir().body_owners().collect();
 
-    // Additionally, tuple struct/variant constructors have MIR, but
-    // they don't have a BodyId, so we need to build them separately.
-    struct GatherCtors {
-        set: FxIndexSet<LocalDefId>,
+    // Coroutine-closures (e.g. async closures) have an additional by-move MIR
+    // body that isn't in the HIR.
+    for body_owner in tcx.hir().body_owners() {
+        if let DefKind::Closure = tcx.def_kind(body_owner)
+            && tcx.needs_coroutine_by_move_body_def_id(body_owner.to_def_id())
+        {
+            set.insert(tcx.coroutine_by_move_body_def_id(body_owner).expect_local());
+        }
     }
-    impl<'tcx> Visitor<'tcx> for GatherCtors {
-        fn visit_variant_data(&mut self, v: &'tcx hir::VariantData<'tcx>) {
-            if let hir::VariantData::Tuple(_, _, def_id) = *v {
-                self.set.insert(def_id);
+
+    // tuple struct/variant constructors have MIR, but they don't have a BodyId,
+    // so we need to build them separately.
+    for item in tcx.hir_crate_items(()).free_items() {
+        if let DefKind::Struct | DefKind::Enum = tcx.def_kind(item.owner_id) {
+            for variant in tcx.adt_def(item.owner_id).variants() {
+                if let Some((CtorKind::Fn, ctor_def_id)) = variant.ctor {
+                    set.insert(ctor_def_id.expect_local());
+                }
             }
-            intravisit::walk_struct_def(self, v)
         }
     }
 
-    let mut gather_ctors = GatherCtors { set };
-    tcx.hir().visit_all_item_likes_in_crate(&mut gather_ctors);
-
-    gather_ctors.set
+    set
 }
 
 fn mir_const_qualif(tcx: TyCtxt<'_>, def: LocalDefId) -> ConstQualifs {
diff --git a/compiler/rustc_symbol_mangling/src/lib.rs b/compiler/rustc_symbol_mangling/src/lib.rs
index 0c97cda81c8..78e6b9ec6e8 100644
--- a/compiler/rustc_symbol_mangling/src/lib.rs
+++ b/compiler/rustc_symbol_mangling/src/lib.rs
@@ -227,7 +227,11 @@ fn compute_symbol_name<'tcx>(
     // and we want to be sure to avoid any symbol conflicts here.
     let is_globally_shared_function = matches!(
         tcx.def_kind(instance.def_id()),
-        DefKind::Fn | DefKind::AssocFn | DefKind::Closure | DefKind::Ctor(..)
+        DefKind::Fn
+            | DefKind::AssocFn
+            | DefKind::Closure
+            | DefKind::SyntheticCoroutineBody
+            | DefKind::Ctor(..)
     ) && matches!(
         MonoItem::Fn(instance).instantiation_mode(tcx),
         InstantiationMode::GloballyShared { may_conflict: true }
diff --git a/tests/ui/async-await/async-closures/foreign.rs b/tests/ui/async-await/async-closures/foreign.rs
index 50bef9cf11d..ab6fe06a3f4 100644
--- a/tests/ui/async-await/async-closures/foreign.rs
+++ b/tests/ui/async-await/async-closures/foreign.rs
@@ -12,8 +12,13 @@ extern crate foreign;
 
 struct NoCopy;
 
+async fn call_once(f: impl async FnOnce()) {
+    f().await;
+}
+
 fn main() {
     block_on::block_on(async {
         foreign::closure()().await;
+        call_once(foreign::closure()).await;
     });
 }
diff --git a/tests/ui/async-await/async-closures/inline-body.rs b/tests/ui/async-await/async-closures/inline-body.rs
new file mode 100644
index 00000000000..a842d98d1de
--- /dev/null
+++ b/tests/ui/async-await/async-closures/inline-body.rs
@@ -0,0 +1,34 @@
+//@ edition: 2021
+//@ compile-flags: -Zinline-mir
+//@ build-pass
+
+// Ensure that we don't hit a Steal ICE because we forgot to ensure
+// `mir_inliner_callees` for the synthetic by-move coroutine body since
+// its def-id wasn't previously being considered.
+
+#![feature(async_closure, noop_waker)]
+
+use std::future::Future;
+use std::pin::pin;
+use std::task::*;
+
+pub fn block_on<T>(fut: impl Future<Output = T>) -> T {
+    let mut fut = pin!(fut);
+    let ctx = &mut Context::from_waker(Waker::noop());
+
+    loop {
+        match fut.as_mut().poll(ctx) {
+            Poll::Pending => {}
+            Poll::Ready(t) => break t,
+        }
+    }
+}
+
+async fn call_once<T>(f: impl async FnOnce() -> T) -> T {
+    f().await
+}
+
+fn main() {
+    let c = async || {};
+    block_on(call_once(c));
+}