Auto merge of #134522 - bjorn3:weak_linkage_allocator, r=<try>

[WIP] Use weak linkage instead of compiler generated shims

#86844 (unstably) made it possible to avoid the allocator shim when using `#[global_allocator]`. This PR makes it possible to also avoid the allocator shim when using the default allocator in libstd by making use of weak linkage. Eventual stabilization of avoiding the allocator shim may be blocked on it working with the default allocator too.

This is still keeping the `__rust_no_alloc_shim_is_unstable` symbol for now until #123015 gets stabilized.

TODO: Update comments everywhere, test on macOS and Windows and write a better PR description why we want this. Also prevent codegen of weak symbols when there is a non-weak definition of the same symbol in the same codegen unit.

try-zzz-job: x86_64-msvc
try-job: x86_64-apple-1
try-job: x86_64-apple-2

Author: bors

compiler/rustc_ast/src/expand/allocator.rs

@@ -11,17 +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",
-        AllocatorKind::Default => "__rdl_oom",
-    }
-}
-
 pub const NO_ALLOC_SHIM_IS_UNSTABLE: &str = "__rust_no_alloc_shim_is_unstable";
 
 pub enum AllocatorTy {

compiler/rustc_builtin_macros/src/alloc_error_handler.rs

@@ -56,7 +56,7 @@ pub(crate) fn expand(
 }
 
 // #[rustc_std_internal_symbol]
-// unsafe fn __rg_oom(size: usize, align: usize) -> ! {
+// unsafe fn __rust_alloc_error_handler(size: usize, align: usize) -> ! {
 //     handler(core::alloc::Layout::from_size_align_unchecked(size, align))
 // }
 fn generate_handler(cx: &ExtCtxt<'_>, handler: Ident, span: Span, sig_span: Span) -> Stmt {
@@ -90,6 +90,7 @@ fn generate_handler(cx: &ExtCtxt<'_>, handler: Ident, span: Span, sig_span: Span
 
     let attrs = thin_vec![cx.attr_word(sym::rustc_std_internal_symbol, span)];
 
-    let item = cx.item(span, Ident::from_str_and_span("__rg_oom", span), attrs, kind);
+    let item =
+        cx.item(span, Ident::from_str_and_span("__rust_alloc_error_handler", span), attrs, kind);
     cx.stmt_item(sig_span, item)
 }

compiler/rustc_codegen_cranelift/src/allocator.rs

@@ -1,85 +1,23 @@
 //! Allocator shim
 // 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,
-};
-use rustc_codegen_ssa::base::allocator_kind_for_codegen;
+use rustc_ast::expand::allocator::NO_ALLOC_SHIM_IS_UNSTABLE;
+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
-}
-
-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),
-            );
-        }
+    if needs_allocator_shim(tcx) {
+        codegen_inner(module, tcx.sess.opts.unstable_opts.oom);
+        true
+    } else {
+        false
     }
+}
 
-    let sig = Signature {
-        call_conv: module.target_config().default_call_conv,
-        params: vec![AbiParam::new(usize_ty), AbiParam::new(usize_ty)],
-        returns: vec![],
-    };
-    crate::common::create_wrapper_function(
-        module,
-        sig,
-        "__rust_alloc_error_handler",
-        &alloc_error_handler_name(alloc_error_handler_kind),
-    );
-
+fn codegen_inner(module: &mut dyn Module, oom_strategy: OomStrategy) {
     let data_id = module.declare_data(OomStrategy::SYMBOL, Linkage::Export, false, false).unwrap();
     let mut data = DataDescription::new();
     data.set_align(1);

compiler/rustc_codegen_gcc/src/allocator.rs

@@ -1,72 +1,13 @@
-#[cfg(feature = "master")]
-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,
-};
-use rustc_middle::bug;
+use gccjit::GlobalKind;
+use rustc_ast::expand::allocator::NO_ALLOC_SHIM_IS_UNSTABLE;
 use rustc_middle::ty::TyCtxt;
 use rustc_session::config::OomStrategy;
 
 use crate::GccContext;
 
