[WIP] Use weak aliases for the allocator shim

Author: bjorn3

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

@@ -2671,11 +2671,11 @@ unsafe extern "C" {
     pub(crate) fn LLVMRustSetNoSanitizeAddress(Global: &Value);
     pub(crate) fn LLVMRustSetNoSanitizeHWAddress(Global: &Value);
 
-    pub(crate) fn LLVMAddAlias2(
-        M: &Module,
+    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

@@ -351,17 +351,14 @@ impl Intrinsic {
 }
 
 /// Safe wrapper for `LLVMAddAlias2`
-pub(crate) fn add_alias(
-    module: &Module,
+pub(crate) fn add_alias<'ll>(
+    module: &'ll Module,
     ty: &Type,
     address_space: AddressSpace,
     aliasee: &Value,
-    name: &[u8],
-) {
-    unsafe {
-        let data = name.as_c_char_ptr();
-        LLVMAddAlias2(value, data, address_space, aliasee, name.len());
-    }
+    name: &CStr,
+) -> &'ll Value {
+    unsafe { LLVMAddAlias2(module, ty, address_space.0, aliasee, name.as_ptr()) }
 }
 
 /// Safe wrapper for `LLVMSetValueName2` from a byte slice

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,10 @@ fn add_linked_symbol_object(
     cmd: &mut dyn Linker,
     sess: &Session,
     tmpdir: &Path,
-    symbols: &[(String, SymbolExportKind)],
+    linked_symbols: &[(String, SymbolExportKind)],
+    exported_symbols: &[(String, SymbolExportKind)],
 ) {
-    if symbols.is_empty() {
+    if linked_symbols.is_empty() {
         return;
     }
 
@@ -2015,7 +2016,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 +2074,23 @@ fn add_linked_symbol_object(
         }
     }
 
+    if sess.target.is_like_msvc {
+        let drectve = exported_symbols
+            .into_iter()
+            .map(|(sym, kind)| {
+                if *kind == SymbolExportKind::Text {
+                    format!(" /EXPORT:\"{sym}\"")
+                } else {
+                    format!(" /EXPORT:\"{sym}\",DATA")
+                }
+            })
+            .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 {
@@ -2248,6 +2266,7 @@ fn linker_with_args(
         sess,
         tmpdir,
         &codegen_results.crate_info.linked_symbols[&crate_type],
+        &codegen_results.crate_info.exported_symbols[&crate_type],
     );
 
     // Sanitizer libraries.