Eliminate a 256-byte memcpy in a hot-path

This one would be a nasty bug if it wasn't caught. The problem was using
`MoveCell` which doesn't allow modification inplace, leading to memcpys
for getting ownership of the value. We replace it by `RefCell` and
remove the `MoveCell` implementation, resulting in a huge performance
boost.

- Move `Block::sbrk` into the main `impl` block for `Block` to allow
  usage in `Arena`'s tests.

src/allocator.rs

@@ -2,8 +2,7 @@
 //!
 //! This contains primitives for the cross-thread allocator.
 
-use prelude::*;
-
+use core::cell::RefCell;
 use core::mem;
 
 use {brk, sync};
@@ -15,8 +14,10 @@ use shim::config;
 use tls;
 
 /// Alias for the wrapper type of the thread-local variable holding the local allocator.
+// FIXME: Replace RefCell with a 2-state primitive (used and not used). Ideally, the mutual
+//        exclusion should be statically ensured (see RFC #1106).
 #[cfg(feature = "tls")]
-type ThreadLocalAllocator = MoveCell<Option<LazyInit<fn() -> LocalAllocator, LocalAllocator>>>;
+type ThreadLocalAllocator = RefCell<Option<LazyInit<fn() -> LocalAllocator, LocalAllocator>>>;
 
 /// The global default allocator.
 // TODO: Remove these filthy function pointers.
@@ -25,7 +26,7 @@ static GLOBAL_ALLOCATOR: sync::Mutex<LazyInit<fn() -> GlobalAllocator, GlobalAll
 #[cfg(feature = "tls")]
 tls! {
     /// The thread-local allocator.
-    static THREAD_ALLOCATOR: ThreadLocalAllocator = MoveCell::new(Some(LazyInit::new(LocalAllocator::init)));
+    static THREAD_ALLOCATOR: ThreadLocalAllocator = RefCell::new(Some(LazyInit::new(LocalAllocator::init)));
 }
 
 /// Temporarily get the allocator.