-pub(crate) unsafe fn codegen(
-    tcx: TyCtxt<'_>,
-    mods: &mut GccContext,
-    _module_name: &str,
-    kind: AllocatorKind,
-    alloc_error_handler_kind: AllocatorKind,
-) {
+pub(crate) unsafe fn codegen(tcx: TyCtxt<'_>, mods: &mut GccContext, _module_name: &str) {
     let context = &mods.context;
-    let usize = match tcx.sess.target.pointer_width {
-        16 => context.new_type::<u16>(),
-        32 => context.new_type::<u32>(),
-        64 => context.new_type::<u64>(),
-        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(
-        tcx,
-        context,
-        "__rust_alloc_error_handler",
-        alloc_error_handler_name(alloc_error_handler_kind),
-        &[usize, usize],
-        None,
-    );
 
     let name = OomStrategy::SYMBOL.to_string();
     let global = context.new_global(None, GlobalKind::Exported, i8, name);
@@ -79,77 +20,3 @@ pub(crate) unsafe fn codegen(
     let value = context.new_rvalue_from_int(i8, 0);
     global.global_set_initializer_rvalue(value);
 }
-
-fn create_wrapper_function(
-    tcx: TyCtxt<'_>,
-    context: &Context<'_>,
-    from_name: &str,
-    to_name: &str,
-    types: &[Type<'_>],
-    output: Option<Type<'_>>,
-) {
-    let void = context.new_type::<()>();
-
-    let args: Vec<_> = types
-        .iter()
-        .enumerate()
-        .map(|(index, typ)| context.new_parameter(None, *typ, format!("param{}", index)))
-        .collect();
-    let func = context.new_function(
-        None,
-        FunctionType::Exported,
-        output.unwrap_or(void),
-        &args,
-        from_name,
-        false,
-    );
-
-    #[cfg(feature = "master")]
-    match tcx.sess.default_visibility() {
-        rustc_target::spec::SymbolVisibility::Hidden => {
-            func.add_attribute(FnAttribute::Visibility(gccjit::Visibility::Hidden))
-        }
-        rustc_target::spec::SymbolVisibility::Protected => {
-            func.add_attribute(FnAttribute::Visibility(gccjit::Visibility::Protected))
-        }
-        rustc_target::spec::SymbolVisibility::Interposable => {}
-    }
-
-    if tcx.sess.must_emit_unwind_tables() {
-        // TODO(antoyo): emit unwind tables.
-    }
-
-    let args: Vec<_> = types
-        .iter()
-        .enumerate()
-        .map(|(index, typ)| context.new_parameter(None, *typ, format!("param{}", index)))
-        .collect();
-    let callee = context.new_function(
-        None,
-        FunctionType::Extern,
-        output.unwrap_or(void),
-        &args,
-        to_name,
-        false,
-    );
-    #[cfg(feature = "master")]
-    callee.add_attribute(FnAttribute::Visibility(gccjit::Visibility::Hidden));
-
-    let block = func.new_block("entry");
-
-    let args = args
-        .iter()
-        .enumerate()
-        .map(|(i, _)| func.get_param(i as i32).to_rvalue())
-        .collect::<Vec<_>>();
-    let ret = context.new_call(None, callee, &args);
-    //llvm::LLVMSetTailCall(ret, True);
-    if output.is_some() {
-        block.end_with_return(None, ret);
-    } else {
-        block.end_with_void_return(None);
-    }
-
-    // TODO(@Commeownist): Check if we need to emit some extra debugging info in certain circumstances
-    // as described in https://github.com/rust-lang/rust/commit/77a96ed5646f7c3ee8897693decc4626fe380643
-}

compiler/rustc_codegen_gcc/src/lib.rs

@@ -92,7 +92,6 @@ use errors::LTONotSupported;
 use gccjit::{CType, Context, OptimizationLevel};
 #[cfg(feature = "master")]
 use gccjit::{TargetInfo, Version};
-use rustc_ast::expand::allocator::AllocatorKind;
 use rustc_codegen_ssa::back::lto::{LtoModuleCodegen, SerializedModule, ThinModule};
 use rustc_codegen_ssa::back::write::{
     CodegenContext, FatLtoInput, ModuleConfig, TargetMachineFactoryFn,
@@ -287,21 +286,15 @@ fn new_context<'gcc, 'tcx>(tcx: TyCtxt<'tcx>) -> Context<'gcc> {
 }
 
 impl ExtraBackendMethods for GccCodegenBackend {
-    fn codegen_allocator(
-        &self,
-        tcx: TyCtxt<'_>,
-        module_name: &str,
-        kind: AllocatorKind,
-        alloc_error_handler_kind: AllocatorKind,
-    ) -> Self::Module {
+    fn codegen_allocator(&self, tcx: TyCtxt<'_>, module_name: &str) -> Self::Module {
         let mut mods = GccContext {
             context: Arc::new(SyncContext::new(new_context(tcx))),
             should_combine_object_files: false,
             temp_dir: None,
         };
 
         unsafe {
-            allocator::codegen(tcx, &mut mods, module_name, kind, alloc_error_handler_kind);
+            allocator::codegen(tcx, &mut mods, module_name);
         }
         mods
     }