Ensure we codegen and don't internalize the entrypoint

Author: saethlin

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

compiler/rustc_middle/src/ty/context.rs

@@ -3402,6 +3402,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

compiler/rustc_monomorphize/src/collector.rs

@@ -1582,6 +1582,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);
         };

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;
     }

tests/ui/entry-point/auxiliary/main_functions.rs

@@ -1 +1,4 @@
 pub fn boilerplate() {}
+
+#[inline]
+pub fn local_codegen() {}

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;