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).

compiler/rustc_metadata/src/creader.rs

@@ -1,8 +1,10 @@
 //! Validates all used crates and extern libraries and loads their metadata
 
 use crate::errors::{
-    ConflictingGlobalAlloc, CrateNotPanicRuntime, GlobalAllocRequired, NoMultipleGlobalAlloc,
-    NoPanicStrategy, NoTransitiveNeedsDep, NotProfilerRuntime, ProfilerBuiltinsNeedsCore,
+    AllocFuncRequired, ConflictingAllocErrorHandler, ConflictingGlobalAlloc, CrateNotPanicRuntime,
+    GlobalAllocRequired, MissingAllocErrorHandler, NoMultipleAllocErrorHandler,
+    NoMultipleGlobalAlloc, NoPanicStrategy, NoTransitiveNeedsDep, NotProfilerRuntime,
+    ProfilerBuiltinsNeedsCore,
 };
 use crate::locator::{CrateError, CrateLocator, CratePaths};
 use crate::rmeta::{CrateDep, CrateMetadata, CrateNumMap, CrateRoot, MetadataBlob};
@@ -41,8 +43,13 @@ pub struct CStore {
     /// This crate needs an allocator and either provides it itself, or finds it in a dependency.
     /// If the above is true, then this field denotes the kind of the found allocator.
     allocator_kind: Option<AllocatorKind>,
+    /// This crate needs an allocation error handler and either provides it itself, or finds it in a dependency.
+    /// If the above is true, then this field denotes the kind of the found allocator.
+    alloc_error_handler_kind: Option<AllocatorKind>,
     /// This crate has a `#[global_allocator]` item.
     has_global_allocator: bool,
+    /// This crate has a `#[alloc_error_handler]` item.
+    has_alloc_error_handler: bool,
 
     /// This map is used to verify we get no hash conflicts between
     /// `StableCrateId` values.
@@ -197,10 +204,18 @@ impl CStore {
         self.allocator_kind
     }
 
+    pub(crate) fn alloc_error_handler_kind(&self) -> Option<AllocatorKind> {
+        self.alloc_error_handler_kind
+    }
+
     pub(crate) fn has_global_allocator(&self) -> bool {
         self.has_global_allocator
     }
 
+    pub(crate) fn has_alloc_error_handler(&self) -> bool {
+        self.has_alloc_error_handler
+    }
+
     pub fn report_unused_deps(&self, tcx: TyCtxt<'_>) {
         let json_unused_externs = tcx.sess.opts.json_unused_externs;
 
@@ -247,7 +262,9 @@ impl<'a> CrateLoader<'a> {
                 metas: IndexVec::from_elem_n(None, 1),
                 injected_panic_runtime: None,
                 allocator_kind: None,
+                alloc_error_handler_kind: None,
                 has_global_allocator: false,
+                has_alloc_error_handler: false,
                 stable_crate_ids,
                 unused_externs: Vec::new(),
             },
@@ -792,6 +809,13 @@ impl<'a> CrateLoader<'a> {
             }
             spans => !spans.is_empty(),
         };
+        self.cstore.has_alloc_error_handler = match &*alloc_error_handler_spans(&self.sess, krate) {
+            [span1, span2, ..] => {
+                self.sess.emit_err(NoMultipleAllocErrorHandler { span2: *span2, span1: *span1 });
+                true
+            }
+            spans => !spans.is_empty(),
+        };
 
         // Check to see if we actually need an allocator. This desire comes
         // about through the `#![needs_allocator]` attribute and is typically
@@ -832,22 +856,48 @@ impl<'a> CrateLoader<'a> {
                 }
             }
         }
