Rollup merge of #135314 - compiler-errors:eagerly-mono-closures, r=wesleywiser

Eagerly collect mono items for non-generic closures

This allows users to use `-Zprint-mono-items=eager` to eagerly monomorphize closures and coroutine bodies, in case they want to inspect the LLVM or ASM for those items.

`-Zprint-mono-items`, which used to be called `-Zprint-trans-items`, was originally added in #30900:

> Eager mode is meant to be used in conjunction with incremental compilation
> where a stable set of translation items is more important than a minimal
> one. Thus, eager mode will instantiate drop-glue for every drop-able type
> in the crate, even of no drop call for that type exists (yet). It will
> also instantiate default implementations of trait methods, something that
> otherwise is only done on demand.

Although it remains an unstable option, its purpose has somewhat expanded since then, and as far as I can tell it's generally useful for cases when you want to monomorphize as many items as possible, even if they're unreachable. Specifically, it's useful for debugging since you can look at the codegen'd body of a function, since we don't emit items that are not reachable in monomorphization.

And even more specifically, it would be very to monomorphize the coroutine body of an async fn, since those you can't easily call those without a runtime. This PR enables this usecase since we now monomorphize `DefKind::Closure`.

Author: matthiaskrgr

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

@@ -1377,6 +1377,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();
     }
 
@@ -1459,6 +1463,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;
     };