Eagerly collect mono items for non-generic closures

Author: compiler-errors

compiler/rustc_middle/src/hir/map/mod.rs

@@ -1246,6 +1246,7 @@ pub(super) fn hir_module_items(tcx: TyCtxt<'_>, module_id: LocalModDefId) -> Mod
         foreign_items,
         body_owners,
         opaques,
+        nested_bodies,
         ..
     } = collector;
     ModuleItems {
@@ -1256,6 +1257,7 @@ pub(super) fn hir_module_items(tcx: TyCtxt<'_>, module_id: LocalModDefId) -> Mod
         foreign_items: foreign_items.into_boxed_slice(),
         body_owners: body_owners.into_boxed_slice(),
         opaques: opaques.into_boxed_slice(),
+        nested_bodies: nested_bodies.into_boxed_slice(),
     }
 }
 
@@ -1276,6 +1278,7 @@ pub(crate) fn hir_crate_items(tcx: TyCtxt<'_>, _: ()) -> ModuleItems {
         foreign_items,
         body_owners,
         opaques,
+        nested_bodies,
         ..
     } = collector;
 
@@ -1287,6 +1290,7 @@ pub(crate) fn hir_crate_items(tcx: TyCtxt<'_>, _: ()) -> ModuleItems {
         foreign_items: foreign_items.into_boxed_slice(),
         body_owners: body_owners.into_boxed_slice(),
         opaques: opaques.into_boxed_slice(),
+        nested_bodies: nested_bodies.into_boxed_slice(),
     }
 }
 
@@ -1302,6 +1306,7 @@ struct ItemCollector<'tcx> {
     foreign_items: Vec<ForeignItemId>,
     body_owners: Vec<LocalDefId>,
     opaques: Vec<LocalDefId>,
+    nested_bodies: Vec<LocalDefId>,
 }
 
 impl<'tcx> ItemCollector<'tcx> {
@@ -1316,6 +1321,7 @@ impl<'tcx> ItemCollector<'tcx> {
             foreign_items: Vec::default(),
             body_owners: Vec::default(),
             opaques: Vec::default(),
+            nested_bodies: Vec::default(),
         }
     }
 }
@@ -1358,6 +1364,7 @@ impl<'hir> Visitor<'hir> for ItemCollector<'hir> {
 
     fn visit_inline_const(&mut self, c: &'hir ConstBlock) {
         self.body_owners.push(c.def_id);
+        self.nested_bodies.push(c.def_id);
         intravisit::walk_inline_const(self, c)
     }
 
@@ -1369,6 +1376,7 @@ impl<'hir> Visitor<'hir> for ItemCollector<'hir> {
     fn visit_expr(&mut self, ex: &'hir Expr<'hir>) {
         if let ExprKind::Closure(closure) = ex.kind {
             self.body_owners.push(closure.def_id);
+            self.nested_bodies.push(closure.def_id);
         }
         intravisit::walk_expr(self, ex)
     }

compiler/rustc_middle/src/hir/mod.rs

@@ -30,6 +30,7 @@ pub struct ModuleItems {
     foreign_items: Box<[ForeignItemId]>,
     opaques: Box<[LocalDefId]>,
     body_owners: Box<[LocalDefId]>,
+    nested_bodies: Box<[LocalDefId]>,
 }
 
 impl ModuleItems {
@@ -70,6 +71,10 @@ impl ModuleItems {
         self.opaques.iter().copied()
     }
 
+    pub fn nested_bodies(&self) -> impl Iterator<Item = LocalDefId> + '_ {
+        self.nested_bodies.iter().copied()
+    }
+
     pub fn definitions(&self) -> impl Iterator<Item = LocalDefId> + '_ {
         self.owners().map(|id| id.def_id)
     }

compiler/rustc_monomorphize/src/collector.rs

@@ -1368,6 +1368,10 @@ fn collect_roots(tcx: TyCtxt<'_>, mode: MonoItemCollectionStrategy) -> Vec<MonoI
             collector.process_impl_item(id);
         }
 
+        for id in crate_items.nested_bodies() {
+            collector.process_nested_body(id);
+        }
+
         collector.push_extra_entry_roots();
     }
 