+        let mut alloc_error_handler =
+            self.cstore.has_alloc_error_handler.then(|| Symbol::intern("this crate"));
+        for (_, data) in self.cstore.iter_crate_data() {
+            if data.has_alloc_error_handler() {
+                match alloc_error_handler {
+                    Some(other_crate) => {
+                        self.sess.emit_err(ConflictingAllocErrorHandler {
+                            crate_name: data.name(),
+                            other_crate_name: other_crate,
+                        });
+                    }
+                    None => alloc_error_handler = Some(data.name()),
+                }
+            }
+        }
 
         if global_allocator.is_some() {
             self.cstore.allocator_kind = Some(AllocatorKind::Global);
-            return;
+        } else {
+            // Ok we haven't found a global allocator but we still need an
+            // allocator. At this point our allocator request is typically fulfilled
+            // by the standard library, denoted by the `#![default_lib_allocator]`
+            // attribute.
+            if !self.sess.contains_name(&krate.attrs, sym::default_lib_allocator)
+                && !self.cstore.iter_crate_data().any(|(_, data)| data.has_default_lib_allocator())
+            {
+                self.sess.emit_err(GlobalAllocRequired);
+            }
+            self.cstore.allocator_kind = Some(AllocatorKind::Default);
         }
 
-        // Ok we haven't found a global allocator but we still need an
-        // allocator. At this point our allocator request is typically fulfilled
-        // by the standard library, denoted by the `#![default_lib_allocator]`
-        // attribute.
-        if !self.sess.contains_name(&krate.attrs, sym::default_lib_allocator)
-            && !self.cstore.iter_crate_data().any(|(_, data)| data.has_default_lib_allocator())
-        {
-            self.sess.emit_err(GlobalAllocRequired);
+        if alloc_error_handler.is_some() {
+            self.cstore.alloc_error_handler_kind = Some(AllocatorKind::Global);
+        } else {
+            // The alloc crate provides a default allocation error handler if
+            // one isn't specified.
+            if !self.sess.features_untracked().default_alloc_error_handler {
+                self.sess.emit_err(AllocFuncRequired);
+                self.sess.emit_note(MissingAllocErrorHandler);
+            }
+            self.cstore.alloc_error_handler_kind = Some(AllocatorKind::Default);
         }
-        self.cstore.allocator_kind = Some(AllocatorKind::Default);
     }
 
     fn inject_dependency_if(
@@ -1023,3 +1073,26 @@ fn global_allocator_spans(sess: &Session, krate: &ast::Crate) -> Vec<Span> {
     visit::walk_crate(&mut f, krate);
     f.spans
 }
+
+fn alloc_error_handler_spans(sess: &Session, krate: &ast::Crate) -> Vec<Span> {
+    struct Finder<'a> {
+        sess: &'a Session,
+        name: Symbol,
+        spans: Vec<Span>,
+    }
+    impl<'ast, 'a> visit::Visitor<'ast> for Finder<'a> {
+        fn visit_item(&mut self, item: &'ast ast::Item) {
+            if item.ident.name == self.name
+                && self.sess.contains_name(&item.attrs, sym::rustc_std_internal_symbol)
+            {
+                self.spans.push(item.span);
+            }
+            visit::walk_item(self, item)
+        }
+    }
+
+    let name = Symbol::intern(&AllocatorKind::Global.fn_name(sym::oom));
+    let mut f = Finder { sess, name, spans: Vec::new() };
+    visit::walk_crate(&mut f, krate);
+    f.spans
+}

compiler/rustc_metadata/src/errors.rs

@@ -343,6 +343,16 @@ pub struct NoMultipleGlobalAlloc {
     pub span1: Span,
 }
 
