Rework how the codegen coordinator code handles the allocator shim

compiler/rustc_codegen_cranelift/src/driver/aot.rs

@@ -12,9 +12,7 @@ use cranelift_object::{ObjectBuilder, ObjectModule};
 use rustc_codegen_ssa::assert_module_sources::CguReuse;
 use rustc_codegen_ssa::back::link::ensure_removed;
 use rustc_codegen_ssa::base::determine_cgu_reuse;
-use rustc_codegen_ssa::{
-    CodegenResults, CompiledModule, CrateInfo, ModuleKind, errors as ssa_errors,
-};
+use rustc_codegen_ssa::{CodegenResults, CompiledModule, CrateInfo, errors as ssa_errors};
 use rustc_data_structures::profiling::SelfProfilerRef;
 use rustc_data_structures::stable_hasher::{HashStable, StableHasher};
 use rustc_data_structures::sync::{IntoDynSyncSend, par_map};
@@ -363,7 +361,6 @@ fn emit_cgu(
         invocation_temp,
         prof,
         product.object,
-        ModuleKind::Regular,
         name.clone(),
         producer,
     )?;
@@ -372,7 +369,6 @@ fn emit_cgu(
         module_regular,
         module_global_asm: global_asm_object_file.map(|global_asm_object_file| CompiledModule {
             name: format!("{name}.asm"),
-            kind: ModuleKind::Regular,
             object: Some(global_asm_object_file),
             dwarf_object: None,
             bytecode: None,
@@ -389,7 +385,6 @@ fn emit_module(
     invocation_temp: Option<&str>,
     prof: &SelfProfilerRef,
     mut object: cranelift_object::object::write::Object<'_>,
-    kind: ModuleKind,
     name: String,
     producer_str: &str,
 ) -> Result<CompiledModule, String> {
@@ -430,7 +425,6 @@ fn emit_module(
 
     Ok(CompiledModule {
         name,
-        kind,
         object: Some(tmp_file),
         dwarf_object: None,
         bytecode: None,
@@ -485,7 +479,6 @@ fn reuse_workproduct_for_cgu(
     Ok(ModuleCodegenResult {
         module_regular: CompiledModule {
             name: cgu.name().to_string(),
-            kind: ModuleKind::Regular,
             object: Some(obj_out_regular),
             dwarf_object: None,
             bytecode: None,
@@ -495,7 +488,6 @@ fn reuse_workproduct_for_cgu(
         },
         module_global_asm: source_file_global_asm.map(|source_file| CompiledModule {
             name: cgu.name().to_string(),
-            kind: ModuleKind::Regular,
             object: Some(obj_out_global_asm),
             dwarf_object: None,
             bytecode: None,
@@ -651,7 +643,6 @@ fn emit_allocator_module(tcx: TyCtxt<'_>) -> Option<CompiledModule> {
             tcx.sess.invocation_temp.as_deref(),
             &tcx.sess.prof,
             product.object,
-            ModuleKind::Allocator,
             "allocator_shim".to_owned(),
             &crate::debuginfo::producer(tcx.sess),
         ) {

compiler/rustc_codegen_gcc/src/back/lto.rs

@@ -204,7 +204,7 @@ fn fat_lto(
             let path = tmp_path.path().to_path_buf().join(&module.name);
             let path = path.to_str().expect("path");
             let context = &module.module_llvm.context;
-            let config = cgcx.config(module.kind);
+            let config = &cgcx.module_config;
             // NOTE: we need to set the optimization level here in order for LTO to do its job.
             context.set_optimization_level(to_gcc_opt_level(config.opt_level));
             context.add_command_line_option("-flto=auto");

compiler/rustc_codegen_llvm/src/back/lto.rs

@@ -11,7 +11,7 @@ use object::{Object, ObjectSection};
 use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule, ThinShared};
 use rustc_codegen_ssa::back::write::{CodegenContext, FatLtoInput};
 use rustc_codegen_ssa::traits::*;
-use rustc_codegen_ssa::{ModuleCodegen, ModuleKind, looks_like_rust_object_file};
+use rustc_codegen_ssa::{ModuleCodegen, looks_like_rust_object_file};
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::memmap::Mmap;
 use rustc_errors::DiagCtxtHandle;
@@ -43,9 +43,7 @@ fn prepare_lto(
         .map(|symbol| CString::new(symbol.to_owned()).unwrap())
         .collect::<Vec<CString>>();
 
-    if cgcx.regular_module_config.instrument_coverage
-        || cgcx.regular_module_config.pgo_gen.enabled()
-    {
+    if cgcx.module_config.instrument_coverage || cgcx.module_config.pgo_gen.enabled() {
         // These are weak symbols that point to the profile version and the
         // profile name, which need to be treated as exported so LTO doesn't nix
         // them.
@@ -55,15 +53,15 @@ fn prepare_lto(
         symbols_below_threshold.extend(PROFILER_WEAK_SYMBOLS.iter().map(|&sym| sym.to_owned()));
     }
 
-    if cgcx.regular_module_config.sanitizer.contains(SanitizerSet::MEMORY) {
+    if cgcx.module_config.sanitizer.contains(SanitizerSet::MEMORY) {
         let mut msan_weak_symbols = Vec::new();
 
         // Similar to profiling, preserve weak msan symbol during LTO.
-        if cgcx.regular_module_config.sanitizer_recover.contains(SanitizerSet::MEMORY) {
+        if cgcx.module_config.sanitizer_recover.contains(SanitizerSet::MEMORY) {
             msan_weak_symbols.push(c"__msan_keep_going");
         }
 
-        if cgcx.regular_module_config.sanitizer_memory_track_origins != 0 {
+        if cgcx.module_config.sanitizer_memory_track_origins != 0 {
             msan_weak_symbols.push(c"__msan_track_origins");
         }
 
@@ -227,15 +225,9 @@ fn fat_lto(
     // All the other modules will be serialized and reparsed into the new
     // context, so this hopefully avoids serializing and parsing the largest
     // codegen unit.
-    //
-    // Additionally use a regular module as the base here to ensure that various
-    // file copy operations in the backend work correctly. The only other kind
-    // of module here should be an allocator one, and if your crate is smaller
-    // than the allocator module then the size doesn't really matter anyway.
     let costliest_module = in_memory
         .iter()
         .enumerate()
-        .filter(|&(_, module)| module.kind == ModuleKind::Regular)
         .map(|(i, module)| {
             let cost = unsafe { llvm::LLVMRustModuleCost(module.module_llvm.llmod()) };
             (cost, i)
@@ -583,7 +575,7 @@ pub(crate) fn run_pass_manager(
     thin: bool,
 ) {
     let _timer = cgcx.prof.generic_activity_with_arg("LLVM_lto_optimize", &*module.name);
-    let config = cgcx.config(module.kind);
+    let config = &cgcx.module_config;
 
     // Now we have one massive module inside of llmod. Time to run the
     // LTO-specific optimization passes that LLVM provides.
@@ -745,7 +737,7 @@ pub(crate) fn optimize_thin_module(
     let module_llvm = ModuleLlvm::parse(cgcx, module_name, thin_module.data(), dcx);
     let mut module = ModuleCodegen::new_regular(thin_module.name(), module_llvm);
     // Given that the newly created module lacks a thinlto buffer for embedding, we need to re-add it here.
-    if cgcx.config(ModuleKind::Regular).embed_bitcode() {
+    if cgcx.module_config.embed_bitcode() {
         module.thin_lto_buffer = Some(thin_module.data().to_vec());
     }
     {

compiler/rustc_codegen_ssa/src/back/lto.rs

@@ -1,6 +1,7 @@
 use std::ffi::CString;
 use std::sync::Arc;
 
+use rustc_ast::expand::allocator::AllocatorKind;
 use rustc_data_structures::memmap::Mmap;
 use rustc_hir::def_id::{CrateNum, LOCAL_CRATE};
 use rustc_middle::middle::exported_symbols::{ExportedSymbol, SymbolExportInfo, SymbolExportLevel};
@@ -95,6 +96,19 @@ pub(super) fn exported_symbols_for_lto(
             .filter_map(|&(s, info): &(ExportedSymbol<'_>, SymbolExportInfo)| {
                 if info.level.is_below_threshold(export_threshold) || info.used {
                     Some(symbol_name_for_instance_in_crate(tcx, s, cnum))
+                } else if export_threshold == SymbolExportLevel::C
+                    && info.rustc_std_internal_symbol
+                    && let Some(AllocatorKind::Default) = allocator_kind_for_codegen(tcx)
+                {
+                    // Export the __rdl_* exports for usage by the allocator shim when not using
+                    // #[global_allocator]. Most of the conditions above are only used to avoid
+                    // unnecessary expensive symbol_name_for_instance_in_crate calls.
+                    let sym = symbol_name_for_instance_in_crate(tcx, s, cnum);
+                    if sym.contains("__rdl_") || sym.contains("__rg_oom") {
+                        Some(sym)
+                    } else {
+                        None
+                    }
                 } else {
                     None
                 }