Unrolled build for #147526

Rollup merge of #147526 - bjorn3:alloc_shim_weak_shape, r=petrochenkov,RalfJung

Move computation of allocator shim contents to cg_ssa

In the future this should make it easier to use weak symbols for the allocator shim on platforms that properly support weak symbols. And it would allow reusing the allocator shim code for handling default implementations of the upcoming externally implementable items feature on platforms that don't properly support weak symbols.

In addition to make this possible, the alloc error handler is now handled in a way such that it is possible to avoid using the allocator shim when liballoc is compiled without `no_global_oom_handling` if you use `#[alloc_error_handler]`. Previously this was only possible if you avoided liballoc entirely or compiled it with `no_global_oom_handling`. You still need to avoid libstd and to define the symbol that indicates that avoiding the allocator shim is unstable.

Author: rust-timer

compiler/rustc_ast/src/expand/allocator.rs

@@ -3,7 +3,9 @@ use rustc_span::{Symbol, sym};
 
 #[derive(Clone, Debug, Copy, Eq, PartialEq, HashStable_Generic)]
 pub enum AllocatorKind {
+    /// Use `#[global_allocator]` as global allocator.
     Global,
+    /// Use the default implementation in libstd as global allocator.
     Default,
 }
 
@@ -15,23 +17,22 @@ 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 ALLOC_ERROR_HANDLER: Symbol = sym::alloc_error_handler;
 pub const NO_ALLOC_SHIM_IS_UNSTABLE: &str = "__rust_no_alloc_shim_is_unstable_v2";
 
+/// Argument or return type for methods in the allocator shim
+#[derive(Copy, Clone)]
 pub enum AllocatorTy {
     Layout,
+    Never,
     Ptr,
     ResultPtr,
     Unit,
     Usize,
 }
 
+/// A method that will be codegened in the allocator shim.
+#[derive(Copy, Clone)]
 pub struct AllocatorMethod {
     pub name: Symbol,
     pub inputs: &'static [AllocatorMethodInput],

compiler/rustc_builtin_macros/src/alloc_error_handler.rs

@@ -1,3 +1,4 @@
+use rustc_ast::expand::allocator::{ALLOC_ERROR_HANDLER, global_fn_name};
 use rustc_ast::{
     self as ast, Fn, FnHeader, FnSig, Generics, ItemKind, Safety, Stmt, StmtKind, TyKind,
 };
@@ -55,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 {
@@ -84,7 +85,7 @@ fn generate_handler(cx: &ExtCtxt<'_>, handler: Ident, span: Span, sig_span: Span
     let kind = ItemKind::Fn(Box::new(Fn {
         defaultness: ast::Defaultness::Final,
         sig,
-        ident: Ident::from_str_and_span("__rg_oom", span),
+        ident: Ident::from_str_and_span(&global_fn_name(ALLOC_ERROR_HANDLER), span),
         generics: Generics::default(),
         contract: None,
         body,

compiler/rustc_builtin_macros/src/global_allocator.rs

@@ -151,7 +151,7 @@ impl AllocFnFactory<'_, '_> {
                 self.cx.expr_ident(self.span, ident)
             }
 
-            AllocatorTy::ResultPtr | AllocatorTy::Unit => {
+            AllocatorTy::Never | AllocatorTy::ResultPtr | AllocatorTy::Unit => {
                 panic!("can't convert AllocatorTy to an argument")
             }
         }
@@ -163,7 +163,7 @@ impl AllocFnFactory<'_, '_> {
 
             AllocatorTy::Unit => self.cx.ty(self.span, TyKind::Tup(ThinVec::new())),
 
-            AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
+            AllocatorTy::Layout | AllocatorTy::Never | AllocatorTy::Usize | AllocatorTy::Ptr => {
                 panic!("can't convert `AllocatorTy` to an output")
             }
         }

compiler/rustc_codegen_cranelift/src/allocator.rs

@@ -3,10 +3,9 @@
 
 use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
 use rustc_ast::expand::allocator::{
-    ALLOCATOR_METHODS, AllocatorKind, AllocatorTy, NO_ALLOC_SHIM_IS_UNSTABLE,
-    alloc_error_handler_name, default_fn_name, global_fn_name,
+    AllocatorMethod, AllocatorTy, NO_ALLOC_SHIM_IS_UNSTABLE, default_fn_name, global_fn_name,
 };
-use rustc_codegen_ssa::base::allocator_kind_for_codegen;
+use rustc_codegen_ssa::base::{allocator_kind_for_codegen, allocator_shim_contents};
 use rustc_session::config::OomStrategy;
 use rustc_symbol_mangling::mangle_internal_symbol;
 
@@ -15,75 +14,57 @@ 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(
-        tcx,
-        module,
-        kind,
-        tcx.alloc_error_handler_kind(()).unwrap(),
-        tcx.sess.opts.unstable_opts.oom,
-    );
+    let methods = allocator_shim_contents(tcx, kind);
+    codegen_inner(tcx, module, &methods, tcx.sess.opts.unstable_opts.oom);
     true
 }
 
 fn codegen_inner(
     tcx: TyCtxt<'_>,
     module: &mut dyn Module,
-    kind: AllocatorKind,
-    alloc_error_handler_kind: AllocatorKind,
+    methods: &[AllocatorMethod],
     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),
+    for method in 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"),
+                AllocatorTy::Never | AllocatorTy::ResultPtr | AllocatorTy::Unit => {
+                    panic!("invalid allocator arg")
                 }
             }
-            let output = match method.output {
-                AllocatorTy::ResultPtr => Some(usize_ty),
-                AllocatorTy::Unit => None,
+        }
+        let output = match method.output {
+            AllocatorTy::ResultPtr => Some(usize_ty),
+            AllocatorTy::Never | AllocatorTy::Unit => None,
 
-                AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
-                    panic!("invalid allocator output")
-                }
-            };
+            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,
-                &mangle_internal_symbol(tcx, &global_fn_name(method.name)),
-                &mangle_internal_symbol(tcx, &default_fn_name(method.name)),
-            );
-        }
+        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,
+            &mangle_internal_symbol(tcx, &global_fn_name(method.name)),
+            &mangle_internal_symbol(tcx, &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)],
-        returns: vec![],
-    };
-    crate::common::create_wrapper_function(
-        module,
-        sig,
-        &mangle_internal_symbol(tcx, "__rust_alloc_error_handler"),
-        &mangle_internal_symbol(tcx, alloc_error_handler_name(alloc_error_handler_kind)),
-    );
-
     {
         let sig = Signature {
             call_conv: module.target_config().default_call_conv,

compiler/rustc_codegen_gcc/src/allocator.rs

@@ -2,8 +2,7 @@
 use gccjit::FnAttribute;
 use gccjit::{Context, FunctionType, RValue, 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,
+    AllocatorMethod, AllocatorTy, NO_ALLOC_SHIM_IS_UNSTABLE, default_fn_name, global_fn_name,
 };
 use rustc_middle::bug;
 use rustc_middle::ty::TyCtxt;
@@ -18,8 +17,7 @@ pub(crate) unsafe fn codegen(
     tcx: TyCtxt<'_>,
     mods: &mut GccContext,
     _module_name: &str,
-    kind: AllocatorKind,
-    alloc_error_handler_kind: AllocatorKind,
+    methods: &[AllocatorMethod],
 ) {
     let context = &mods.context;
     let usize = match tcx.sess.target.pointer_width {
@@ -31,45 +29,35 @@ pub(crate) unsafe fn codegen(
     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"),
+    for method in 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);
                 }
-            }
-            let output = match method.output {
-                AllocatorTy::ResultPtr => Some(i8p),
-                AllocatorTy::Unit => None,
+                AllocatorTy::Ptr => types.push(i8p),
+                AllocatorTy::Usize => types.push(usize),
 
-                AllocatorTy::Layout | AllocatorTy::Usize | AllocatorTy::Ptr => {
-                    panic!("invalid allocator output")
+                AllocatorTy::Never | AllocatorTy::ResultPtr | AllocatorTy::Unit => {
+                    panic!("invalid allocator arg")
                 }
-            };
-            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, context, &from_name, Some(&to_name), &types, output);
+            }
         }
