Move OOM handling to liballoc and remove the oom lang item

src/liballoc/alloc.rs

@@ -15,6 +15,8 @@
 use core::intrinsics::{min_align_of_val, size_of_val};
 use core::ptr::{NonNull, Unique};
 use core::usize;
+#[cfg(target_has_atomic = "ptr")] use core::{mem, ptr};
+#[cfg(target_has_atomic = "ptr")] use core::sync::atomic::{AtomicPtr, Ordering};
 
 #[stable(feature = "alloc_module", since = "1.28.0")]
 #[doc(inline)]
@@ -191,12 +193,89 @@ pub(crate) unsafe fn box_free<T: ?Sized>(ptr: Unique<T>) {
 #[stable(feature = "global_alloc", since = "1.28.0")]
 #[rustc_allocator_nounwind]
 pub fn handle_alloc_error(layout: Layout) -> ! {
-    #[allow(improper_ctypes)]
-    extern "Rust" {
-        #[lang = "oom"]
-        fn oom_impl(layout: Layout) -> !;
+    #[cfg(target_has_atomic = "ptr")]
+    {
+        let hook = HOOK.load(Ordering::SeqCst);
+        let hook: fn(Layout) = if hook.is_null() {
+            default_alloc_error_hook
+        } else {
+            unsafe { mem::transmute(hook) }
+        };
+        hook(layout);
+    }
+    #[cfg(not(target_has_atomic = "ptr"))]
+    {
+        default_alloc_error_hook(layout)
+    }
+
+    unsafe {
+        #[cfg(not(stage0))]
+        {
+            #[allow(improper_ctypes)]
+            extern "Rust" {
+                // This function is generated by the compiler
+                // and calls __rust_abort_internal in src/libstd/alloc.rs
+                // if the `std` crate is linked.
+                fn __rust_maybe_abort_internal();
+            }
+            __rust_maybe_abort_internal();
+        }
+        // The above does nothing and returns when `std` is not linked.
+        // In that case we abort "naively" here:
+        ::core::intrinsics::abort()
+    }
+}
+
+fn default_alloc_error_hook(layout: Layout) {
+    #[cfg(stage0)]
+    let _ = layout;
+    #[cfg(not(stage0))]
+    unsafe {
+        #[allow(improper_ctypes)]
+        extern "Rust" {
+            // This function is generated by the compiler
+            // and calls __rust_default_alloc_error_hook in src/libstd/alloc.rs
+            // if the `std` crate is linked.
+            fn __rust_maybe_default_alloc_error_hook(size: usize, align: usize);
+        }
+        __rust_maybe_default_alloc_error_hook(layout.size(), layout.align())
+    }
+}
+
+#[cfg(target_has_atomic = "ptr")]
+static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
+
+/// Registers a custom allocation error hook, replacing any that was previously registered.
+///
+/// The allocation error hook is invoked when an infallible memory allocation fails, before
+/// the runtime aborts. The default hook prints a message to standard error,
+/// but this behavior can be customized with the [`set_alloc_error_hook`] and
+/// [`take_alloc_error_hook`] functions.
+///
+/// The hook is provided with a `Layout` struct which contains information
+/// about the allocation that failed.
+///
+/// The allocation error hook is a global resource.
+#[unstable(feature = "alloc_error_hook", issue = "51245")]
+#[cfg(target_has_atomic = "ptr")]
+pub fn set_alloc_error_hook(hook: fn(Layout)) {
+    HOOK.store(hook as *mut (), Ordering::SeqCst);
+}
+
+/// Unregisters the current allocation error hook, returning it.
+///
+/// *See also the function [`set_alloc_error_hook`].*
+///
+/// If no custom hook is registered, the default hook will be returned.
+#[unstable(feature = "alloc_error_hook", issue = "51245")]
+#[cfg(target_has_atomic = "ptr")]
+pub fn take_alloc_error_hook() -> fn(Layout) {
+    let hook = HOOK.swap(ptr::null_mut(), Ordering::SeqCst);
+    if hook.is_null() {
+        default_alloc_error_hook
+    } else {
+        unsafe { mem::transmute(hook) }
     }
-    unsafe { oom_impl(layout) }
 }
 
 #[cfg(test)]

src/librustc/middle/lang_items.rs

@@ -307,7 +307,6 @@ language_item_table! {
     ExchangeMallocFnLangItem,        "exchange_malloc",         exchange_malloc_fn;
     BoxFreeFnLangItem,               "box_free",                box_free_fn;
     DropInPlaceFnLangItem,           "drop_in_place",           drop_in_place_fn;
-    OomLangItem,                     "oom",                     oom;
 
     StartFnLangItem,                 "start",                   start_fn;
 

src/librustc/middle/weak_lang_items.rs

@@ -151,5 +151,4 @@ weak_lang_items! {
     panic_impl,         PanicImplLangItem,          rust_begin_unwind;
     eh_personality,     EhPersonalityLangItem,      rust_eh_personality;
     eh_unwind_resume,   EhUnwindResumeLangItem,     rust_eh_unwind_resume;
-    oom,                OomLangItem,                rust_oom;
 }

src/librustc/session/mod.rs

@@ -118,6 +118,7 @@ pub struct Session {
     /// injected.
     pub injected_allocator: Once<Option<CrateNum>>,
     pub allocator_kind: Once<Option<AllocatorKind>>,
+    pub injected_default_alloc_error_hook: Once<InjectedDefaultOomHook>,
     pub injected_panic_runtime: Once<Option<CrateNum>>,
 
     /// Map from imported macro spans (which consist of
@@ -175,6 +176,12 @@ pub struct PerfStats {
     pub normalize_projection_ty: AtomicUsize,
 }
 
+pub enum InjectedDefaultOomHook {
+    None,
+    Noop,
+    Platform,
+}
+
 /// Enum to support dispatch of one-time diagnostics (in Session.diag_once)
 enum DiagnosticBuilderMethod {
     Note,
@@ -1119,6 +1126,7 @@ pub fn build_session_(
         next_node_id: OneThread::new(Cell::new(NodeId::new(1))),
         injected_allocator: Once::new(),
         allocator_kind: Once::new(),
+        injected_default_alloc_error_hook: Once::new(),
         injected_panic_runtime: Once::new(),
         imported_macro_spans: OneThread::new(RefCell::new(HashMap::new())),
         incr_comp_session: OneThread::new(RefCell::new(IncrCompSession::NotInitialized)),

src/librustc_allocator/lib.rs

@@ -42,6 +42,19 @@ pub static ALLOCATOR_METHODS: &[AllocatorMethod] = &[
     },
 ];
 
+pub static OOM_HANDLING_METHODS: &[AllocatorMethod] = &[
+    AllocatorMethod {
+        name: "default_alloc_error_hook",
+        inputs: &[AllocatorTy::Layout],
+        output: AllocatorTy::Unit,
+    },
+    AllocatorMethod {
+        name: "abort_internal",
+        inputs: &[],
+        output: AllocatorTy::Unit,
+    },
+];
+
 pub struct AllocatorMethod {
     pub name: &'static str,
     pub inputs: &'static [AllocatorTy],

src/librustc_codegen_llvm/allocator.rs

@@ -14,8 +14,9 @@ use std::ptr;
 use attributes;
 use libc::c_uint;
 use rustc::middle::allocator::AllocatorKind;
+use rustc::session::InjectedDefaultOomHook;
 use rustc::ty::TyCtxt;
-use rustc_allocator::{ALLOCATOR_METHODS, AllocatorTy};
+use rustc_allocator::{ALLOCATOR_METHODS, OOM_HANDLING_METHODS, AllocatorTy};
 
 use ModuleLlvm;
 use llvm::{self, False, True};
@@ -33,9 +34,10 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt, mods: &ModuleLlvm, kind: AllocatorKind
     let i8p = llvm::LLVMPointerType(i8, 0);
     let void = llvm::LLVMVoidTypeInContext(llcx);
 
-    for method in ALLOCATOR_METHODS {
+    let build = |name: String, inputs: &[AllocatorTy], output: &AllocatorTy,
+                 callee: Option<String>| {
         let mut args = Vec::new();
-        for ty in method.inputs.iter() {
+        for ty in inputs.iter() {
             match *ty {
                 AllocatorTy::Layout => {
                     args.push(usize); // size
@@ -48,7 +50,7 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt, mods: &ModuleLlvm, kind: AllocatorKind
                 AllocatorTy::Unit => panic!("invalid allocator arg"),
             }
         }
-        let output = match method.output {
+        let output = match *output {
             AllocatorTy::ResultPtr => Some(i8p),
             AllocatorTy::Unit => None,
 
@@ -60,29 +62,40 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt, mods: &ModuleLlvm, kind: AllocatorKind
                                         args.as_ptr(),
                                         args.len() as c_uint,
                                         False);
-        let name = CString::new(format!("__rust_{}", method.name)).unwrap();
+        let name = CString::new(name).unwrap();
         let llfn = llvm::LLVMRustGetOrInsertFunction(llmod,
                                                      name.as_ptr(),
                                                      ty);
 
         if tcx.sess.target.target.options.default_hidden_visibility {
             llvm::LLVMRustSetVisibility(llfn, llvm::Visibility::Hidden);
         }
-       if tcx.sess.target.target.options.requires_uwtable {
-           attributes::emit_uwtable(llfn, true);
-       }
-
-        let callee = CString::new(kind.fn_name(method.name)).unwrap();
-        let callee = llvm::LLVMRustGetOrInsertFunction(llmod,
-                                                       callee.as_ptr(),
-                                                       ty);
+        if tcx.sess.target.target.options.requires_uwtable {
+            attributes::emit_uwtable(llfn, true);
+        }
 
         let llbb = llvm::LLVMAppendBasicBlockInContext(llcx,
                                                        llfn,
                                                        "entry\0".as_ptr() as *const _);
 
         let llbuilder = llvm::LLVMCreateBuilderInContext(llcx);
         llvm::LLVMPositionBuilderAtEnd(llbuilder, llbb);
+
+        let callee = if let Some(callee) = callee {
+            callee
+        } else {
+            // Generate a no-op function
+            llvm::LLVMBuildRetVoid(llbuilder);
+            llvm::LLVMDisposeBuilder(llbuilder);
+            return
+        };
+
+        // Forward the call to another function
+        let callee = CString::new(callee).unwrap();
+        let callee = llvm::LLVMRustGetOrInsertFunction(llmod,
+                                                       callee.as_ptr(),
+                                                       ty);
+
         let args = args.iter().enumerate().map(|(i, _)| {
             llvm::LLVMGetParam(llfn, i as c_uint)
         }).collect::<Vec<_>>();
@@ -98,6 +111,28 @@ pub(crate) unsafe fn codegen(tcx: TyCtxt, mods: &ModuleLlvm, kind: AllocatorKind
         } else {
             llvm::LLVMBuildRetVoid(llbuilder);
         }
+
         llvm::LLVMDisposeBuilder(llbuilder);
+    };
+
+    for method in ALLOCATOR_METHODS {
+        let name = format!("__rust_{}", method.name);
+        build(name, method.inputs, &method.output, Some(kind.fn_name(method.name)))
+    }
+
+    let has_plaftom_functions = match tcx.sess.injected_default_alloc_error_hook.get() {
+        InjectedDefaultOomHook::None => return,
+        InjectedDefaultOomHook::Noop => false,
+        InjectedDefaultOomHook::Platform => true,
+    };
+
+    for method in OOM_HANDLING_METHODS {
+        let callee = if has_plaftom_functions {
+            Some(format!("__rust_{}", method.name))
+        } else {
+            None
+        };
+        let name = format!("__rust_maybe_{}", method.name);
+        build(name, method.inputs, &method.output, callee)
     }
 }

src/librustc_codegen_llvm/back/symbol_export.rs

@@ -17,12 +17,12 @@ use rustc::hir::CodegenFnAttrFlags;
 use rustc::hir::def_id::{CrateNum, DefId, LOCAL_CRATE, CRATE_DEF_INDEX};
 use rustc::ich::Fingerprint;
 use rustc::middle::exported_symbols::{SymbolExportLevel, ExportedSymbol, metadata_symbol_name};
-use rustc::session::config;
+use rustc::session::{config, InjectedDefaultOomHook};
 use rustc::ty::{TyCtxt, SymbolName};
 use rustc::ty::query::Providers;
 use rustc::ty::subst::Substs;
 use rustc::util::nodemap::{FxHashMap, DefIdMap};
-use rustc_allocator::ALLOCATOR_METHODS;
+use rustc_allocator::{ALLOCATOR_METHODS, OOM_HANDLING_METHODS};
 use rustc_data_structures::indexed_vec::IndexVec;
 use std::collections::hash_map::Entry::*;
 
@@ -222,6 +222,19 @@ fn exported_symbols_provider_local<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
         }
     }
 
+    match tcx.sess.injected_default_alloc_error_hook.get() {
+        InjectedDefaultOomHook::None => {}
+        InjectedDefaultOomHook::Noop |
+        InjectedDefaultOomHook::Platform => {
+            for method in OOM_HANDLING_METHODS {
+                let symbol_name = format!("__rust_{}", method.name);
+                let exported_symbol = ExportedSymbol::NoDefId(SymbolName::new(&symbol_name));
+
+                symbols.push((exported_symbol, SymbolExportLevel::Rust));
+            }
+        }
+    }
+
     if tcx.sess.opts.debugging_opts.pgo_gen.is_some() {
         // These are weak symbols that point to the profile version and the
         // profile name, which need to be treated as exported so LTO doesn't nix

src/librustc_metadata/creader.rs

@@ -20,7 +20,7 @@ use rustc::hir::svh::Svh;
 use rustc::middle::allocator::AllocatorKind;
 use rustc::middle::cstore::DepKind;
 use rustc::mir::interpret::AllocDecodingState;
-use rustc::session::{Session, CrateDisambiguator};
+use rustc::session::{Session, CrateDisambiguator, InjectedDefaultOomHook};
 use rustc::session::config::{Sanitizer, self};
 use rustc_target::spec::{PanicStrategy, TargetTriple};
 use rustc::session::search_paths::PathKind;
@@ -841,6 +841,7 @@ impl<'a> CrateLoader<'a> {
         if !needs_allocator {
             self.sess.injected_allocator.set(None);
             self.sess.allocator_kind.set(None);
+            self.sess.injected_default_alloc_error_hook.set(InjectedDefaultOomHook::None);
             return
         }
 
@@ -862,9 +863,22 @@ impl<'a> CrateLoader<'a> {
         if !need_lib_alloc && !need_exe_alloc {
             self.sess.injected_allocator.set(None);
             self.sess.allocator_kind.set(None);
+            self.sess.injected_default_alloc_error_hook.set(InjectedDefaultOomHook::None);
             return
         }
 
+        let mut has_default_lib_allocator = attr::contains_name(&krate.attrs,
+                                                                "default_lib_allocator");;
+        self.cstore.iter_crate_data(|_, data| {
+            has_default_lib_allocator = has_default_lib_allocator ||
+                                        data.root.has_default_lib_allocator;
+        });
+        self.sess.injected_default_alloc_error_hook.set(if has_default_lib_allocator {
+            InjectedDefaultOomHook::Platform
+        } else {
+            InjectedDefaultOomHook::Noop
+        });
+
         // Ok, we need an allocator. Not only that but we're actually going to
         // create an artifact that needs one linked in. Let's go find the one
         // that we're going to link in.