[WIP] Use weak linkage instead of compiler generated shims

bjorn3 · bjorn3 · commit 1a516634ffb3 · 2024-12-20T09:04:34.000Z
This is still keeping the allocator shim for the oom handler and the
__rust_no_alloc_shim_is_unstable symbol for now.

TODO: Update comments everywhere and test on macOS and Windows
diff --git a/compiler/rustc_ast/src/expand/allocator.rs b/compiler/rustc_ast/src/expand/allocator.rs
@@ -11,10 +11,6 @@ pub fn global_fn_name(base: Symbol) -> String {
     format!("__rust_{base}")
 }
 
-pub fn default_fn_name(base: Symbol) -> String {
-    format!("__rdl_{base}")
-}
-
 pub fn alloc_error_handler_name(alloc_error_handler_kind: AllocatorKind) -> &'static str {
     match alloc_error_handler_kind {
         AllocatorKind::Global => "__rg_oom",
diff --git a/compiler/rustc_codegen_cranelift/src/allocator.rs b/compiler/rustc_codegen_cranelift/src/allocator.rs
@@ -2,72 +2,34 @@
 // Adapted from rustc
 
 use rustc_ast::expand::allocator::{
-    ALLOCATOR_METHODS, AllocatorKind, AllocatorTy, NO_ALLOC_SHIM_IS_UNSTABLE,
-    alloc_error_handler_name, default_fn_name, global_fn_name,
+    AllocatorKind, NO_ALLOC_SHIM_IS_UNSTABLE, alloc_error_handler_name,
 };
-use rustc_codegen_ssa::base::allocator_kind_for_codegen;
+use rustc_codegen_ssa::base::needs_allocator_shim;
 use rustc_session::config::OomStrategy;
 
 use crate::prelude::*;
 
 /// Returns whether an allocator shim was created
 pub(crate) fn codegen(tcx: TyCtxt<'_>, module: &mut dyn Module) -> bool {
-    let Some(kind) = allocator_kind_for_codegen(tcx) else { return false };
-    codegen_inner(
-        module,
-        kind,
-        tcx.alloc_error_handler_kind(()).unwrap(),
-        tcx.sess.opts.unstable_opts.oom,
-    );
-    true
+    if needs_allocator_shim(tcx) {
+        codegen_inner(
+            module,
+            tcx.alloc_error_handler_kind(()).unwrap(),
+            tcx.sess.opts.unstable_opts.oom,
+        );
+        true
+    } else {
+        false
+    }
 }
 
 fn codegen_inner(
     module: &mut dyn Module,
-    kind: AllocatorKind,
     alloc_error_handler_kind: AllocatorKind,
     oom_strategy: OomStrategy,
 ) {
     let usize_ty = module.target_config().pointer_type();
 
-    if kind == AllocatorKind::Default {
-        for method in ALLOCATOR_METHODS {
-            let mut arg_tys = Vec::with_capacity(method.inputs.len());
-            for input in method.inputs.iter() {
-                match input.ty {
-                    AllocatorTy::Layout => {
-                        arg_tys.push(usize_ty); // size
-                        arg_tys.push(usize_ty); // align
-                    }
-                    AllocatorTy::Ptr => arg_tys.push(usize_ty),
-                    AllocatorTy::Usize => arg_tys.push(usize_ty),
-
-                    AllocatorTy::ResultPtr | AllocatorTy::Unit => panic!("invalid allocator arg"),
-                }
-            }
-            let output = match method.output {
-                AllocatorTy::ResultPtr => Some(usize_ty),
-                AllocatorTy::Unit => None,
-
-                AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
-                    panic!("invalid allocator output")
-                }
-            };
-
-            let sig = Signature {
-                call_conv: module.target_config().default_call_conv,
-                params: arg_tys.iter().cloned().map(AbiParam::new).collect(),
-                returns: output.into_iter().map(AbiParam::new).collect(),
-            };
-            crate::common::create_wrapper_function(
-                module,
-                sig,
-                &global_fn_name(method.name),
-                &default_fn_name(method.name),
-            );
-        }
-    }
-
     let sig = Signature {
         call_conv: module.target_config().default_call_conv,
         params: vec![AbiParam::new(usize_ty), AbiParam::new(usize_ty)],
diff --git a/compiler/rustc_codegen_gcc/src/allocator.rs b/compiler/rustc_codegen_gcc/src/allocator.rs
@@ -2,8 +2,7 @@
 use gccjit::FnAttribute;
 use gccjit::{Context, FunctionType, GlobalKind, ToRValue, Type};
 use rustc_ast::expand::allocator::{
-    ALLOCATOR_METHODS, AllocatorKind, AllocatorTy, NO_ALLOC_SHIM_IS_UNSTABLE,
-    alloc_error_handler_name, default_fn_name, global_fn_name,
+    AllocatorKind, NO_ALLOC_SHIM_IS_UNSTABLE, alloc_error_handler_name,
 };
 use rustc_middle::bug;
 use rustc_middle::ty::TyCtxt;
@@ -15,7 +14,6 @@ pub(crate) unsafe fn codegen(
     tcx: TyCtxt<'_>,
     mods: &mut GccContext,
     _module_name: &str,
-    kind: AllocatorKind,
     alloc_error_handler_kind: AllocatorKind,
 ) {
     let context = &mods.context;
@@ -26,37 +24,6 @@ pub(crate) unsafe fn codegen(
         tws => bug!("Unsupported target word size for int: {}", tws),
     };
     let i8 = context.new_type::<i8>();
-    let i8p = i8.make_pointer();
-
-    if kind == AllocatorKind::Default {
-        for method in ALLOCATOR_METHODS {
-            let mut types = Vec::with_capacity(method.inputs.len());
-            for input in method.inputs.iter() {
-                match input.ty {
-                    AllocatorTy::Layout => {
-                        types.push(usize);
-                        types.push(usize);
-                    }
-                    AllocatorTy::Ptr => types.push(i8p),
-                    AllocatorTy::Usize => types.push(usize),
-
-                    AllocatorTy::ResultPtr | AllocatorTy::Unit => panic!("invalid allocator arg"),
-                }
-            }
-            let output = match method.output {
-                AllocatorTy::ResultPtr => Some(i8p),
-                AllocatorTy::Unit => None,
-
-                AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
-                    panic!("invalid allocator output")
-                }
-            };
-            let from_name = global_fn_name(method.name);
-            let to_name = default_fn_name(method.name);
-
-            create_wrapper_function(tcx, context, &from_name, &to_name, &types, output);
-        }
-    }
 
     // FIXME(bjorn3): Add noreturn attribute
     create_wrapper_function(
diff --git a/compiler/rustc_codegen_gcc/src/lib.rs b/compiler/rustc_codegen_gcc/src/lib.rs
@@ -291,7 +291,6 @@ impl ExtraBackendMethods for GccCodegenBackend {
         &self,
         tcx: TyCtxt<'_>,
         module_name: &str,
-        kind: AllocatorKind,
         alloc_error_handler_kind: AllocatorKind,
     ) -> Self::Module {
         let mut mods = GccContext {
@@ -301,7 +300,7 @@ impl ExtraBackendMethods for GccCodegenBackend {
         };
 
         unsafe {
-            allocator::codegen(tcx, &mut mods, module_name, kind, alloc_error_handler_kind);
+            allocator::codegen(tcx, &mut mods, module_name, alloc_error_handler_kind);
         }
         mods
     }
diff --git a/compiler/rustc_codegen_llvm/src/allocator.rs b/compiler/rustc_codegen_llvm/src/allocator.rs
@@ -1,7 +1,6 @@
 use libc::c_uint;
 use rustc_ast::expand::allocator::{
-    ALLOCATOR_METHODS, AllocatorKind, AllocatorTy, NO_ALLOC_SHIM_IS_UNSTABLE,
-    alloc_error_handler_name, default_fn_name, global_fn_name,
+    AllocatorKind, NO_ALLOC_SHIM_IS_UNSTABLE, alloc_error_handler_name,
 };
 use rustc_middle::bug;
 use rustc_middle::ty::TyCtxt;
@@ -15,7 +14,6 @@ pub(crate) unsafe fn codegen(
     tcx: TyCtxt<'_>,
     module_llvm: &mut ModuleLlvm,
     module_name: &str,
-    kind: AllocatorKind,
     alloc_error_handler_kind: AllocatorKind,
 ) {
     let llcx = &*module_llvm.llcx;
@@ -29,38 +27,6 @@ pub(crate) unsafe fn codegen(
         }
     };
     let i8 = unsafe { llvm::LLVMInt8TypeInContext(llcx) };
-    let i8p = unsafe { llvm::LLVMPointerTypeInContext(llcx, 0) };
-
-    if kind == AllocatorKind::Default {
-        for method in ALLOCATOR_METHODS {
-            let mut args = Vec::with_capacity(method.inputs.len());
-            for input in method.inputs.iter() {
-                match input.ty {
-                    AllocatorTy::Layout => {
-                        args.push(usize); // size
-                        args.push(usize); // align
-                    }
-                    AllocatorTy::Ptr => args.push(i8p),
-                    AllocatorTy::Usize => args.push(usize),
-
-                    AllocatorTy::ResultPtr | AllocatorTy::Unit => panic!("invalid allocator arg"),
-                }
-            }
-            let output = match method.output {
-                AllocatorTy::ResultPtr => Some(i8p),
-                AllocatorTy::Unit => None,
-
-                AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
-                    panic!("invalid allocator output")
-                }
-            };
-
-            let from_name = global_fn_name(method.name);
-            let to_name = default_fn_name(method.name);
-
-            create_wrapper_function(tcx, llcx, llmod, &from_name, &to_name, &args, output, false);
-        }
-    }
 
     // rust alloc error handler
     create_wrapper_function(
diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/compiler/rustc_codegen_llvm/src/lib.rs
@@ -114,12 +114,11 @@ impl ExtraBackendMethods for LlvmCodegenBackend {
         &self,
         tcx: TyCtxt<'tcx>,
         module_name: &str,
-        kind: AllocatorKind,
         alloc_error_handler_kind: AllocatorKind,
     ) -> ModuleLlvm {
         let mut module_llvm = ModuleLlvm::new_metadata(tcx, module_name);
         unsafe {
-            allocator::codegen(tcx, &mut module_llvm, module_name, kind, alloc_error_handler_kind);
+            allocator::codegen(tcx, &mut module_llvm, module_name, alloc_error_handler_kind);
         }
         module_llvm
     }
diff --git a/compiler/rustc_codegen_ssa/src/back/symbol_export.rs b/compiler/rustc_codegen_ssa/src/back/symbol_export.rs
@@ -16,7 +16,7 @@ use rustc_session::config::{CrateType, OomStrategy};
 use rustc_target::spec::{SanitizerSet, TlsModel};
 use tracing::debug;
 
-use crate::base::allocator_kind_for_codegen;
+use crate::base::needs_allocator_shim;
 
 fn threshold(tcx: TyCtxt<'_>) -> SymbolExportLevel {
     crates_export_threshold(tcx.crate_types())
@@ -206,7 +206,7 @@ fn exported_symbols_provider_local(
     }
 
     // Mark allocator shim symbols as exported only if they were generated.
-    if allocator_kind_for_codegen(tcx).is_some() {
+    if needs_allocator_shim(tcx) {
         for symbol_name in ALLOCATOR_METHODS
             .iter()
             .map(|method| format!("__rust_{}", method.name))
diff --git a/compiler/rustc_codegen_ssa/src/base.rs b/compiler/rustc_codegen_ssa/src/base.rs
@@ -4,7 +4,7 @@ use std::time::{Duration, Instant};
 
 use itertools::Itertools;
 use rustc_abi::FIRST_VARIANT;
-use rustc_ast::expand::allocator::{ALLOCATOR_METHODS, AllocatorKind, global_fn_name};
+use rustc_ast::expand::allocator::{ALLOCATOR_METHODS, global_fn_name};
 use rustc_data_structures::fx::{FxHashMap, FxIndexSet};
 use rustc_data_structures::profiling::{get_resident_set_size, print_time_passes_entry};
 use rustc_data_structures::sync::{Lrc, par_map};
@@ -584,7 +584,7 @@ pub fn collect_debugger_visualizers_transitive(
 /// Decide allocator kind to codegen. If `Some(_)` this will be the same as
 /// `tcx.allocator_kind`, but it may be `None` in more cases (e.g. if using
 /// allocator definitions from a dylib dependency).
-pub fn allocator_kind_for_codegen(tcx: TyCtxt<'_>) -> Option<AllocatorKind> {
+pub fn needs_allocator_shim(tcx: TyCtxt<'_>) -> bool {
     // If the crate doesn't have an `allocator_kind` set then there's definitely
     // no shim to generate. Otherwise we also check our dependency graph for all
     // our output crate types. If anything there looks like its a `Dynamic`
@@ -595,7 +595,7 @@ pub fn allocator_kind_for_codegen(tcx: TyCtxt<'_>) -> Option<AllocatorKind> {
         use rustc_middle::middle::dependency_format::Linkage;
         list.iter().any(|&linkage| linkage == Linkage::Dynamic)
     });
-    if any_dynamic_crate { None } else { tcx.allocator_kind(()) }
+    if any_dynamic_crate { false } else { tcx.allocator_kind(()).is_some() }
 }
 
 pub fn codegen_crate<B: ExtraBackendMethods>(
@@ -664,14 +664,13 @@ pub fn codegen_crate<B: ExtraBackendMethods>(
         start_async_codegen(backend.clone(), tcx, target_cpu, metadata, metadata_module);
 
     // Codegen an allocator shim, if necessary.
-    if let Some(kind) = allocator_kind_for_codegen(tcx) {
+    if needs_allocator_shim(tcx) {
         let llmod_id =
             cgu_name_builder.build_cgu_name(LOCAL_CRATE, &["crate"], Some("allocator")).to_string();
         let module_llvm = tcx.sess.time("write_allocator_module", || {
             backend.codegen_allocator(
                 tcx,
                 &llmod_id,
-                kind,
                 // If allocator_kind is Some then alloc_error_handler_kind must
                 // also be Some.
                 tcx.alloc_error_handler_kind(()).unwrap(),
diff --git a/compiler/rustc_codegen_ssa/src/traits/backend.rs b/compiler/rustc_codegen_ssa/src/traits/backend.rs
@@ -104,7 +104,6 @@ pub trait ExtraBackendMethods:
         &self,
         tcx: TyCtxt<'tcx>,
         module_name: &str,
-        kind: AllocatorKind,
         alloc_error_handler_kind: AllocatorKind,
     ) -> Self::Module;
 
diff --git a/compiler/rustc_metadata/src/creader.rs b/compiler/rustc_metadata/src/creader.rs
@@ -1127,6 +1127,8 @@ fn global_allocator_spans(krate: &ast::Crate) -> Vec<Span> {
         fn visit_item(&mut self, item: &'ast ast::Item) {
             if item.ident.name == self.name
                 && attr::contains_name(&item.attrs, sym::rustc_std_internal_symbol)
+                // Ignore the default allocator in libstd with weak linkage
+                && attr::find_by_name(&item.attrs, sym::linkage).is_none()
             {
                 self.spans.push(item.span);
             }
diff --git a/compiler/rustc_monomorphize/src/partitioning.rs b/compiler/rustc_monomorphize/src/partitioning.rs
@@ -887,6 +887,7 @@ fn mono_item_visibility<'tcx>(
         //   Removal of these functions can't be done by LLVM but rather must be
         //   done by the linker as it's a non-local decision.
         //
+        // FIXME update comment
         // * Second is "std internal symbols". Currently this is primarily used
         //   for allocator symbols. Allocators are a little weird in their
         //   implementation, but the idea is that the compiler, at the last
diff --git a/library/alloc/src/alloc.rs b/library/alloc/src/alloc.rs
@@ -14,7 +14,7 @@ extern "Rust" {
     // These are the magic symbols to call the global allocator. rustc generates
     // them to call `__rg_alloc` etc. if there is a `#[global_allocator]` attribute
     // (the code expanding that attribute macro generates those functions), or to call
-    // the default implementations in std (`__rdl_alloc` etc. in `library/std/src/alloc.rs`)
+    // the default implementations in std (weak symbols in `library/std/src/alloc.rs`)
     // otherwise.
     // The rustc fork of LLVM 14 and earlier also special-cases these function names to be able to optimize them
     // like `malloc`, `realloc`, and `free`, respectively.
diff --git a/library/std/src/alloc.rs b/library/std/src/alloc.rs
@@ -382,9 +382,8 @@ pub fn rust_oom(layout: Layout) -> ! {
 #[unstable(feature = "alloc_internals", issue = "none")]
 pub mod __default_lib_allocator {
     use super::{GlobalAlloc, Layout, System};
-    // These magic symbol names are used as a fallback for implementing the
-    // `__rust_alloc` etc symbols (see `src/liballoc/alloc.rs`) when there is
-    // no `#[global_allocator]` attribute.
+    // These are used as a fallback for implementing the `__rust_alloc`, etc symbols
+    // (see `src/liballoc/alloc.rs`) when there is no `#[global_allocator]` attribute.
 
     // for symbol names src/librustc_ast/expand/allocator.rs
     // for signatures src/librustc_allocator/lib.rs
@@ -393,7 +392,8 @@ pub mod __default_lib_allocator {
     // ABI
 
     #[rustc_std_internal_symbol]
-    pub unsafe extern "C" fn __rdl_alloc(size: usize, align: usize) -> *mut u8 {
+    #[linkage = "weak"]
+    pub unsafe extern "Rust" fn __rust_alloc(size: usize, align: usize) -> *mut u8 {
         // SAFETY: see the guarantees expected by `Layout::from_size_align` and
         // `GlobalAlloc::alloc`.
         unsafe {
@@ -403,14 +403,16 @@ pub mod __default_lib_allocator {
     }
 
     #[rustc_std_internal_symbol]
-    pub unsafe extern "C" fn __rdl_dealloc(ptr: *mut u8, size: usize, align: usize) {
+    #[linkage = "weak"]
+    pub unsafe extern "Rust" fn __rust_dealloc(ptr: *mut u8, size: usize, align: usize) {
         // SAFETY: see the guarantees expected by `Layout::from_size_align` and
         // `GlobalAlloc::dealloc`.
         unsafe { System.dealloc(ptr, Layout::from_size_align_unchecked(size, align)) }
     }
 
     #[rustc_std_internal_symbol]
-    pub unsafe extern "C" fn __rdl_realloc(
+    #[linkage = "weak"]
+    pub unsafe extern "Rust" fn __rust_realloc(
         ptr: *mut u8,
         old_size: usize,
         align: usize,
@@ -425,7 +427,8 @@ pub mod __default_lib_allocator {
     }
 
     #[rustc_std_internal_symbol]
-    pub unsafe extern "C" fn __rdl_alloc_zeroed(size: usize, align: usize) -> *mut u8 {
+    #[linkage = "weak"]
+    pub unsafe extern "Rust" fn __rust_alloc_zeroed(size: usize, align: usize) -> *mut u8 {
         // SAFETY: see the guarantees expected by `Layout::from_size_align` and
         // `GlobalAlloc::alloc_zeroed`.
         unsafe {
diff --git a/src/tools/miri/src/shims/alloc.rs b/src/tools/miri/src/shims/alloc.rs
diff --git a/src/tools/miri/src/shims/foreign_items.rs b/src/tools/miri/src/shims/foreign_items.rs

Original file line number	Diff line number	Diff line change
`@@ -114,12 +114,11 @@ impl ExtraBackendMethods for LlvmCodegenBackend {`
`114`	`114`	`&self,`
`115`	`115`	`tcx: TyCtxt<'tcx>,`
`116`	`116`	`module_name: &str,`
`117`		`- kind: AllocatorKind,`
`118`	`117`	`alloc_error_handler_kind: AllocatorKind,`
`119`	`118`	`) -> ModuleLlvm {`
`120`	`119`	`let mut module_llvm = ModuleLlvm::new_metadata(tcx, module_name);`
`121`	`120`	`unsafe {`
`122`		`- allocator::codegen(tcx, &mut module_llvm, module_name, kind, alloc_error_handler_kind);`
	`121`	`+ allocator::codegen(tcx, &mut module_llvm, module_name, alloc_error_handler_kind);`
`123`	`122`	`}`
`124`	`123`	`module_llvm`
`125`	`124`	`}`
Original file line number	Diff line number	Diff line change
`@@ -1127,6 +1127,8 @@ fn global_allocator_spans(krate: &ast::Crate) -> Vec<Span> {`
`1127`	`1127`	`fn visit_item(&mut self, item: &'ast ast::Item) {`
`1128`	`1128`	`if item.ident.name == self.name`
`1129`	`1129`	`&& attr::contains_name(&item.attrs, sym::rustc_std_internal_symbol)`
	`1130`	`+ // Ignore the default allocator in libstd with weak linkage`
	`1131`	`+ && attr::find_by_name(&item.attrs, sym::linkage).is_none()`
`1130`	`1132`	`{`
`1131`	`1133`	`self.spans.push(item.span);`
`1132`	`1134`	`}`