Rollup merge of #144094 - saethlin:codegen-the-main-fn, r=petrochenkov

Ensure we codegen the main fn

This fixes two bugs. The one that was identified in the linked issue is that when we have a `main` function, mono collection didn't consider it as an extra collection root.

The other is that since CGU partitioning doesn't  know about the call edges between the entrypoint functions, naively it can put them in different CGUs and mark them all as internal. Which would result in LLVM just deleting all of them. There was an existing hack to exclude `lang = "start"` from internalization, which I've extended to include `main`.

Fixes https://github.com/rust-lang/rust/issues/144052

6 files changed, 44 insertions, 14 deletions

diff --git a/compiler/rustc_middle/src/mir/mono.rs b/compiler/rustc_middle/src/mir/mono.rs
index 2d7ddd105bd..105736b9e24 100644
--- a/compiler/rustc_middle/src/mir/mono.rs
+++ b/compiler/rustc_middle/src/mir/mono.rs
@@ -143,10 +143,8 @@ impl<'tcx> MonoItem<'tcx> {
         };
 
         // Similarly, the executable entrypoint must be instantiated exactly once.
-        if let Some((entry_def_id, _)) = tcx.entry_fn(()) {
-            if instance.def_id() == entry_def_id {
-                return InstantiationMode::GloballyShared { may_conflict: false };
-            }
+        if tcx.is_entrypoint(instance.def_id()) {
+            return InstantiationMode::GloballyShared { may_conflict: false };
         }
 
         // If the function is #[naked] or contains any other attribute that requires exactly-once
diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
index 9d9cc3403dc..66d1335e763 100644
--- a/compiler/rustc_middle/src/ty/context.rs
+++ b/compiler/rustc_middle/src/ty/context.rs
@@ -3419,6 +3419,20 @@ impl<'tcx> TyCtxt<'tcx> {
     pub fn do_not_recommend_impl(self, def_id: DefId) -> bool {
         self.get_diagnostic_attr(def_id, sym::do_not_recommend).is_some()
     }
+
+    /// Whether this def is one of the special bin crate entrypoint functions that must have a
+    /// monomorphization and also not be internalized in the bin crate.
+    pub fn is_entrypoint(self, def_id: DefId) -> bool {
+        if self.is_lang_item(def_id, LangItem::Start) {
+            return true;
+        }
+        if let Some((entry_def_id, _)) = self.entry_fn(())
+            && entry_def_id == def_id
+        {
+            return true;
+        }
+        false
+    }
 }
 
 /// Parameter attributes that can only be determined by examining the body of a function instead
diff --git a/compiler/rustc_monomorphize/src/collector.rs b/compiler/rustc_monomorphize/src/collector.rs
index d435e4e77b9..1bfd83d97ac 100644
--- a/compiler/rustc_monomorphize/src/collector.rs
+++ b/compiler/rustc_monomorphize/src/collector.rs
@@ -1575,6 +1575,15 @@ impl<'v> RootCollector<'_, 'v> {
             return;
         };
 
+        let main_instance = Instance::mono(self.tcx, main_def_id);
+        if self.tcx.should_codegen_locally(main_instance) {
+            self.output.push(create_fn_mono_item(
+                self.tcx,
+                main_instance,
+                self.tcx.def_span(main_def_id),
+            ));
+        }
+
         let Some(start_def_id) = self.tcx.lang_items().start_fn() else {
             self.tcx.dcx().emit_fatal(errors::StartNotFound);
         };
diff --git a/compiler/rustc_monomorphize/src/partitioning.rs b/compiler/rustc_monomorphize/src/partitioning.rs
index 69851511fb1..ca8228de57e 100644
--- a/compiler/rustc_monomorphize/src/partitioning.rs
+++ b/compiler/rustc_monomorphize/src/partitioning.rs
@@ -223,11 +223,7 @@ where
         // So even if its mode is LocalCopy, we need to treat it like a root.
         match mono_item.instantiation_mode(cx.tcx) {
             InstantiationMode::GloballyShared { .. } => {}
-            InstantiationMode::LocalCopy => {
-                if !cx.tcx.is_lang_item(mono_item.def_id(), LangItem::Start) {
-                    continue;
-                }
-            }
+            InstantiationMode::LocalCopy => continue,
         }
 
         let characteristic_def_id = characteristic_def_id_of_mono_item(cx.tcx, mono_item);
@@ -821,10 +817,9 @@ fn mono_item_visibility<'tcx>(
         | InstanceKind::FnPtrAddrShim(..) => return Visibility::Hidden,
     };
 
-    // The `start_fn` lang item is actually a monomorphized instance of a
-    // function in the standard library, used for the `main` function. We don't
-    // want to export it so we tag it with `Hidden` visibility but this symbol
-    // is only referenced from the actual `main` symbol which we unfortunately
+    // Both the `start_fn` lang item and `main` itself should not be exported,
+    // so we give them with `Hidden` visibility but these symbols are
+    // only referenced from the actual `main` symbol which we unfortunately
     // don't know anything about during partitioning/collection. As a result we
     // forcibly keep this symbol out of the `internalization_candidates` set.
     //
@@ -834,7 +829,7 @@ fn mono_item_visibility<'tcx>(
     //        from the `main` symbol we'll generate later.
     //
     //        This may be fixable with a new `InstanceKind` perhaps? Unsure!
-    if tcx.is_lang_item(def_id, LangItem::Start) {
+    if tcx.is_entrypoint(def_id) {
         *can_be_internalized = false;
         return Visibility::Hidden;
     }
diff --git a/tests/ui/entry-point/auxiliary/main_functions.rs b/tests/ui/entry-point/auxiliary/main_functions.rs
index cc7992a42c1..ab4a09b6331 100644
--- a/tests/ui/entry-point/auxiliary/main_functions.rs
+++ b/tests/ui/entry-point/auxiliary/main_functions.rs
@@ -1 +1,4 @@
 pub fn boilerplate() {}
+
+#[inline]
+pub fn local_codegen() {}
diff --git a/tests/ui/entry-point/imported_main_local_codegen.rs b/tests/ui/entry-point/imported_main_local_codegen.rs
new file mode 100644
index 00000000000..1e46c109373
--- /dev/null
+++ b/tests/ui/entry-point/imported_main_local_codegen.rs
@@ -0,0 +1,11 @@
+//@ run-pass
+//@ aux-build:main_functions.rs
+//@ compile-flags: -Ccodegen-units=1024
+
+// This is a regression test for https://github.com/rust-lang/rust/issues/144052.
+// Entrypoint functions call each other in ways that CGU partitioning doesn't know about. So there
+// is a special check to not internalize any of them. But internalizing them can be okay if there
+// are few enough CGUs, so we use a lot of CGUs in this test to hit the bad case.
+
+extern crate main_functions;
+pub use main_functions::local_codegen as main;