[WIP] Use weak aliases for the allocator shim

compiler/rustc_codegen_gcc/src/mono_item.rs

@@ -76,4 +76,6 @@ impl<'gcc, 'tcx> PreDefineCodegenMethods<'tcx> for CodegenCx<'gcc, 'tcx> {
         self.functions.borrow_mut().insert(symbol_name.to_string(), decl);
         self.function_instances.borrow_mut().insert(instance, decl);
     }
+
+    fn weak_alias(&self, _aliasee: Self::Function, _aliasee_name: &str, _name: &str) {}
 }

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_codegen_ssa::traits::BaseTypeCodegenMethods as _;
 use rustc_middle::bug;
@@ -18,7 +17,6 @@ pub(crate) unsafe fn codegen(
     tcx: TyCtxt<'_>,
     cx: SimpleCx<'_>,
     module_name: &str,
-    kind: AllocatorKind,
     alloc_error_handler_kind: AllocatorKind,
 ) {
     let usize = match tcx.sess.target.pointer_width {
@@ -28,38 +26,6 @@ pub(crate) unsafe fn codegen(
         tws => bug!("Unsupported target word size for int: {}", tws),
     };
     let i8 = cx.type_i8();
-    let i8p = cx.type_ptr();
-
-    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 = mangle_internal_symbol(tcx, &global_fn_name(method.name));
-            let to_name = mangle_internal_symbol(tcx, &default_fn_name(method.name));
-
-            create_wrapper_function(tcx, &cx, &from_name, &to_name, &args, output, false);
-        }
-    }
 
     // rust alloc error handler
     create_wrapper_function(

compiler/rustc_codegen_llvm/src/lib.rs

@@ -108,14 +108,14 @@ impl ExtraBackendMethods for LlvmCodegenBackend {
         &self,
         tcx: TyCtxt<'tcx>,
         module_name: &str,
-        kind: AllocatorKind,
+        _kind: AllocatorKind,
         alloc_error_handler_kind: AllocatorKind,
     ) -> ModuleLlvm {
         let module_llvm = ModuleLlvm::new_metadata(tcx, module_name);
         let cx =
             SimpleCx::new(module_llvm.llmod(), &module_llvm.llcx, tcx.data_layout.pointer_size);
         unsafe {
-            allocator::codegen(tcx, cx, module_name, kind, alloc_error_handler_kind);
+            allocator::codegen(tcx, cx, module_name, alloc_error_handler_kind);
         }
         module_llvm
     }

compiler/rustc_codegen_llvm/src/llvm/ffi.rs

@@ -2670,4 +2670,12 @@ unsafe extern "C" {
 
     pub(crate) fn LLVMRustSetNoSanitizeAddress(Global: &Value);
     pub(crate) fn LLVMRustSetNoSanitizeHWAddress(Global: &Value);
+
+    pub(crate) fn LLVMAddAlias2<'ll>(
+        M: &'ll Module,
+        ValueTy: &Type,
+        AddressSpace: c_uint,
+        Aliasee: &Value,
+        Name: *const c_char,
+    ) -> &'ll Value;
 }

compiler/rustc_codegen_llvm/src/llvm/mod.rs

@@ -7,7 +7,7 @@ use std::str::FromStr;
 use std::string::FromUtf8Error;
 
 use libc::c_uint;
-use rustc_abi::{Align, Size, WrappingRange};
+use rustc_abi::{AddressSpace, Align, Size, WrappingRange};
 use rustc_llvm::RustString;
 
 pub(crate) use self::CallConv::*;
@@ -350,6 +350,17 @@ impl Intrinsic {
     }
 }
 
