From 89a67051a76f8511372d7b9f14610474b1dba6aa Mon Sep 17 00:00:00 2001
From: Charles Lew <crlf0710@gmail.com>
Date: Sat, 24 Apr 2021 13:16:34 +0800
Subject: [PATCH] Add primary marker on codegen unit to take charge of
 main_wrapper for non-local cases.

---
 compiler/rustc_codegen_ssa/src/base.rs          | 17 +++--------------
 compiler/rustc_middle/src/mir/mono.rs           | 12 +++++++++++-
 .../src/monomorphize/partitioning/mod.rs        |  6 ++++--
 .../imported_main_from_extern_crate.rs          |  6 ++----
 .../imported_main_from_extern_crate.stderr      | 10 ----------
 5 files changed, 20 insertions(+), 31 deletions(-)
 delete mode 100644 src/test/ui/entry-point/imported_main_from_extern_crate.stderr

diff --git a/compiler/rustc_codegen_ssa/src/base.rs b/compiler/rustc_codegen_ssa/src/base.rs
index e045a23eb0ce3..7a19b0e4d5ac3 100644
--- a/compiler/rustc_codegen_ssa/src/base.rs
+++ b/compiler/rustc_codegen_ssa/src/base.rs
@@ -357,20 +357,9 @@ pub fn maybe_create_entry_wrapper<'a, 'tcx, Bx: BuilderMethods<'a, 'tcx>>(
         if !cx.codegen_unit().contains_item(&MonoItem::Fn(instance)) {
             return None;
         }
-    } else {
-        // FIXME: Add support for non-local main fn codegen
-        let span = cx.tcx().main_def.unwrap().span;
-        let n = 28937;
-        cx.sess()
-            .struct_span_err(span, "entry symbol `main` from foreign crate is not yet supported.")
-            .note(&format!(
-                "see issue #{} <https://github.com/rust-lang/rust/issues/{}> \
-                 for more information",
-                n, n,
-            ))
-            .emit();
-        cx.sess().abort_if_errors();
-        bug!();
+    } else if !cx.codegen_unit().is_primary() {
+        // We want to create the wrapper only when the codegen unit is the primary one
+        return None;
     }
 
     let main_llfn = cx.get_fn_addr(instance);
diff --git a/compiler/rustc_middle/src/mir/mono.rs b/compiler/rustc_middle/src/mir/mono.rs
index 77f38e52ad2e4..67440e6e0edf0 100644
--- a/compiler/rustc_middle/src/mir/mono.rs
+++ b/compiler/rustc_middle/src/mir/mono.rs
@@ -229,6 +229,7 @@ pub struct CodegenUnit<'tcx> {
     name: Symbol,
     items: FxHashMap<MonoItem<'tcx>, (Linkage, Visibility)>,
     size_estimate: Option<usize>,
+    primary: bool,
 }
 
 /// Specifies the linkage type for a `MonoItem`.
@@ -258,7 +259,7 @@ pub enum Visibility {
 
 impl<'tcx> CodegenUnit<'tcx> {
     pub fn new(name: Symbol) -> CodegenUnit<'tcx> {
-        CodegenUnit { name, items: Default::default(), size_estimate: None }
+        CodegenUnit { name, items: Default::default(), size_estimate: None, primary: false }
     }
 
     pub fn name(&self) -> Symbol {
@@ -269,6 +270,14 @@ impl<'tcx> CodegenUnit<'tcx> {
         self.name = name;
     }
 
+    pub fn is_primary(&self) -> bool {
+        self.primary
+    }
+
+    pub fn make_primary(&mut self) {
+        self.primary = true;
+    }
+
     pub fn items(&self) -> &FxHashMap<MonoItem<'tcx>, (Linkage, Visibility)> {
         &self.items
     }
@@ -378,6 +387,7 @@ impl<'a, 'tcx> HashStable<StableHashingContext<'a>> for CodegenUnit<'tcx> {
             name,
             // The size estimate is not relevant to the hash
             size_estimate: _,
+            primary: _,
         } = *self;
 
         name.hash_stable(hcx, hasher);
diff --git a/compiler/rustc_mir/src/monomorphize/partitioning/mod.rs b/compiler/rustc_mir/src/monomorphize/partitioning/mod.rs
index dc2379fd92b83..333cb30159093 100644
--- a/compiler/rustc_mir/src/monomorphize/partitioning/mod.rs
+++ b/compiler/rustc_mir/src/monomorphize/partitioning/mod.rs
@@ -350,12 +350,14 @@ fn collect_and_partition_mono_items<'tcx>(
     let (codegen_units, _) = tcx.sess.time("partition_and_assert_distinct_symbols", || {
         sync::join(
             || {
-                &*tcx.arena.alloc_from_iter(partition(
+                let mut codegen_units = partition(
                     tcx,
                     &mut items.iter().cloned(),
                     tcx.sess.codegen_units(),
                     &inlining_map,
-                ))
+                );
+                codegen_units[0].make_primary();
+                &*tcx.arena.alloc_from_iter(codegen_units)
             },
             || assert_symbols_are_distinct(tcx, items.iter()),
         )
diff --git a/src/test/ui/entry-point/imported_main_from_extern_crate.rs b/src/test/ui/entry-point/imported_main_from_extern_crate.rs
index 6bbf67fa5408d..4fddfc44ac60a 100644
--- a/src/test/ui/entry-point/imported_main_from_extern_crate.rs
+++ b/src/test/ui/entry-point/imported_main_from_extern_crate.rs
@@ -1,9 +1,7 @@
-// build-fail
+// run-pass
 // aux-build:main_functions.rs
 
 #![feature(imported_main)]
 
 extern crate main_functions;
-pub use main_functions::boilerplate as main; //~ ERROR entry symbol `main` from foreign crate
-
-// FIXME: Should be run-pass
+pub use main_functions::boilerplate as main;
diff --git a/src/test/ui/entry-point/imported_main_from_extern_crate.stderr b/src/test/ui/entry-point/imported_main_from_extern_crate.stderr
deleted file mode 100644
index 8792e1e414244..0000000000000
--- a/src/test/ui/entry-point/imported_main_from_extern_crate.stderr
+++ /dev/null
@@ -1,10 +0,0 @@
-error: entry symbol `main` from foreign crate is not yet supported.
-  --> $DIR/imported_main_from_extern_crate.rs:7:9
-   |
-LL | pub use main_functions::boilerplate as main;
-   |         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-   |
-   = note: see issue #28937 <https://github.com/rust-lang/rust/issues/28937> for more information
-
-error: aborting due to previous error
-