Revert "Remove #[alloc_error_handler] from the compiler and library"

This reverts commit abc0660118cc95f47445fd33502a11dd448f5968.

7 files changed, 104 insertions, 0 deletions

diff --git a/compiler/rustc_metadata/messages.ftl b/compiler/rustc_metadata/messages.ftl
index f0158aeae85..79b8b417257 100644
--- a/compiler/rustc_metadata/messages.ftl
+++ b/compiler/rustc_metadata/messages.ftl
@@ -155,9 +155,19 @@ metadata_no_multiple_global_alloc =
 metadata_prev_global_alloc =
     previous global allocator defined here
 
+metadata_no_multiple_alloc_error_handler =
+    cannot define multiple allocation error handlers
+    .label = cannot define a new allocation error handler
+
+metadata_prev_alloc_error_handler =
+    previous allocation error handler defined here
+
 metadata_conflicting_global_alloc =
     the `#[global_allocator]` in {$other_crate_name} conflicts with global allocator in: {$crate_name}
 
+metadata_conflicting_alloc_error_handler =
+    the `#[alloc_error_handler]` in {$other_crate_name} conflicts with allocation error handler in: {$crate_name}
+
 metadata_global_alloc_required =
     no global memory allocator found but one is required; link to std or add `#[global_allocator]` to a static item that implements the GlobalAlloc trait
 
diff --git a/compiler/rustc_metadata/src/creader.rs b/compiler/rustc_metadata/src/creader.rs
index 8bacd136647..179453238f2 100644
--- a/compiler/rustc_metadata/src/creader.rs
+++ b/compiler/rustc_metadata/src/creader.rs
@@ -38,8 +38,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,
 
     /// The interned [StableCrateId]s.
     pub(crate) stable_crate_ids: StableCrateIdMap,
@@ -216,10 +221,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;
 
@@ -255,7 +268,9 @@ impl CStore {
             metas: IndexVec::from_iter(iter::once(None)),
             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(),
         }
@@ -761,6 +776,14 @@ impl<'a, 'tcx> CrateLoader<'a, 'tcx> {
             }
             spans => !spans.is_empty(),
         };
+        self.cstore.has_alloc_error_handler = match &*alloc_error_handler_spans(krate) {
+            [span1, span2, ..] => {
+                self.sess
+                    .emit_err(errors::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
@@ -801,6 +824,21 @@ impl<'a, 'tcx> CrateLoader<'a, 'tcx> {
                 }
             }
         }
+        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(errors::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);
@@ -816,6 +854,14 @@ impl<'a, 'tcx> CrateLoader<'a, 'tcx> {
             }
             self.cstore.allocator_kind = Some(AllocatorKind::Default);
         }
+
+        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.
+            self.cstore.alloc_error_handler_kind = Some(AllocatorKind::Default);
+        }
     }
 
     fn inject_dependency_if(
@@ -991,6 +1037,28 @@ fn global_allocator_spans(krate: &ast::Crate) -> Vec<Span> {
     f.spans
 }
 
+fn alloc_error_handler_spans(krate: &ast::Crate) -> Vec<Span> {
+    struct Finder {
+        name: Symbol,
+        spans: Vec<Span>,
+    }
+    impl<'ast> visit::Visitor<'ast> for Finder {
+        fn visit_item(&mut self, item: &'ast ast::Item) {
+            if item.ident.name == self.name
+                && attr::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 { name, spans: Vec::new() };
+    visit::walk_crate(&mut f, krate);
+    f.spans
+}
+
 // On Windows the compiler would sometimes intermittently fail to open the
 // proc-macro DLL with `Error::LoadLibraryExW`. It is suspected that something in the
 // system still holds a lock on the file, so we retry a few times before calling it
diff --git a/compiler/rustc_metadata/src/errors.rs b/compiler/rustc_metadata/src/errors.rs
index 7ecb551a3e5..51b41b5f6a2 100644
--- a/compiler/rustc_metadata/src/errors.rs
+++ b/compiler/rustc_metadata/src/errors.rs
@@ -353,6 +353,16 @@ pub struct NoMultipleGlobalAlloc {
 }
 
 #[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 {
     pub crate_name: Symbol,
@@ -360,6 +370,13 @@ pub struct ConflictingGlobalAlloc {
 }
 
 #[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;
 
diff --git a/compiler/rustc_metadata/src/rmeta/decoder.rs b/compiler/rustc_metadata/src/rmeta/decoder.rs
index 820a502d04c..64245a0460d 100644
--- a/compiler/rustc_metadata/src/rmeta/decoder.rs
+++ b/compiler/rustc_metadata/src/rmeta/decoder.rs
@@ -1683,6 +1683,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
     }
diff --git a/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs b/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
index 4aa3768fc3b..141980912b1 100644
--- a/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
+++ b/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
@@ -290,6 +290,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 }
@@ -378,6 +379,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, LocalCrate| false,
         native_library: |tcx, id| {
             tcx.native_libraries(id.krate)
@@ -494,6 +496,7 @@ pub(in crate::rmeta) fn provide(providers: &mut Providers) {
 
         dependency_formats: |tcx, ()| Lrc::new(crate::dependency_format::calculate(tcx)),
         has_global_allocator: |tcx, LocalCrate| CStore::from_tcx(tcx).has_global_allocator(),
+        has_alloc_error_handler: |tcx, LocalCrate| 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))
diff --git a/compiler/rustc_metadata/src/rmeta/encoder.rs b/compiler/rustc_metadata/src/rmeta/encoder.rs
index fd8e49efea0..17a9daee7d0 100644
--- a/compiler/rustc_metadata/src/rmeta/encoder.rs
+++ b/compiler/rustc_metadata/src/rmeta/encoder.rs
@@ -676,6 +676,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: attr::contains_name(&attrs, sym::default_lib_allocator),
                 proc_macro_data,
diff --git a/compiler/rustc_metadata/src/rmeta/mod.rs b/compiler/rustc_metadata/src/rmeta/mod.rs
index 1bb0eabc64f..dd02463e16a 100644
--- a/compiler/rustc_metadata/src/rmeta/mod.rs
+++ b/compiler/rustc_metadata/src/rmeta/mod.rs
@@ -225,6 +225,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,