Rollup merge of rust-lang#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: matthiaskrgr

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,