+#[derive(Diagnostic)]
+#[diag(metadata_no_multiple_alloc_error_handler)]
+pub struct NoMultipleAllocErrorHandler {
+    #[primary_span]
+    #[label]
+    pub span2: Span,
+    #[label(metadata_prev_alloc_error_handler)]
+    pub span1: Span,
+}
+
 #[derive(Diagnostic)]
 #[diag(metadata_conflicting_global_alloc)]
 pub struct ConflictingGlobalAlloc {
@@ -350,10 +360,25 @@ pub struct ConflictingGlobalAlloc {
     pub other_crate_name: Symbol,
 }
 
+#[derive(Diagnostic)]
+#[diag(metadata_conflicting_alloc_error_handler)]
+pub struct ConflictingAllocErrorHandler {
+    pub crate_name: Symbol,
+    pub other_crate_name: Symbol,
+}
+
 #[derive(Diagnostic)]
 #[diag(metadata_global_alloc_required)]
 pub struct GlobalAllocRequired;
 
+#[derive(Diagnostic)]
+#[diag(metadata_alloc_func_required)]
+pub struct AllocFuncRequired;
+
+#[derive(Diagnostic)]
+#[diag(metadata_missing_alloc_error_handler)]
+pub struct MissingAllocErrorHandler;
+
 #[derive(Diagnostic)]
 #[diag(metadata_no_transitive_needs_dep)]
 pub struct NoTransitiveNeedsDep<'a> {

compiler/rustc_metadata/src/rmeta/decoder.rs

@@ -1764,6 +1764,10 @@ impl CrateMetadata {
         self.root.has_global_allocator
     }
 
+    pub(crate) fn has_alloc_error_handler(&self) -> bool {
+        self.root.has_alloc_error_handler
+    }
+
     pub(crate) fn has_default_lib_allocator(&self) -> bool {
         self.root.has_default_lib_allocator
     }

compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs

@@ -255,6 +255,7 @@ provide! { tcx, def_id, other, cdata,
     is_panic_runtime => { cdata.root.panic_runtime }
     is_compiler_builtins => { cdata.root.compiler_builtins }
     has_global_allocator => { cdata.root.has_global_allocator }
+    has_alloc_error_handler => { cdata.root.has_alloc_error_handler }
     has_panic_handler => { cdata.root.has_panic_handler }
     is_profiler_runtime => { cdata.root.profiler_runtime }
     required_panic_strategy => { cdata.root.required_panic_strategy }
@@ -339,6 +340,7 @@ pub(in crate::rmeta) fn provide(providers: &mut Providers) {
     // resolve! Does this work? Unsure! That's what the issue is about
     *providers = Providers {
         allocator_kind: |tcx, ()| CStore::from_tcx(tcx).allocator_kind(),
+        alloc_error_handler_kind: |tcx, ()| CStore::from_tcx(tcx).alloc_error_handler_kind(),
         is_private_dep: |_tcx, cnum| {
             assert_eq!(cnum, LOCAL_CRATE);
             false
@@ -464,6 +466,10 @@ pub(in crate::rmeta) fn provide(providers: &mut Providers) {
             assert_eq!(cnum, LOCAL_CRATE);
             CStore::from_tcx(tcx).has_global_allocator()
         },
+        has_alloc_error_handler: |tcx, cnum| {
+            assert_eq!(cnum, LOCAL_CRATE);
+            CStore::from_tcx(tcx).has_alloc_error_handler()
+        },
         postorder_cnums: |tcx, ()| {
             tcx.arena
                 .alloc_slice(&CStore::from_tcx(tcx).crate_dependencies_in_postorder(LOCAL_CRATE))

compiler/rustc_metadata/src/rmeta/encoder.rs

@@ -670,6 +670,7 @@ impl<'a, 'tcx> EncodeContext<'a, 'tcx> {
                 panic_in_drop_strategy: tcx.sess.opts.unstable_opts.panic_in_drop,
                 edition: tcx.sess.edition(),
                 has_global_allocator: tcx.has_global_allocator(LOCAL_CRATE),
+                has_alloc_error_handler: tcx.has_alloc_error_handler(LOCAL_CRATE),
                 has_panic_handler: tcx.has_panic_handler(LOCAL_CRATE),
                 has_default_lib_allocator: tcx
                     .sess

compiler/rustc_metadata/src/rmeta/mod.rs

@@ -223,6 +223,7 @@ pub(crate) struct CrateRoot {
     panic_in_drop_strategy: PanicStrategy,
     edition: Edition,
     has_global_allocator: bool,
+    has_alloc_error_handler: bool,
     has_panic_handler: bool,
     has_default_lib_allocator: bool,