Stop using a special inner body for the coroutine by-move body for async closures

Author: compiler-errors

compiler/rustc_const_eval/src/interpret/terminator.rs

@@ -608,7 +608,6 @@ impl<'tcx, M: Machine<'tcx>> InterpCx<'tcx, M> {
             | ty::InstanceKind::ReifyShim(..)
             | ty::InstanceKind::ClosureOnceShim { .. }
             | ty::InstanceKind::ConstructCoroutineInClosureShim { .. }
-            | ty::InstanceKind::CoroutineKindShim { .. }
             | ty::InstanceKind::FnPtrShim(..)
             | ty::InstanceKind::DropGlue(..)
             | ty::InstanceKind::CloneShim(..)

compiler/rustc_hir_analysis/src/lib.rs

@@ -199,10 +199,6 @@ pub fn check_crate(tcx: TyCtxt<'_>) {
         }
     });
 
-    // Freeze definitions as we don't add new ones at this point. This improves performance by
-    // allowing lock-free access to them.
-    tcx.untracked().definitions.freeze();
-
     // FIXME: Remove this when we implement creating `DefId`s
     // for anon constants during their parents' typeck.
     // Typeck all body owners in parallel will produce queries

compiler/rustc_interface/src/passes.rs

@@ -816,6 +816,13 @@ fn run_required_analyses(tcx: TyCtxt<'_>) {
             );
         }
     });
+    // Freeze definitions as we don't add new ones at this point.
+    // We need to wait until now since we synthesize a by-move body
+    // This improves performance by allowing lock-free access to them.
+    // FIXME(async_closures): We could force `coroutine_by_move_body_def_id`
+    // immediately after typeck, then freeze after that.
+    tcx.untracked().definitions.freeze();
+
     sess.time("layout_testing", || layout_test::test_layout(tcx));
     sess.time("abi_testing", || abi_test::test_abi(tcx));
 }

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -1366,7 +1366,11 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
                 let def_span = tcx.def_span(local_id);
                 record!(self.tables.def_span[def_id] <- def_span);
             }
