1 files changed, 49 insertions, 64 deletions

diff --git a/ctr-std/src/alloc.rs b/ctr-std/src/alloc.rs
index a857840..3b1a3a4 100644
--- a/ctr-std/src/alloc.rs
+++ b/ctr-std/src/alloc.rs
@@ -13,16 +13,61 @@
 #![unstable(issue = "32838", feature = "allocator_api")]
 
 #[doc(inline)] #[allow(deprecated)] pub use alloc_crate::alloc::Heap;
-#[doc(inline)] pub use alloc_crate::alloc::{Global, oom};
+#[doc(inline)] pub use alloc_crate::alloc::{Global, Layout, oom};
 #[doc(inline)] pub use alloc_system::System;
 #[doc(inline)] pub use core::alloc::*;
 
-#[cfg(not(stage0))]
+use core::sync::atomic::{AtomicPtr, Ordering};
+use core::{mem, ptr};
+use sys_common::util::dumb_print;
+
+static HOOK: AtomicPtr<()> = AtomicPtr::new(ptr::null_mut());
+
+/// Registers a custom OOM hook, replacing any that was previously registered.
+///
+/// The OOM 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_oom_hook`] and
+/// [`take_oom_hook`] functions.
+///
+/// The hook is provided with a `Layout` struct which contains information
+/// about the allocation that failed.
+///
+/// The OOM hook is a global resource.
+pub fn set_oom_hook(hook: fn(Layout)) {
+    HOOK.store(hook as *mut (), Ordering::SeqCst);
+}
+
+/// Unregisters the current OOM hook, returning it.
+///
+/// *See also the function [`set_oom_hook`].*
+///
+/// If no custom hook is registered, the default hook will be returned.
+pub fn take_oom_hook() -> fn(Layout) {
+    let hook = HOOK.swap(ptr::null_mut(), Ordering::SeqCst);
+    if hook.is_null() {
+        default_oom_hook
+    } else {
+        unsafe { mem::transmute(hook) }
+    }
+}
+
+fn default_oom_hook(layout: Layout) {
+    dumb_print(format_args!("memory allocation of {} bytes failed", layout.size()));
+}
+
 #[cfg(not(test))]
 #[doc(hidden)]
 #[lang = "oom"]
-pub extern fn rust_oom() -> ! {
-    rtabort!("memory allocation failed");
+pub extern fn rust_oom(layout: Layout) -> ! {
+    let hook = HOOK.load(Ordering::SeqCst);
+    let hook: fn(Layout) = if hook.is_null() {
+        default_oom_hook
+    } else {
+        unsafe { mem::transmute(hook) }
+    };
+    hook(layout);
+    unsafe { ::sys::abort_internal(); }
 }
 
 #[cfg(not(test))]
@@ -43,13 +88,6 @@ pub mod __default_lib_allocator {
         System.alloc(layout) as *mut u8
     }
 
-    #[cfg(stage0)]
-    #[no_mangle]
-    #[rustc_std_internal_symbol]
-    pub unsafe extern fn __rdl_oom() -> ! {
-        super::oom()
-    }
-
     #[no_mangle]
     #[rustc_std_internal_symbol]
     pub unsafe extern fn __rdl_dealloc(ptr: *mut u8,
@@ -74,57 +112,4 @@ pub mod __default_lib_allocator {
         let layout = Layout::from_size_align_unchecked(size, align);
         System.alloc_zeroed(layout) as *mut u8
     }
-
-    #[cfg(stage0)]
-    pub mod stage0 {
-        #[no_mangle]
-        #[rustc_std_internal_symbol]
-        pub unsafe extern fn __rdl_usable_size(_layout: *const u8,
-                                               _min: *mut usize,
-                                               _max: *mut usize) {
-            unimplemented!()
-        }
-
-        #[no_mangle]
-        #[rustc_std_internal_symbol]
-        pub unsafe extern fn __rdl_alloc_excess(_size: usize,
-                                                _align: usize,
-                                                _excess: *mut usize,
-                                                _err: *mut u8) -> *mut u8 {
-            unimplemented!()
-        }
-
-        #[no_mangle]
-        #[rustc_std_internal_symbol]
-        pub unsafe extern fn __rdl_realloc_excess(_ptr: *mut u8,
-                                                  _old_size: usize,
-                                                  _old_align: usize,
-                                                  _new_size: usize,
-                                                  _new_align: usize,
-                                                  _excess: *mut usize,
-                                                  _err: *mut u8) -> *mut u8 {
-            unimplemented!()
-        }
-
-        #[no_mangle]
-        #[rustc_std_internal_symbol]
-        pub unsafe extern fn __rdl_grow_in_place(_ptr: *mut u8,
-                                                 _old_size: usize,
-                                                 _old_align: usize,
-                                                 _new_size: usize,
-                                                 _new_align: usize) -> u8 {
-            unimplemented!()
-        }
-
-        #[no_mangle]
-        #[rustc_std_internal_symbol]
-        pub unsafe extern fn __rdl_shrink_in_place(_ptr: *mut u8,
-                                                   _old_size: usize,
-                                                   _old_align: usize,
-                                                   _new_size: usize,
-                                                   _new_align: usize) -> u8 {
-            unimplemented!()
-        }
-
-    }
 }