Improve TypedArena's chunk allocation scheme

Closes rust-lang#17931
Fixes rust-lang#18037

Author: pczarn

src/libarena/lib.rs

@@ -31,9 +31,10 @@
 #![feature(alloc)]
 #![feature(box_syntax)]
 #![feature(core_intrinsics)]
+#![feature(drop_in_place)]
+#![feature(raw)]
 #![feature(heap_api)]
 #![feature(oom)]
-#![feature(ptr_as_ref)]
 #![feature(raw)]
 #![feature(staged_api)]
 #![feature(dropck_parametricity)]
@@ -47,9 +48,13 @@ use std::intrinsics;
 use std::marker;
 use std::mem;
 use std::ptr;
+use std::raw;
+use std::raw::Repr;
 use std::rc::Rc;
+use std::slice;
 
-use alloc::heap::{allocate, deallocate};
+use alloc::heap;
+use alloc::raw_vec::RawVec;
 
 // The way arena uses arrays is really deeply awful. The arrays are
 // allocated, and have capacities reserved, but the fill for the array
@@ -366,115 +371,87 @@ fn test_arena_destructors_fail() {
 /// A faster arena that can hold objects of only one type.
 pub struct TypedArena<T> {
     /// A pointer to the next object to be allocated.
-    ptr: Cell<*const T>,
+    ptr: Cell<*mut T>,
 
     /// A pointer to the end of the allocated area. When this pointer is
     /// reached, a new chunk is allocated.
-    end: Cell<*const T>,
+    end: Cell<*mut T>,
 
-    /// A pointer to the first arena segment.
-    first: RefCell<*mut TypedArenaChunk<T>>,
+    /// A vector arena segments.
+    chunks: RefCell<Vec<TypedArenaChunk<T>>>,
 
     /// Marker indicating that dropping the arena causes its owned
     /// instances of `T` to be dropped.
     _own: marker::PhantomData<T>,
 }
 
 struct TypedArenaChunk<T> {
-    marker: marker::PhantomData<T>,
-
     /// Pointer to the next arena segment.
-    next: *mut TypedArenaChunk<T>,
-
-    /// The number of elements that this chunk can hold.
-    capacity: usize,
-
-    // Objects follow here, suitably aligned.
-}
-
-fn calculate_size<T>(capacity: usize) -> usize {
-    let mut size = mem::size_of::<TypedArenaChunk<T>>();
-    size = round_up(size, mem::align_of::<T>());
-    let elem_size = mem::size_of::<T>();
-    let elems_size = elem_size.checked_mul(capacity).unwrap();
-    size = size.checked_add(elems_size).unwrap();
-    size
+    storage: RawVec<T>,
 }
 
 impl<T> TypedArenaChunk<T> {
     #[inline]
-    unsafe fn new(next: *mut TypedArenaChunk<T>, capacity: usize) -> *mut TypedArenaChunk<T> {
-        let size = calculate_size::<T>(capacity);
-        let chunk =
-            allocate(size, mem::align_of::<TypedArenaChunk<T>>()) as *mut TypedArenaChunk<T>;
-        if chunk.is_null() {
-            alloc::oom()
-        }
-        (*chunk).next = next;
-        (*chunk).capacity = capacity;
-        chunk
+    unsafe fn new(capacity: usize) -> TypedArenaChunk<T> {
+        TypedArenaChunk { storage: RawVec::with_capacity(capacity) }
     }
 
-    /// Destroys this arena chunk. If the type descriptor is supplied, the
-    /// drop glue is called; otherwise, drop glue is not called.
+    /// Destroys this arena chunk.
     #[inline]
     unsafe fn destroy(&mut self, len: usize) {
-        // Destroy all the allocated objects.
+        // The branch on needs_drop() is an -O1 performance optimization.
+        // Without the branch, dropping TypedArena<u8> takes linear time.
         if intrinsics::needs_drop::<T>() {
             let mut start = self.start();
+            // Destroy all allocated objects.
             for _ in 0..len {
-                ptr::read(start as *const T); // run the destructor on the pointer
-                start = start.offset(mem::size_of::<T>() as isize)
+                ptr::drop_in_place(start);
+                start = start.offset(1);
             }
         }
-
-        // Destroy the next chunk.
-        let next = self.next;
-        let size = calculate_size::<T>(self.capacity);
-        let self_ptr: *mut TypedArenaChunk<T> = self;
-        deallocate(self_ptr as *mut u8,
-                   size,
-                   mem::align_of::<TypedArenaChunk<T>>());
-        if !next.is_null() {
-            let capacity = (*next).capacity;
-            (*next).destroy(capacity);
-        }
     }
 
     // Returns a pointer to the first allocated object.
     #[inline]
-    fn start(&self) -> *const u8 {
-        let this: *const TypedArenaChunk<T> = self;
-        unsafe { round_up(this.offset(1) as usize, mem::align_of::<T>()) as *const u8 }
+    fn start(&self) -> *mut T {
+        self.storage.ptr()
     }
 
     // Returns a pointer to the end of the allocated space.
     #[inline]
-    fn end(&self) -> *const u8 {
+    fn end(&self) -> *mut T {
         unsafe {
-            let size = mem::size_of::<T>().checked_mul(self.capacity).unwrap();
-            self.start().offset(size as isize)
+            if mem::size_of::<T>() == 0 {
+                // A pointer as large as possible for zero-sized elements.
+                !0 as *mut T
+            } else {
+                self.start().offset(self.storage.cap() as isize)
+            }
         }
     }
 }
 
+const PAGE: usize = 4096;
+
 impl<T> TypedArena<T> {
-    /// Creates a new `TypedArena` with preallocated space for eight objects.
+    /// Creates a new `TypedArena` with preallocated space for many objects.
     #[inline]
     pub fn new() -> TypedArena<T> {
-        TypedArena::with_capacity(8)
+        // Reserve at least one page.
+        let elem_size = cmp::max(1, mem::size_of::<T>());
+        TypedArena::with_capacity(PAGE / elem_size)
     }
 
     /// Creates a new `TypedArena` with preallocated space for the given number of
     /// objects.
     #[inline]
     pub fn with_capacity(capacity: usize) -> TypedArena<T> {
         unsafe {
-            let chunk = TypedArenaChunk::<T>::new(ptr::null_mut(), capacity);
+            let chunk = TypedArenaChunk::<T>::new(cmp::max(1, capacity));
             TypedArena {
-                ptr: Cell::new((*chunk).start() as *const T),
-                end: Cell::new((*chunk).end() as *const T),
-                first: RefCell::new(chunk),
+                ptr: Cell::new(chunk.start()),
+                end: Cell::new(chunk.end()),
+                chunks: RefCell::new(vec![chunk]),
                 _own: marker::PhantomData,
             }
         }
@@ -488,23 +465,39 @@ impl<T> TypedArena<T> {
         }
 
         unsafe {
-            let ptr: &mut T = &mut *(self.ptr.get() as *mut T);
-            ptr::write(ptr, object);
-            self.ptr.set(self.ptr.get().offset(1));
-            ptr
+            if mem::size_of::<T>() == 0 {
+                self.ptr.set(intrinsics::arith_offset(self.ptr.get() as *mut u8, 1) as *mut T);
+                let ptr = heap::EMPTY as *mut T;
+                // Don't drop the object. This `write` is equivalent to `forget`.
+                ptr::write(ptr, object);
+                &mut *ptr
+            } else {
+                let ptr = self.ptr.get();
+                // Advance the pointer.
+                self.ptr.set(self.ptr.get().offset(1));
+                // Write into uninitialized memory.
+                ptr::write(ptr, object);
+                &mut *ptr
+            }
         }
     }
 
     /// Grows the arena.
     #[inline(never)]
+    #[cold]
     fn grow(&self) {
         unsafe {
-            let chunk = *self.first.borrow_mut();
-            let new_capacity = (*chunk).capacity.checked_mul(2).unwrap();
-            let chunk = TypedArenaChunk::<T>::new(chunk, new_capacity);
-            self.ptr.set((*chunk).start() as *const T);
-            self.end.set((*chunk).end() as *const T);
-            *self.first.borrow_mut() = chunk
+            let mut chunks = self.chunks.borrow_mut();
+            let prev_capacity = chunks.last().unwrap().storage.cap();
+            let new_capacity = prev_capacity.checked_mul(2).unwrap();
+            if chunks.last_mut().unwrap().storage.double_in_place() {
+                self.end.set(chunks.last().unwrap().end());
+            } else {
+                let chunk = TypedArenaChunk::<T>::new(new_capacity);
+                self.ptr.set(chunk.start());
+                self.end.set(chunk.end());
+                chunks.push(chunk);
+            }
         }
     }
 }
@@ -514,12 +507,26 @@ impl<T> Drop for TypedArena<T> {
     fn drop(&mut self) {
         unsafe {
             // Determine how much was filled.
-            let start = self.first.borrow().as_ref().unwrap().start() as usize;
+            let mut chunks_borrow = self.chunks.borrow_mut();
+            let mut last_chunk = chunks_borrow.pop().unwrap();
+            let start = last_chunk.start() as usize;
             let end = self.ptr.get() as usize;
-            let diff = (end - start) / mem::size_of::<T>();
+            let diff = if mem::size_of::<T>() == 0 {
+                // Avoid division by zero.
+                end - start
+            } else {
+                (end - start) / mem::size_of::<T>()
+            };
 
             // Pass that to the `destroy` method.
-            (**self.first.borrow_mut()).destroy(diff)
+            last_chunk.destroy(diff);
+            // Destroy this chunk.
+            let _: RawVec<T> = mem::transmute(last_chunk);
+
+            for chunk in chunks_borrow.iter_mut() {
+                let cap = chunk.storage.cap();
+                chunk.destroy(cap);
+            }
         }
     }
 }
@@ -657,4 +664,12 @@ mod tests {
             })
         })
     }
+
+    #[test]
+    pub fn test_zero_sized() {
+        let arena = TypedArena::new();
+        for _ in 0..100000 {
+            arena.alloc(());
+        }
+    }
 }