-            if should_encode_attrs(def_kind) {
+            // FIXME(async_closures): We should just use `tcx.attrs` rather than going
+            // through the HIR. Historically, though, this has been inefficient apparently.
+            // For now, it's kind of pointless to fix, because coroutine-closures' coroutine
+            // bodies have no attrs anyways.
+            if should_encode_attrs(def_kind) && !tcx.is_synthetic_mir(local_id) {
                 self.encode_attrs(local_id);
             }
             if should_encode_expn_that_defined(def_kind) {

compiler/rustc_middle/src/mir/mod.rs

@@ -267,18 +267,6 @@ pub struct CoroutineInfo<'tcx> {
     /// Coroutine drop glue. This field is populated after the state transform pass.
     pub coroutine_drop: Option<Body<'tcx>>,
 
-    /// The body of the coroutine, modified to take its upvars by move rather than by ref.
-    ///
-    /// This is used by coroutine-closures, which must return a different flavor of coroutine
-    /// when called using `AsyncFnOnce::call_once`. It is produced by the `ByMoveBody` pass which
-    /// is run right after building the initial MIR, and will only be populated for coroutines
-    /// which come out of the async closure desugaring.
-    ///
-    /// This body should be processed in lockstep with the containing body -- any optimization
-    /// passes, etc, should be applied to this body as well. This is done automatically if
-    /// using `run_passes`.
-    pub by_move_body: Option<Body<'tcx>>,
-
     /// The layout of a coroutine. This field is populated after the state transform pass.
     pub coroutine_layout: Option<CoroutineLayout<'tcx>>,
 
@@ -298,7 +286,6 @@ impl<'tcx> CoroutineInfo<'tcx> {
             coroutine_kind,
             yield_ty: Some(yield_ty),
             resume_ty: Some(resume_ty),
-            by_move_body: None,
             coroutine_drop: None,
             coroutine_layout: None,
         }
@@ -663,10 +650,6 @@ impl<'tcx> Body<'tcx> {
         self.coroutine.as_ref().and_then(|coroutine| coroutine.coroutine_drop.as_ref())
     }
 
-    pub fn coroutine_by_move_body(&self) -> Option<&Body<'tcx>> {
-        self.coroutine.as_ref()?.by_move_body.as_ref()
-    }
-
     #[inline]
     pub fn coroutine_kind(&self) -> Option<CoroutineKind> {
         self.coroutine.as_ref().map(|coroutine| coroutine.coroutine_kind)

compiler/rustc_middle/src/mir/mono.rs

@@ -415,7 +415,6 @@ impl<'tcx> CodegenUnit<'tcx> {
                             | InstanceKind::Virtual(..)
                             | InstanceKind::ClosureOnceShim { .. }
                             | InstanceKind::ConstructCoroutineInClosureShim { .. }
-                            | InstanceKind::CoroutineKindShim { .. }
                             | InstanceKind::DropGlue(..)
                             | InstanceKind::CloneShim(..)
                             | InstanceKind::ThreadLocalShim(..)

compiler/rustc_middle/src/mir/visit.rs

@@ -349,7 +349,6 @@ macro_rules! make_mir_visitor {
                             coroutine_closure_def_id: _def_id,
                             receiver_by_ref: _,
                         } |
-                        ty::InstanceKind::CoroutineKindShim { coroutine_def_id: _def_id } |
                         ty::InstanceKind::AsyncDropGlueCtorShim(_def_id, None) |
                         ty::InstanceKind::DropGlue(_def_id, None) => {}
 

compiler/rustc_middle/src/query/mod.rs

@@ -327,6 +327,7 @@ rustc_queries! {
     query predicates_of(key: DefId) -> ty::GenericPredicates<'tcx> {
         desc { |tcx| "computing predicates of `{}`", tcx.def_path_str(key) }
         cache_on_disk_if { key.is_local() }
+        feedable
     }
 
     query opaque_types_defined_by(
@@ -488,13 +489,19 @@ rustc_queries! {
         separate_provide_extern
     }
 
+    query is_synthetic_mir(key: LocalDefId) -> bool {
+        desc { |tcx| "checking if `{}` is synthetic", tcx.def_path_str(key) }
+        feedable
+    }
+
     /// Build the MIR for a given `DefId` and prepare it for const qualification.
     ///
     /// See the [rustc dev guide] for more info.
     ///
     /// [rustc dev guide]: https://rustc-dev-guide.rust-lang.org/mir/construction.html
     query mir_built(key: LocalDefId) -> &'tcx Steal<mir::Body<'tcx>> {
         desc { |tcx| "building MIR for `{}`", tcx.def_path_str(key) }
+        feedable
     }
 
     /// Try to build an abstract representation of the given constant.
@@ -739,6 +746,7 @@ rustc_queries! {
     query constness(key: DefId) -> hir::Constness {
         desc { |tcx| "checking if item is const: `{}`", tcx.def_path_str(key) }
         separate_provide_extern
+        feedable
     }
 
     query asyncness(key: DefId) -> ty::Asyncness {
@@ -757,10 +765,22 @@ rustc_queries! {
         desc { |tcx| "checking if item is promotable: `{}`", tcx.def_path_str(key) }
     }
 
+    /// The body of the coroutine, modified to take its upvars by move rather than by ref.
+    ///
+    /// This is used by coroutine-closures, which must return a different flavor of coroutine
+    /// when called using `AsyncFnOnce::call_once`. It is produced by the `ByMoveBody` pass which
+    /// is run right after building the initial MIR, and will only be populated for coroutines
+    /// which come out of the async closure desugaring.
+    query coroutine_by_move_body_def_id(def_id: DefId) -> DefId {
+        desc { |tcx| "looking up the coroutine by-move body for `{}`", tcx.def_path_str(def_id) }
+        separate_provide_extern
+    }
+
     /// Returns `Some(coroutine_kind)` if the node pointed to by `def_id` is a coroutine.
     query coroutine_kind(def_id: DefId) -> Option<hir::CoroutineKind> {
         desc { |tcx| "looking up coroutine kind of `{}`", tcx.def_path_str(def_id) }
         separate_provide_extern
+        feedable
     }
 
     query coroutine_for_closure(def_id: DefId) -> DefId {

compiler/rustc_middle/src/ty/instance.rs

@@ -141,14 +141,6 @@ pub enum InstanceKind<'tcx> {
         receiver_by_ref: bool,
     },
 
-    /// `<[coroutine] as Future>::poll`, but for coroutines produced when `AsyncFnOnce`
-    /// is called on a coroutine-closure whose closure kind greater than `FnOnce`, or
-    /// similarly for `AsyncFnMut`.
-    ///
-    /// This will select the body that is produced by the `ByMoveBody` transform, and thus
-    /// take and use all of its upvars by-move rather than by-ref.
-    CoroutineKindShim { coroutine_def_id: DefId },
-
     /// Compiler-generated accessor for thread locals which returns a reference to the thread local
     /// the `DefId` defines. This is used to export thread locals from dylibs on platforms lacking
     /// native support.
@@ -248,7 +240,6 @@ impl<'tcx> InstanceKind<'tcx> {
                 coroutine_closure_def_id: def_id,
                 receiver_by_ref: _,
             }
-            | ty::InstanceKind::CoroutineKindShim { coroutine_def_id: def_id }
             | InstanceKind::DropGlue(def_id, _)
             | InstanceKind::CloneShim(def_id, _)
             | InstanceKind::FnPtrAddrShim(def_id, _)
@@ -270,7 +261,6 @@ impl<'tcx> InstanceKind<'tcx> {
             | InstanceKind::Intrinsic(..)
             | InstanceKind::ClosureOnceShim { .. }
             | ty::InstanceKind::ConstructCoroutineInClosureShim { .. }
-            | ty::InstanceKind::CoroutineKindShim { .. }
             | InstanceKind::DropGlue(..)
             | InstanceKind::AsyncDropGlueCtorShim(..)
             | InstanceKind::CloneShim(..)
@@ -377,7 +367,6 @@ impl<'tcx> InstanceKind<'tcx> {
             | InstanceKind::AsyncDropGlueCtorShim(_, Some(_)) => false,
             InstanceKind::ClosureOnceShim { .. }
             | InstanceKind::ConstructCoroutineInClosureShim { .. }
-            | InstanceKind::CoroutineKindShim { .. }
             | InstanceKind::DropGlue(..)
             | InstanceKind::AsyncDropGlueCtorShim(..)
             | InstanceKind::Item(_)
@@ -452,7 +441,6 @@ pub fn fmt_instance(
         InstanceKind::FnPtrShim(_, ty) => write!(f, " - shim({ty})"),
         InstanceKind::ClosureOnceShim { .. } => write!(f, " - shim"),
         InstanceKind::ConstructCoroutineInClosureShim { .. } => write!(f, " - shim"),
-        InstanceKind::CoroutineKindShim { .. } => write!(f, " - shim"),
         InstanceKind::DropGlue(_, None) => write!(f, " - shim(None)"),
         InstanceKind::DropGlue(_, Some(ty)) => write!(f, " - shim(Some({ty}))"),
         InstanceKind::CloneShim(_, ty) => write!(f, " - shim({ty})"),
@@ -850,7 +838,9 @@ impl<'tcx> Instance<'tcx> {
                 Some(Instance { def: ty::InstanceKind::Item(coroutine_def_id), args })
             } else {
                 Some(Instance {
-                    def: ty::InstanceKind::CoroutineKindShim { coroutine_def_id },
+                    def: ty::InstanceKind::Item(
+                        tcx.coroutine_by_move_body_def_id(coroutine_def_id),
+                    ),
                     args,
                 })
             }

compiler/rustc_middle/src/ty/mod.rs

@@ -1755,7 +1755,6 @@ impl<'tcx> TyCtxt<'tcx> {
             | ty::InstanceKind::Virtual(..)
             | ty::InstanceKind::ClosureOnceShim { .. }
             | ty::InstanceKind::ConstructCoroutineInClosureShim { .. }
-            | ty::InstanceKind::CoroutineKindShim { .. }
             | ty::InstanceKind::DropGlue(..)
             | ty::InstanceKind::CloneShim(..)
             | ty::InstanceKind::ThreadLocalShim(..)
@@ -1873,7 +1872,8 @@ impl<'tcx> TyCtxt<'tcx> {
                     identity_kind_ty.to_opt_closure_kind(),
                     Some(ClosureKind::Fn | ClosureKind::FnMut)
                 );
-                mir.coroutine_by_move_body().unwrap().coroutine_layout_raw()
+                self.optimized_mir(self.coroutine_by_move_body_def_id(def_id))
+                    .coroutine_layout_raw()
             }
         }
     }

compiler/rustc_mir_transform/src/coroutine.rs

@@ -53,7 +53,7 @@
 mod by_move_body;
 use std::{iter, ops};
 
-pub use by_move_body::ByMoveBody;
+pub use by_move_body::coroutine_by_move_body_def_id;
 use rustc_data_structures::fx::FxHashSet;
 use rustc_errors::pluralize;
 use rustc_hir as hir;