+/// Safe wrapper for `LLVMAddAlias2`
+pub(crate) fn add_alias<'ll>(
+    module: &'ll Module,
+    ty: &Type,
+    address_space: AddressSpace,
+    aliasee: &Value,
+    name: &CStr,
+) -> &'ll Value {
+    unsafe { LLVMAddAlias2(module, ty, address_space.0, aliasee, name.as_ptr()) }
+}
+
 /// Safe wrapper for `LLVMSetValueName2` from a byte slice
 pub(crate) fn set_value_name(value: &Value, name: &[u8]) {
     unsafe {

compiler/rustc_codegen_llvm/src/mono_item.rs

@@ -1,3 +1,6 @@
+use std::ffi::CString;
+
+use rustc_abi::AddressSpace;
 use rustc_codegen_ssa::traits::*;
 use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LOCAL_CRATE};
@@ -6,6 +9,7 @@ use rustc_middle::mir::mono::{Linkage, Visibility};
 use rustc_middle::ty::layout::{FnAbiOf, HasTypingEnv, LayoutOf};
 use rustc_middle::ty::{self, Instance, TypeVisitableExt};
 use rustc_session::config::CrateType;
+use rustc_symbol_mangling::mangle_internal_symbol;
 use rustc_target::spec::RelocModel;
 use tracing::debug;
 
@@ -77,8 +81,51 @@ impl<'tcx> PreDefineCodegenMethods<'tcx> for CodegenCx<'_, 'tcx> {
 
         self.assume_dso_local(lldecl, false);
 
+        let symbol_name = self.tcx.symbol_name(instance);
+        if symbol_name.name.contains("__rdl_alloc") {
+            self.weak_alias(
+                lldecl,
+                symbol_name.name,
+                &mangle_internal_symbol(self.tcx, "__rust_alloc"),
+            );
+        }
+        if symbol_name.name.contains("__rdl_dealloc") {
+            self.weak_alias(
+                lldecl,
+                symbol_name.name,
+                &mangle_internal_symbol(self.tcx, "__rust_dealloc"),
+            );
+        }
+        if symbol_name.name.contains("__rdl_realloc") {
+            self.weak_alias(
+                lldecl,
+                symbol_name.name,
+                &mangle_internal_symbol(self.tcx, "__rust_realloc"),
+            );
+        }
+        if symbol_name.name.contains("__rdl_alloc_zeroed") {
+            self.weak_alias(
+                lldecl,
+                symbol_name.name,
+                &mangle_internal_symbol(self.tcx, "__rust_alloc_zeroed"),
+            );
+        }
+
         self.instances.borrow_mut().insert(instance, lldecl);
     }
+
+    fn weak_alias(&self, aliasee: Self::Function, _aliasee_name: &str, name: &str) {
+        let ty = self.get_type_of_global(aliasee);
+        let alias = llvm::add_alias(
+            self.llmod,
+            ty,
+            AddressSpace::DATA,
+            aliasee,
+            &CString::new(name).unwrap(),
+        );
+
+        llvm::set_linkage(alias, llvm::Linkage::WeakAnyLinkage);
+    }
 }
 
 impl CodegenCx<'_, '_> {

compiler/rustc_codegen_ssa/src/back/link.rs

@@ -1976,9 +1976,11 @@ fn add_linked_symbol_object(
     cmd: &mut dyn Linker,
     sess: &Session,
     tmpdir: &Path,
-    symbols: &[(String, SymbolExportKind)],
+    crate_type: CrateType,
+    linked_symbols: &[(String, SymbolExportKind)],
+    exported_symbols: &[String],
 ) {
-    if symbols.is_empty() {
+    if linked_symbols.is_empty() && exported_symbols.is_empty() {
         return;
     }
 
@@ -2015,7 +2017,7 @@ fn add_linked_symbol_object(
         None
     };
 
-    for (sym, kind) in symbols.iter() {
+    for (sym, kind) in linked_symbols.iter() {
         let symbol = file.add_symbol(object::write::Symbol {
             name: sym.clone().into(),
             value: 0,
@@ -2073,6 +2075,38 @@ fn add_linked_symbol_object(
         }
     }
 
+    if sess.target.is_like_msvc {
+        // Symbol visibility takes care of this for executables typically
+        let should_filter_symbols = if crate_type == CrateType::Executable {
+            sess.opts.unstable_opts.export_executable_symbols
+        } else {
+            true
+        };
+        if should_filter_symbols {
+            // Currently the compiler doesn't use `dllexport` (an LLVM attribute) to
+            // export symbols from a dynamic library. When building a dynamic library,
+            // however, we're going to want some symbols exported, so this adds a
+            // `.drectve` section which lists all the symbols using /EXPORT arguments.
+            //
+            // The linker will read these arguments from the `.drectve` section and
+            // export all the symbols from the dynamic library. Note that this is not
+            // as simple as just exporting all the symbols in the current crate (as
+            // specified by `codegen.reachable`) but rather we also need to possibly
+            // export the symbols of upstream crates. Upstream rlibs may be linked
+            // statically to this dynamic library, in which case they may continue to
+            // transitively be used and hence need their symbols exported.
+            let drectve = exported_symbols
+                .into_iter()
+                .map(|sym| format!(" /EXPORT:\"{sym}\""))
+                .collect::<Vec<_>>()
+                .join("");
+
+            let section =
+                file.add_section(vec![], b".drectve".to_vec(), object::SectionKind::Linker);
+            file.append_section_data(section, drectve.as_bytes(), 1);
+        }
+    }
+
     let path = tmpdir.join("symbols.o");
     let result = std::fs::write(&path, file.write().unwrap());
     if let Err(error) = result {
@@ -2247,7 +2281,9 @@ fn linker_with_args(
         cmd,
         sess,
         tmpdir,
+        crate_type,
         &codegen_results.crate_info.linked_symbols[&crate_type],
+        &codegen_results.crate_info.exported_symbols[&crate_type],
     );
 
     // Sanitizer libraries.

compiler/rustc_codegen_ssa/src/back/linker.rs

@@ -1086,19 +1086,10 @@ impl<'a> Linker for MsvcLinker<'a> {
         }
     }
 
-    // Currently the compiler doesn't use `dllexport` (an LLVM attribute) to
-    // export symbols from a dynamic library. When building a dynamic library,
-    // however, we're going to want some symbols exported, so this function
-    // generates a DEF file which lists all the symbols.
-    //
-    // The linker will read this `*.def` file and export all the symbols from
-    // the dynamic library. Note that this is not as simple as just exporting
-    // all the symbols in the current crate (as specified by `codegen.reachable`)
-    // but rather we also need to possibly export the symbols of upstream
-    // crates. Upstream rlibs may be linked statically to this dynamic library,
-    // in which case they may continue to transitively be used and hence need
-    // their symbols exported.
-    fn export_symbols(&mut self, tmpdir: &Path, crate_type: CrateType, symbols: &[String]) {
+    fn export_symbols(&mut self, tmpdir: &Path, crate_type: CrateType, _symbols: &[String]) {
+        // We already add /EXPORT arguments to the .drectve section of symbols.o. We generate
+        // a .DEF file here anyway as it might prevent auto-export of some symbols.
+
         // Symbol visibility takes care of this typically
         if crate_type == CrateType::Executable {
             let should_export_executable_symbols =
@@ -1116,10 +1107,6 @@ impl<'a> Linker for MsvcLinker<'a> {
             // straight to exports.
             writeln!(f, "LIBRARY")?;
             writeln!(f, "EXPORTS")?;
-            for symbol in symbols {
-                debug!("  _{symbol}");
-                writeln!(f, "  {symbol}")?;
-            }
         };
         if let Err(error) = res {
             self.sess.dcx().emit_fatal(errors::LibDefWriteFailure { error });

compiler/rustc_codegen_ssa/src/traits/declare.rs

@@ -2,7 +2,9 @@ use rustc_hir::def_id::DefId;
 use rustc_middle::mir::mono::{Linkage, Visibility};
 use rustc_middle::ty::Instance;
 
-pub trait PreDefineCodegenMethods<'tcx> {
+use crate::traits::BackendTypes;
+
+pub trait PreDefineCodegenMethods<'tcx>: BackendTypes {
     fn predefine_static(
         &mut self,
         def_id: DefId,
@@ -17,4 +19,6 @@ pub trait PreDefineCodegenMethods<'tcx> {
         visibility: Visibility,
         symbol_name: &str,
     );
+
+    fn weak_alias(&self, aliasee: Self::Function, aliasee_name: &str, name: &str);
 }

library/alloc/src/alloc.rs

@@ -31,8 +31,6 @@ unsafe extern "Rust" {
     #[rustc_std_internal_symbol]
     fn __rust_alloc_zeroed(size: usize, align: usize) -> *mut u8;
 
-    #[rustc_std_internal_symbol]
-    static __rust_no_alloc_shim_is_unstable: u8;
 }
 
 /// The global memory allocator.
@@ -85,13 +83,7 @@ pub struct Global;
 #[inline]
 #[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
 pub unsafe fn alloc(layout: Layout) -> *mut u8 {
-    unsafe {
-        // Make sure we don't accidentally allow omitting the allocator shim in
-        // stable code until it is actually stabilized.
-        core::ptr::read_volatile(&__rust_no_alloc_shim_is_unstable);
-
-        __rust_alloc(layout.size(), layout.align())
-    }
+    unsafe { __rust_alloc(layout.size(), layout.align()) }
 }
 
 /// Deallocates memory with the global allocator.
@@ -168,13 +160,7 @@ pub unsafe fn realloc(ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8
 #[inline]
 #[cfg_attr(miri, track_caller)] // even without panics, this helps for Miri backtraces
 pub unsafe fn alloc_zeroed(layout: Layout) -> *mut u8 {
-    unsafe {
-        // Make sure we don't accidentally allow omitting the allocator shim in
-        // stable code until it is actually stabilized.
-        core::ptr::read_volatile(&__rust_no_alloc_shim_is_unstable);
-
-        __rust_alloc_zeroed(layout.size(), layout.align())
-    }
+    unsafe { __rust_alloc_zeroed(layout.size(), layout.align()) }
 }
 
 impl Global {