Rewrite implementation of `#[alloc_error_handler]`

The new implementation doesn't use weak lang items and instead changes
`#[alloc_error_handler]` to an attribute macro just like
`#[global_allocator]`.

The attribute will generate the `__rg_oom` function which is called by
the compiler-generated `__rust_alloc_error_handler`. If no `__rg_oom`
function is defined in any crate then the compiler shim will call
`__rdl_oom` in the alloc crate which will simply panic.

This also fixes link errors with `-C link-dead-code` with
`default_alloc_error_handler`: `__rg_oom` was previously defined in the
alloc crate and would attempt to reference the `oom` lang item, even if
it didn't exist. This worked as long as `__rg_oom` was excluded from
linking since it was not called.

This is a prerequisite for the stabilization of
`default_alloc_error_handler` (#102318).

1 files changed, 5 insertions, 4 deletions

diff --git a/src/allocator.rs b/src/allocator.rs
index bad8a87b9be..12bb00d346d 100644
--- a/src/allocator.rs
+++ b/src/allocator.rs
@@ -5,6 +5,7 @@ use crate::prelude::*;
 
 use rustc_ast::expand::allocator::{AllocatorKind, AllocatorTy, ALLOCATOR_METHODS};
 use rustc_session::config::OomStrategy;
+use rustc_span::symbol::sym;
 
 /// Returns whether an allocator shim was created
 pub(crate) fn codegen(
@@ -23,7 +24,7 @@ pub(crate) fn codegen(
             module,
             unwind_context,
             kind,
-            tcx.lang_items().oom().is_some(),
+            tcx.alloc_error_handler_kind(()).unwrap(),
             tcx.sess.opts.unstable_opts.oom,
         );
         true
@@ -36,7 +37,7 @@ fn codegen_inner(
     module: &mut impl Module,
     unwind_context: &mut UnwindContext,
     kind: AllocatorKind,
-    has_alloc_error_handler: bool,
+    alloc_error_handler_kind: AllocatorKind,
     oom_strategy: OomStrategy,
 ) {
     let usize_ty = module.target_config().pointer_type();
@@ -108,12 +109,12 @@ fn codegen_inner(
         returns: vec![],
     };
 
-    let callee_name = if has_alloc_error_handler { "__rg_oom" } else { "__rdl_oom" };
+    let callee_name = alloc_error_handler_kind.fn_name(sym::oom);
 
     let func_id =
         module.declare_function("__rust_alloc_error_handler", Linkage::Export, &sig).unwrap();
 
-    let callee_func_id = module.declare_function(callee_name, Linkage::Import, &sig).unwrap();
+    let callee_func_id = module.declare_function(&callee_name, Linkage::Import, &sig).unwrap();
 
     let mut ctx = Context::new();
     ctx.func.signature = sig;