Move computation of allocator shim contents to cg_ssa

compiler/rustc_ast/src/expand/allocator.rs

@@ -15,23 +15,21 @@ 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,
 }
 
+#[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
     }