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.

Author: bjorn3

src/allocator.rs

@@ -2,8 +2,7 @@
 use gccjit::FnAttribute;
 use gccjit::{Context, FunctionType, RValue, ToRValue, Type};
 use rustc_ast::expand::allocator::{
-    ALLOC_ERROR_HANDLER, ALLOCATOR_METHODS, AllocatorKind, AllocatorTy, NO_ALLOC_SHIM_IS_UNSTABLE,
-    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,46 +29,34 @@ 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,
 
-    if alloc_error_handler_kind == AllocatorKind::Default {
-        // FIXME(bjorn3): Add noreturn attribute
-        create_wrapper_function(
-            tcx,
-            context,
-            &mangle_internal_symbol(tcx, &global_fn_name(ALLOC_ERROR_HANDLER)),
-            Some(&mangle_internal_symbol(tcx, &default_fn_name(ALLOC_ERROR_HANDLER))),
-            &[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(

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
     }