[WIP] Use weak aliases for the allocator shim

bjorn3 · bjorn3 · commit d4526ac3090e · 2025-06-17T11:42:56.000Z
diff --git a/compiler/rustc_codegen_gcc/src/mono_item.rs b/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) {}
 }
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_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(
diff --git a/compiler/rustc_codegen_llvm/src/lib.rs b/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
     }
diff --git a/compiler/rustc_codegen_llvm/src/llvm/ffi.rs b/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;
 }
diff --git a/compiler/rustc_codegen_llvm/src/llvm/mod.rs b/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
diff --git a/compiler/rustc_codegen_llvm/src/mono_item.rs b/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<'_, '_> {
diff --git a/compiler/rustc_codegen_ssa/src/traits/declare.rs b/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);
 }
diff --git a/library/alloc/src/alloc.rs b/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 {

Original file line number	Diff line number	Diff line change
`@@ -76,4 +76,6 @@ impl<'gcc, 'tcx> PreDefineCodegenMethods<'tcx> for CodegenCx<'gcc, 'tcx> {`
`76`	`76`	`self.functions.borrow_mut().insert(symbol_name.to_string(), decl);`
`77`	`77`	`self.function_instances.borrow_mut().insert(instance, decl);`
`78`	`78`	`}`
	`79`	`+`
	`80`	`+ fn weak_alias(&self, _aliasee: Self::Function, _aliasee_name: &str, _name: &str) {}`
`79`	`81`	`}`