@@ -1450,6 +1454,33 @@ impl<'v> RootCollector<'_, 'v> {
         }
     }
 
+    fn process_nested_body(&mut self, def_id: LocalDefId) {
+        match self.tcx.def_kind(def_id) {
+            DefKind::Closure => {
+                if self.strategy == MonoItemCollectionStrategy::Eager
+                    && !self
+                        .tcx
+                        .generics_of(self.tcx.typeck_root_def_id(def_id.to_def_id()))
+                        .requires_monomorphization(self.tcx)
+                {
+                    let instance = match *self.tcx.type_of(def_id).instantiate_identity().kind() {
+                        ty::Closure(def_id, args)
+                        | ty::Coroutine(def_id, args)
+                        | ty::CoroutineClosure(def_id, args) => {
+                            Instance::new(def_id, self.tcx.erase_regions(args))
+                        }
+                        _ => unreachable!(),
+                    };
+                    let mono_item = create_fn_mono_item(self.tcx, instance, DUMMY_SP);
+                    if mono_item.node.is_instantiable(self.tcx) {
+                        self.output.push(mono_item);
+                    }
+                }
+            }
+            _ => {}
+        }
+    }
+
     fn is_root(&self, def_id: LocalDefId) -> bool {
         !self.tcx.generics_of(def_id).requires_monomorphization(self.tcx)
             && match self.strategy {

tests/codegen-units/item-collection/closures.rs

@@ -0,0 +1,12 @@
+//@ edition: 2021
+//@ compile-flags: -Zprint-mono-items=eager --crate-type=lib
+
+//~ MONO_ITEM fn async_fn @@
+//~ MONO_ITEM fn async_fn::{closure#0} @@
+pub async fn async_fn() {}
+
+//~ MONO_ITEM fn closure @@
+//~ MONO_ITEM fn closure::{closure#0} @@
+pub fn closure() {
+    let _ = || {};
+}

tests/codegen-units/item-collection/non-generic-closures.rs

@@ -13,6 +13,7 @@ fn temporary() {
 
 //~ MONO_ITEM fn assigned_to_variable_but_not_executed @@ non_generic_closures-cgu.0[Internal]
 fn assigned_to_variable_but_not_executed() {
+    //~ MONO_ITEM fn assigned_to_variable_but_not_executed::{closure#0}
     let _x = |a: i16| {
         let _ = a + 1;
     };
@@ -21,9 +22,9 @@ fn assigned_to_variable_but_not_executed() {
 //~ MONO_ITEM fn assigned_to_variable_executed_indirectly @@ non_generic_closures-cgu.0[Internal]
 fn assigned_to_variable_executed_indirectly() {
     //~ MONO_ITEM fn assigned_to_variable_executed_indirectly::{closure#0} @@ non_generic_closures-cgu.0[Internal]
-    //~ MONO_ITEM fn <{closure@TEST_PATH:27:13: 27:21} as std::ops::FnOnce<(i32,)>>::call_once - shim @@ non_generic_closures-cgu.0[Internal]
-    //~ MONO_ITEM fn <{closure@TEST_PATH:27:13: 27:21} as std::ops::FnOnce<(i32,)>>::call_once - shim(vtable) @@ non_generic_closures-cgu.0[Internal]
-    //~ MONO_ITEM fn std::ptr::drop_in_place::<{closure@TEST_PATH:27:13: 27:21}> - shim(None) @@ non_generic_closures-cgu.0[Internal]
+    //~ MONO_ITEM fn <{closure@TEST_PATH:28:13: 28:21} as std::ops::FnOnce<(i32,)>>::call_once - shim @@ non_generic_closures-cgu.0[Internal]
+    //~ MONO_ITEM fn <{closure@TEST_PATH:28:13: 28:21} as std::ops::FnOnce<(i32,)>>::call_once - shim(vtable) @@ non_generic_closures-cgu.0[Internal]
+    //~ MONO_ITEM fn std::ptr::drop_in_place::<{closure@TEST_PATH:28:13: 28:21}> - shim(None) @@ non_generic_closures-cgu.0[Internal]
     let f = |a: i32| {
         let _ = a + 2;
     };