-    }
+        let output = match method.output {
+            AllocatorTy::ResultPtr => Some(i8p),
+            AllocatorTy::Never | AllocatorTy::Unit => None,
 
-    // FIXME(bjorn3): Add noreturn attribute
-    create_wrapper_function(
-        tcx,
-        context,
-        &mangle_internal_symbol(tcx, "__rust_alloc_error_handler"),
-        Some(&mangle_internal_symbol(tcx, alloc_error_handler_name(alloc_error_handler_kind))),
-        &[usize, usize],
-        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, context, &from_name, Some(&to_name), &types, output);
+    }
 
     create_const_value_function(
         tcx,

compiler/rustc_codegen_gcc/src/lib.rs

@@ -92,7 +92,7 @@ use back::lto::{ThinBuffer, ThinData};
 use gccjit::{CType, Context, OptimizationLevel};
 #[cfg(feature = "master")]
 use gccjit::{TargetInfo, Version};
-use rustc_ast::expand::allocator::AllocatorKind;
+use rustc_ast::expand::allocator::AllocatorMethod;
 use rustc_codegen_ssa::back::lto::{SerializedModule, ThinModule};
 use rustc_codegen_ssa::back::write::{
     CodegenContext, FatLtoInput, ModuleConfig, TargetMachineFactoryFn,
@@ -284,8 +284,7 @@ impl ExtraBackendMethods for GccCodegenBackend {
         &self,
         tcx: TyCtxt<'_>,
         module_name: &str,
-        kind: AllocatorKind,
-        alloc_error_handler_kind: AllocatorKind,
+        methods: &[AllocatorMethod],
     ) -> Self::Module {
         let mut mods = GccContext {
             context: Arc::new(SyncContext::new(new_context(tcx))),
@@ -295,7 +294,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, methods);
         }
         mods
     }