Implement #[alloc_error_handler]

This to-be-stable attribute is equivalent to `#[lang = "oom"]`.
It is required when using the alloc crate without the std crate.
It is called by `handle_alloc_error`, which is in turned called
by "infallible" allocations APIs such as `Vec::push`.

src/libcore/alloc.rs

@@ -48,6 +48,7 @@ fn size_align<T>() -> (usize, usize) {
 /// use specific allocators with looser requirements.)
 #[stable(feature = "alloc_layout", since = "1.28.0")]
 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
+#[cfg_attr(not(stage0), lang = "alloc_layout")]
 pub struct Layout {
     // size of the requested block of memory, measured in bytes.
     size_: usize,

src/librustc/middle/dead.rs

@@ -288,7 +288,17 @@ impl<'a, 'tcx> Visitor<'tcx> for MarkSymbolVisitor<'a, 'tcx> {
 fn has_allow_dead_code_or_lang_attr(tcx: TyCtxt,
                                     id: ast::NodeId,
                                     attrs: &[ast::Attribute]) -> bool {
-    if attr::contains_name(attrs, "lang") || attr::contains_name(attrs, "panic_implementation") {
+    if attr::contains_name(attrs, "lang") {
+        return true;
+    }
+
+    // (To be) stable attribute for #[lang = "panic_impl"]
+    if attr::contains_name(attrs, "panic_implementation") {
+        return true;
+    }
+
+    // (To be) stable attribute for #[lang = "oom"]
+    if attr::contains_name(attrs, "alloc_error_handler") {
         return true;
     }
 

src/librustc/middle/lang_items.rs

@@ -187,6 +187,8 @@ pub fn extract(attrs: &[ast::Attribute]) -> Option<(Symbol, Span)> {
             }
         } else if attribute.check_name("panic_implementation") {
             return Some((Symbol::intern("panic_impl"), attribute.span))
+        } else if attribute.check_name("alloc_error_handler") {
+            return Some((Symbol::intern("oom"), attribute.span))
         }
     }
 
@@ -308,6 +310,7 @@ language_item_table! {
     BoxFreeFnLangItem,               "box_free",                box_free_fn;
     DropInPlaceFnLangItem,           "drop_in_place",           drop_in_place_fn;
     OomLangItem,                     "oom",                     oom;
+    AllocLayoutLangItem,             "alloc_layout",            alloc_layout;
 
     StartFnLangItem,                 "start",                   start_fn;
 

src/librustc/middle/weak_lang_items.rs

@@ -115,6 +115,9 @@ fn verify<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
             if lang_items::$item == lang_items::PanicImplLangItem {
                 tcx.sess.err(&format!("`#[panic_implementation]` function required, \
                                         but not found"));
+            } else if lang_items::$item == lang_items::OomLangItem {
+                tcx.sess.err(&format!("`#[alloc_error_handler]` function required, \
+                                        but not found"));
             } else {
                 tcx.sess.err(&format!("language item required, but not found: `{}`",
                                         stringify!($name)));

src/librustc_typeck/check/mod.rs

@@ -1182,7 +1182,54 @@ fn check_fn<'a, 'gcx, 'tcx>(inherited: &'a Inherited<'a, 'gcx, 'tcx>,
                 fcx.tcx.sess.err("language item required, but not found: `panic_info`");
             }
         }
+    }
 
+    // Check that a function marked as `#[alloc_error_handler]` has signature `fn(Layout) -> !`
+    if let Some(alloc_error_handler_did) = fcx.tcx.lang_items().oom() {
+        if alloc_error_handler_did == fcx.tcx.hir.local_def_id(fn_id) {
+            if let Some(alloc_layout_did) = fcx.tcx.lang_items().alloc_layout() {
+                if declared_ret_ty.sty != ty::TyNever {
+                    fcx.tcx.sess.span_err(
+                        decl.output.span(),
+                        "return type should be `!`",
+                    );
+                }
+
+                let inputs = fn_sig.inputs();
+                let span = fcx.tcx.hir.span(fn_id);
+                if inputs.len() == 1 {
+                    let arg_is_alloc_layout = match inputs[0].sty {
+                        ty::TyAdt(ref adt, _) => {
+                            adt.did == alloc_layout_did
+                        },
+                        _ => false,
+                    };
+
+                    if !arg_is_alloc_layout {
+                        fcx.tcx.sess.span_err(
+                            decl.inputs[0].span,
+                            "argument should be `Layout`",
+                        );
+                    }
+
+                    if let Node::NodeItem(item) = fcx.tcx.hir.get(fn_id) {
+                        if let Item_::ItemFn(_, _, ref generics, _) = item.node {
+                            if !generics.params.is_empty() {
+                                fcx.tcx.sess.span_err(
+                                    span,
+                                    "`#[alloc_error_handler]` function should have no type \
+                                     parameters",
+                                );
+                            }
+                        }
+                    }
+                } else {
+                    fcx.tcx.sess.span_err(span, "function should have one argument");
+                }
+            } else {
+                fcx.tcx.sess.err("language item required, but not found: `alloc_layout`");
+            }
+        }
     }
 
     (fcx, gen_ty)

src/libstd/alloc.rs

@@ -125,7 +125,8 @@ fn default_alloc_error_hook(layout: Layout) {
 
 #[cfg(not(test))]
 #[doc(hidden)]
-#[lang = "oom"]
+#[cfg_attr(stage0, lang = "oom")]
+#[cfg_attr(not(stage0), alloc_error_handler)]
 #[unstable(feature = "alloc_internals", issue = "0")]
 pub fn rust_oom(layout: Layout) -> ! {
     let hook = HOOK.load(Ordering::SeqCst);

src/libstd/lib.rs

@@ -233,8 +233,9 @@
 // std is implemented with unstable features, many of which are internal
 // compiler details that will never be stable
 #![feature(alloc)]
-#![feature(allocator_api)]
+#![feature(alloc_error_handler)]
 #![feature(alloc_system)]
+#![feature(allocator_api)]
 #![feature(allocator_internals)]
 #![feature(allow_internal_unsafe)]
 #![feature(allow_internal_unstable)]

src/libsyntax/feature_gate.rs

@@ -481,6 +481,9 @@ declare_features! (
 
     // Allows async and await syntax
     (active, async_await, "1.28.0", Some(50547), None),
+
+    // #[alloc_error_handler]
+    (active, alloc_error_handler, "1.29.0", Some(51540), None),
 );
 
 declare_features! (
@@ -1083,6 +1086,11 @@ pub const BUILTIN_ATTRIBUTES: &'static [(&'static str, AttributeType, AttributeG
                            "#[panic_implementation] is an unstable feature",
                            cfg_fn!(panic_implementation))),
 
+    ("alloc_error_handler", Normal, Gated(Stability::Unstable,
+                           "alloc_error_handler",
+                           "#[alloc_error_handler] is an unstable feature",
+                           cfg_fn!(alloc_error_handler))),
+
     // Crate level attributes
     ("crate_name", CrateLevel, Ungated),
     ("crate_type", CrateLevel, Ungated),

src/test/compile-fail/alloc-error-handler-bad-signature-1.rs

@@ -0,0 +1,28 @@
+// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// compile-flags:-C panic=abort
+
+#![feature(alloc_error_handler, panic_implementation)]
+#![no_std]
+#![no_main]
+
+use core::alloc::Layout;
+
+#[alloc_error_handler]
+fn oom(
+    info: &Layout, //~ ERROR argument should be `Layout`
+) -> () //~ ERROR return type should be `!`
+{
+    loop {}
+}
+
+#[panic_implementation]
+fn panic(_: &core::panic::PanicInfo) -> ! { loop {} }

src/test/compile-fail/alloc-error-handler-bad-signature-2.rs

@@ -0,0 +1,27 @@
+// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// compile-flags:-C panic=abort
+
+#![feature(alloc_error_handler, panic_implementation)]
+#![no_std]
+#![no_main]
+
+struct Layout;
+
+#[alloc_error_handler]
+fn oom(
+    info: Layout, //~ ERROR argument should be `Layout`
+) { //~ ERROR return type should be `!`
+    loop {}
+}
+
+#[panic_implementation]
+fn panic(_: &core::panic::PanicInfo) -> ! { loop {} }

src/test/compile-fail/alloc-error-handler-bad-signature-3.rs

@@ -0,0 +1,25 @@
+// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// compile-flags:-C panic=abort
+
+#![feature(alloc_error_handler, panic_implementation)]
+#![no_std]
+#![no_main]
+
+struct Layout;
+
+#[alloc_error_handler]
+fn oom() -> ! { //~ ERROR function should have one argument
+    loop {}
+}
+
+#[panic_implementation]
+fn panic(_: &core::panic::PanicInfo) -> ! { loop {} }

src/test/compile-fail/feature-gate-alloc-error-handler.rs

@@ -0,0 +1,21 @@
+// Copyright 2018 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// compile-flags:-C panic=abort
+
+#![no_std]
+#![no_main]
+
+use core::alloc::Layout;
+
+#[alloc_error_handler] //~ ERROR #[alloc_error_handler] is an unstable feature (see issue #51540)
+fn oom(info: Layout) -> ! {
+    loop {}
+}

src/test/run-make-fulldeps/issue-51671/app.rs

@@ -14,6 +14,7 @@
 #![no_main]
 #![no_std]
 
+use core::alloc::Layout;
 use core::panic::PanicInfo;
 
 #[panic_implementation]
@@ -25,4 +26,6 @@ fn panic(_: &PanicInfo) -> ! {
 fn eh() {}
 
 #[lang = "oom"]
-fn oom() {}
+fn oom(_: Layout) -> ! {
+    loop {}
+}

src/test/ui/missing-alloc_error_handler.rs

@@ -0,0 +1,33 @@
+// Copyright 2015 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+// compile-flags: -C panic=abort
+// no-prefer-dynamic
+
+#![no_std]
+#![crate_type = "staticlib"]
+#![feature(panic_implementation, alloc_error_handler, alloc)]
+
+#[panic_implementation]
+fn panic(_: &core::panic::PanicInfo) -> ! {
+    loop {}
+}
+
+extern crate alloc;
+
+#[global_allocator]
+static A: MyAlloc = MyAlloc;
+
+struct MyAlloc;
+
+unsafe impl core::alloc::GlobalAlloc for MyAlloc {
+    unsafe fn alloc(&self, _: core::alloc::Layout) -> *mut u8 { 0 as _ }
+    unsafe fn dealloc(&self, _: *mut u8, _: core::alloc::Layout) {}
+}

src/test/ui/missing-alloc_error_handler.stderr

@@ -0,0 +1,4 @@
+error: `#[alloc_error_handler]` function required, but not found
+
+error: aborting due to previous error
+

src/test/ui/missing-allocator.rs

@@ -13,14 +13,16 @@
 
 #![no_std]
 #![crate_type = "staticlib"]
-#![feature(panic_implementation, lang_items, alloc)]
+#![feature(panic_implementation, alloc_error_handler, alloc)]
 
 #[panic_implementation]
 fn panic(_: &core::panic::PanicInfo) -> ! {
     loop {}
 }
 
-#[lang = "oom"]
+#[alloc_error_handler]
-fn oom() {}
+fn oom(_: core::alloc::Layout) -> ! {
+    loop {}
+}
 
 extern crate alloc;