@@ -42,22 +43,21 @@ tls! {
 //       it run after the TLS keys that might be declared.
 macro_rules! get_allocator {
     (|$v:ident| $b:expr) => {{
-        // Get the thread allocator, if TLS is enabled
+        // Get the thread allocator, if TLS is enabled.
         #[cfg(feature = "tls")]
         {
             THREAD_ALLOCATOR.with(|thread_alloc| {
-                if let Some(mut thread_alloc_original) = thread_alloc.replace(None) {
+                if let Ok(mut thread_alloc_original) = thread_alloc.try_borrow_mut() {
                     let ret = {
                         // Call the closure involved.
                         let $v = thread_alloc_original.get();
                         $b
                     };
 
-                    // Put back the original allocator.
-                    thread_alloc.replace(Some(thread_alloc_original));
-
                     ret
                 } else {
+                    // TODO: Test if this fallback can be removed.
+
                     // The local allocator seems to have been deinitialized, for this reason we fallback to
                     // the global allocator.
                     log!(WARNING, "Accessing the allocator after deinitialization of the local allocator.");

src/arena.rs

@@ -169,7 +169,7 @@ mod test {
     /// Helper method to make an artificial arena.
     fn make<T>() -> Arena<T> {
         let mut arena = Arena::new();
-        arena.provide(brk::lock().sbrk(1024).unwrap());
+        arena.provide(Block::sbrk(826));
 
         arena
     }

src/block.rs

@@ -285,6 +285,12 @@ impl Block {
 
         self
     }
+
+    /// Create a new block by extending the program break.
+    #[cfg(test)]
+    pub fn sbrk(size: Size) -> Block {
+        Block::from_raw_parts(brk::lock().sbrk(size.try_into()).unwrap(), size)
+    }
 }
 
 impl From<Block> for Pointer<u8> {
@@ -336,15 +342,6 @@ mod test {
 
     use brk;
 
-    /// Implementation we will use for testing.
-    impl Block {
-        /// Create a new block by extending the program break.
-        #[cfg(test)]
-        pub fn sbrk(size: Size) -> Block {
-            Block::from_raw_parts(brk::lock().sbrk(size.try_into()).unwrap(), size)
-        }
-    }
-
     #[test]
     fn split() {
         let block = Block::sbrk(Size(26));

src/cell.rs

@@ -1,73 +0,0 @@
-//! Cell primitives.
-
-use core::cell::UnsafeCell;
-use core::mem;
-
-use take::take;
-
-/// A move cell.
-///
-/// This allows you to take ownership and replace the internal data with a new value. The
-/// functionality is similar to the one provided by [RFC #1659](https://github.com/rust-lang/rfcs/pull/1659).
-// TODO: Use the features provided by the RFC.
-pub struct MoveCell<T> {
-    /// The inner data.
-    inner: UnsafeCell<T>,
-}
-
-impl<T> MoveCell<T> {
-    /// Create a new cell with some inner data.
-    #[inline]
-    pub const fn new(data: T) -> MoveCell<T> {
-        MoveCell {
-            inner: UnsafeCell::new(data),
-        }
-    }
-
-    /// Replace the inner data and return the old.
-    #[inline]
-    pub fn replace(&self, new: T) -> T {
-        mem::replace(unsafe {
-            // LAST AUDIT: 2016-08-21 (Ticki).
-
-            // This is safe due to never aliasing the value, but simply transferring ownership to
-            // the caller.
-            &mut *self.inner.get()
-        }, new)
-    }
-
-    /// Get a reference to the inner value.
-    ///
-    /// Safety is enforced statically due to the guarantee of mutual exclusion in mutable
-    /// references.
-    pub fn get(&mut self) -> &mut T {
-        mem::replace(unsafe {
-            // LAST AUDIT: 2016-09-01 (Ticki).
-
-            // This is safe due to the `&mut self`, enforcing the guarantee of uniqueness. This
-            // will thus not alias it for the lifetime of that reference.
-            &mut *self.inner.get()
-        })
-    }
-}
-
-#[cfg(test)]
-mod test {
-    use super::MoveCell;
-
-    #[test]
-    fn replace() {
-        let cell = MoveCell::new(200);
-        assert_eq!(cell.replace(300), 200);
-        assert_eq!(cell.replace(4), 300);
-    }
-
-    #[test]
-    fn get() {
-        let mut cell = MoveCell::new(200);
-        assert_eq!(*cell.get(), 200);
-
-        *cell.get() = 300;
-        assert_eq!(*cell.get(), 300);
-    }
-}

src/fail.rs

@@ -1,9 +1,8 @@
 //! General error handling.
 
-use prelude::*;
-
-use core::sync::atomic::{self, AtomicPtr};
+use core::cell::Cell;
 use core::mem;
+use core::sync::atomic::{self, AtomicPtr};
 
 use shim::config;
 
@@ -15,7 +14,7 @@ static OOM_HANDLER: AtomicPtr<()> = AtomicPtr::new(config::default_oom_handler a
 #[cfg(feature = "tls")]
 tls! {
     /// The thread-local OOM handler.
-    static THREAD_OOM_HANDLER: MoveCell<Option<fn() -> !>> = MoveCell::new(None);
+    static THREAD_OOM_HANDLER: Cell<Option<fn() -> !>> = Cell::new(None);
 }
 
 /// Call the OOM handler.
@@ -34,7 +33,7 @@ pub fn oom() -> ! {
     // If TLS is enabled, we will use the thread-local OOM.
     #[cfg(feature = "tls")]
     {
-        if let Some(handler) = THREAD_OOM_HANDLER.with(|x| x.replace(None)) {
+        if let Some(handler) = THREAD_OOM_HANDLER.with(|x| x.get()) {
             log!(DEBUG, "Calling the local OOM handler.");
 
             handler();
@@ -73,12 +72,9 @@ pub fn set_thread_oom_handler(handler: fn() -> !) {
 
     THREAD_OOM_HANDLER.with(|thread_oom| {
         // Replace it with the new handler.
-        let old = thread_oom.replace(Some(handler));
+        thread_oom.set(Some(handler));
 
-        // Throw a warning if it overrides another handler.
-        if old.is_some() {
-            log!(WARNING, "An old thread OOM handler was overriden.");
-        }
+        // TODO: Bring back the overwrite check (i.e. warn when overwriting the handler).
     });
 }
 

src/lib.rs

@@ -44,7 +44,6 @@ mod arena;
 mod bk;
 mod block;
 mod brk;
-mod cell;
 mod fail;
 mod lazy_init;
 mod leak;

src/take.rs

@@ -7,6 +7,8 @@ use core::{mem, intrinsics};
 /// A guarding type which will exit upon drop.
 ///
 /// This is used for catching unwinding and transforming it into abort.
+///
+/// The destructor should never be called naturally (use `mem::forget()`), and only when unwinding.
 struct ExitGuard;
 
 impl Drop for ExitGuard {
@@ -41,7 +43,7 @@ pub fn replace_with<T, F>(val: &mut T, replace: F)
         ptr::write(val, new);
     }
 
-    // Drop the guard.
+    // Forget the guard.
     mem::forget(guard);
 }
 

src/tls.rs

@@ -84,7 +84,7 @@ macro_rules! tls {
 
 #[cfg(test)]
 mod test {
-    use cell::MoveCell;
+    use core::cell::Cell;
 
     #[test]
     fn tls() {
@@ -95,12 +95,16 @@ mod test {
 
     #[test]
     fn mutability() {
-        tls!(static HELLO: MoveCell<u32> = MoveCell::new(3));
-
-        HELLO.with(|x| assert_eq!(x.replace(4), 3));
-        HELLO.with(|x| assert_eq!(x.replace(5), 4));
-        HELLO.with(|x| assert_eq!(x.replace(10), 5));
-        HELLO.with(|x| assert_eq!(x.replace(0), 10));
-        HELLO.with(|x| assert_eq!(x.replace(0), 0));
+        tls!(static HELLO: Cell<u32> = Cell::new(3));
+
+        HELLO.with(|x| assert_eq!(x.get(), 3));
+        HELLO.with(|x| x.set(4));
+        HELLO.with(|x| assert_eq!(x.get(), 4));
+        HELLO.with(|x| x.set(5));
+        HELLO.with(|x| assert_eq!(x.get(), 5));
+        HELLO.with(|x| x.set(10));
+        HELLO.with(|x| assert_eq!(x.get(), 10));
+        HELLO.with(|x| x.set(0));
+        HELLO.with(|x| assert_eq!(x.get(), 0));
     }
 }