Refactor heap allocation

Cargo.toml

@@ -18,6 +18,9 @@ exclude       = [
     "test/**/*"
 ]
 
+[dependencies]
+stable-heap = "0.1.0"
+
 [dev-dependencies]
 rand = "0.3.5"
 

src/alloc/heap.rs

@@ -1,11 +1,17 @@
-use alloc::{Allocator, Mem, MemRef};
-use std::{mem, ptr, usize};
-use std::ops::DerefMut;
+use alloc::{Mem, MemRef};
+use stable_heap as heap;
+use std::{mem, ptr, isize, usize};
+use std::sync::atomic::{self, AtomicUsize, Ordering};
 
 const MAX_ALLOC_SIZE: usize = usize::MAX;
+const MAX_REFCOUNT: usize = (isize::MAX) as usize;
 
 pub struct Heap;
 
+struct Allocation {
+    refs: AtomicUsize,
+}
+
 impl Heap {
     pub fn allocate(&self, len: usize) -> MemRef {
         // Make sure that the allocation is within the permitted range
@@ -13,40 +19,96 @@ impl Heap {
             return MemRef::none();
         }
 
-        let alloc_len = len +
-            mem::size_of::<Mem>() +
-            mem::size_of::<Vec<u8>>();
-
         unsafe {
-            let mut vec: Vec<u8> = Vec::with_capacity(alloc_len);
-            vec.set_len(alloc_len);
+            let mut ptr = heap::allocate(alloc_len(len), align());
+            let mut off = 0;
 
-            let ptr = vec.deref_mut().as_mut_ptr();
+            ptr::write(ptr as *mut Allocation, Allocation::new());
 
-            ptr::write(ptr as *mut Vec<u8>, vec);
+            off += mem::size_of::<Allocation>();
+            ptr::write(ptr.offset(off as isize) as *mut &Mem, &*(ptr as *const Allocation));
 
-            let ptr = ptr.offset(mem::size_of::<Vec<u8>>() as isize);
-            ptr::write(ptr as *mut Mem, Mem::new(len, mem::transmute(self as &Allocator)));
+            off += mem::size_of::<&Mem>();
+            ptr::write(ptr.offset(off as isize) as *mut usize, len);
 
-            // Return the info
-            MemRef::new(ptr as *mut Mem)
+            ptr = ptr.offset(mem::size_of::<Allocation>() as isize);
+
+            MemRef::new(ptr)
         }
     }
 
-    pub fn deallocate(&self, mem: *mut Mem) {
+    fn deallocate(ptr: *mut u8) {
         unsafe {
-            let ptr = mem as *mut u8;
-            let _ = ptr::read(ptr.offset(-(mem::size_of::<Vec<u8>>() as isize)) as *const Vec<u8>);
+            let off = mem::size_of::<Allocation>() + mem::size_of::<&Mem>();
+            let len = ptr::read(ptr.offset(off as isize) as *const usize);
+
+            heap::deallocate(ptr, alloc_len(len), align());
         }
     }
 }
 
-impl Allocator for Heap {
-    fn allocate(&self, len: usize) -> MemRef {
-        Heap::allocate(self, len)
-    }
-
-    fn deallocate(&self, mem: *mut Mem) {
-        Heap::deallocate(self, mem)
+impl Allocation {
+    fn new() -> Allocation {
+        Allocation {
+            refs: AtomicUsize::new(1),
+        }
     }
 }
+
+impl Mem for Allocation {
+    fn ref_inc(&self) {
+        // Using a relaxed ordering is alright here, as knowledge of the
+        // original reference prevents other threads from erroneously deleting
+        // the object.
+        //
+        // As explained in the [Boost documentation][1], Increasing the
+        // reference counter can always be done with memory_order_relaxed: New
+        // references to an object can only be formed from an existing
+        // reference, and passing an existing reference from one thread to
+        // another must already provide any required synchronization.
+        //
+        // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html)
+        let old_size = self.refs.fetch_add(1, Ordering::Relaxed);
+
+        // However we need to guard against massive refcounts in case someone
+        // is `mem::forget`ing Arcs. If we don't do this the count can overflow
+        // and users will use-after free. We racily saturate to `isize::MAX` on
+        // the assumption that there aren't ~2 billion threads incrementing
+        // the reference count at once. This branch will never be taken in
+        // any realistic program.
+        //
+        // We abort because such a program is incredibly degenerate, and we
+        // don't care to support it.
+        if old_size > MAX_REFCOUNT {
+            panic!("too many refs");
+        }
+    }
+
+    fn ref_dec(&self) {
+        if self.refs.fetch_sub(1, Ordering::Release) != 1 {
+            return;
+        }
+
+        atomic::fence(Ordering::Acquire);
+        Heap::deallocate(self as *const Allocation as *const u8 as *mut u8);
+    }
+}
+
+#[inline]
+fn alloc_len(bytes_len: usize) -> usize {
+    let len = bytes_len +
+        mem::size_of::<Allocation>() +
+        mem::size_of::<&Mem>() +
+        mem::size_of::<usize>();
+
+    if len & (align() - 1) == 0 {
+        len
+    } else {
+        (len & !align()) + align()
+    }
+}
+
+#[inline]
+fn align() -> usize {
+    mem::size_of::<usize>()
+}

src/alloc/mod.rs

@@ -1,43 +1,35 @@
 mod heap;
 
-pub use self::heap::{Heap};
+pub use self::heap::Heap;
+pub use self::pool::Pool;
 
 use std::{mem, ptr};
-use std::sync::atomic::{AtomicUsize, Ordering};
 
 pub fn heap(len: usize) -> MemRef {
     Heap.allocate(len)
 }
 
-/// Allocates memory to be used by Bufs or Bytes. Allows allocating memory
-/// using alternate stratgies than the default Rust heap allocator. Also does
-/// not require that allocations are continuous in memory.
-///
-/// For example, an alternate allocator could use a slab of 4kb chunks of
-/// memory and return as many chunks as needed to satisfy the length
-/// requirement.
-pub trait Allocator: Sync + Send {
+pub trait Mem: Send + Sync {
+    /// Increment the ref count
+    fn ref_inc(&self);
 
-  /// Allocate memory. May or may not be contiguous.
-  fn allocate(&self, len: usize) -> MemRef;
-
-  /// Deallocate a chunk of memory
-  fn deallocate(&self, mem: *mut Mem);
+    /// Decrement the ref count
+    fn ref_dec(&self);
 }
 
 pub struct MemRef {
+    // Pointer to the memory
+    // Layout:
+    // - &Mem
+    // - usize (len)
+    // - u8... bytes
     ptr: *mut u8,
 }
 
 impl MemRef {
-    pub fn new(mem: *mut Mem) -> MemRef {
-        let ptr = mem as *mut u8;
-
-        unsafe {
-            MemRef {
-                ptr: ptr.offset(mem::size_of::<Mem>() as isize),
-            }
-        }
+    #[inline]
+    pub unsafe fn new(ptr: *mut u8) -> MemRef {
+        MemRef { ptr: ptr }
     }
 
     #[inline]
@@ -51,14 +43,16 @@ impl MemRef {
     }
 
     #[inline]
-    pub fn ptr(&self) -> *mut u8 {
-        self.ptr
+    pub fn len(&self) -> usize {
+        unsafe { *self.len_ptr() }
     }
 
+    #[inline]
     pub fn bytes(&self) -> &[u8] {
         use std::slice;
+
         unsafe {
-            slice::from_raw_parts(self.ptr(), self.mem().len)
+            slice::from_raw_parts(self.bytes_ptr(), self.len())
         }
     }
 
@@ -66,66 +60,51 @@ impl MemRef {
     pub fn bytes_mut(&mut self) -> &mut [u8] {
         use std::slice;
         unsafe {
-            slice::from_raw_parts_mut(self.ptr(), self.mem().len)
-        }
-    }
-
-    #[inline]
-    fn mem_ptr(&self) -> *mut Mem {
-        unsafe {
-            self.ptr.offset(-(mem::size_of::<Mem>() as isize)) as *mut Mem
+            slice::from_raw_parts_mut(self.bytes_ptr(), self.len())
         }
     }
 
     #[inline]
     fn mem(&self) -> &Mem {
         unsafe {
-            mem::transmute(self.mem_ptr())
+            *(self.ptr as *const &Mem)
         }
     }
+
+    #[inline]
+    unsafe fn len_ptr(&self) -> *mut usize {
+        let off = mem::size_of::<&Mem>();
+        self.ptr.offset(off as isize) as *mut usize
+    }
+
+    #[inline]
+    unsafe fn bytes_ptr(&self) -> *mut u8 {
+        let off = mem::size_of::<&Mem>() + mem::size_of::<usize>();
+        self.ptr.offset(off as isize)
+    }
 }
 
 impl Clone for MemRef {
     #[inline]
     fn clone(&self) -> MemRef {
-        self.mem().refs.fetch_add(1, Ordering::Relaxed);
+        if self.is_none() {
+            return MemRef::none();
+        }
+
+        self.mem().ref_inc();
         MemRef { ptr: self.ptr }
     }
 }
 
 impl Drop for MemRef {
     fn drop(&mut self) {
-        // Guard against the ref having already been dropped
-        if self.ptr.is_null() { return; }
-
-        // Decrement the ref count
-        if 1 == self.mem().refs.fetch_sub(1, Ordering::Relaxed) {
-            // Last ref dropped, free the memory
-            unsafe {
-                let alloc: &Allocator = mem::transmute(self.mem().allocator);
-                alloc.deallocate(self.mem_ptr());
-            }
+        if self.is_none() {
+            return;
         }
+
+        self.mem().ref_dec();
     }
 }
 
 unsafe impl Send for MemRef { }
 unsafe impl Sync for MemRef { }
-
-/// Memory allocated by an Allocator must be prefixed with Mem
-pub struct Mem {
-    // TODO: It should be possible to reduce the size of this struct
-    allocator: *const Allocator,
-    refs: AtomicUsize,
-    len: usize,
-}
-
-impl Mem {
-    pub fn new(len: usize, allocator: *const Allocator) -> Mem {
-        Mem {
-            allocator: allocator,
-            refs: AtomicUsize::new(1),
-            len: len,
-        }
-    }
-}

src/buf/byte.rs

@@ -1,5 +1,5 @@
 use {alloc, Buf, Bytes, MutBuf, SeqByteStr, MAX_CAPACITY};
-use std::{cmp, fmt, ptr};
+use std::{cmp, fmt};
 
 /*
  *
@@ -101,13 +101,9 @@ impl ByteBuf {
     pub fn read_slice(&mut self, dst: &mut [u8]) -> usize {
         let len = cmp::min(dst.len(), self.remaining());
         let cnt = len as u32;
+        let pos = self.pos as usize;
 
-        unsafe {
-            ptr::copy_nonoverlapping(
-                self.mem.ptr().offset(self.pos as isize),
-                dst.as_mut_ptr(),
-                len);
-        }
+        dst[0..len].copy_from_slice(&self.mem.bytes()[pos..pos+len]);
 
         self.pos += cnt;
         len
@@ -295,27 +291,15 @@ impl MutByteBuf {
 
     #[inline]
     pub fn write_slice(&mut self, src: &[u8]) -> usize {
-        let cnt = src.len() as u32;
-        let rem = self.buf.remaining_u32();
+        let cnt = cmp::min(src.len(), self.buf.remaining());
+        let pos = self.buf.pos as usize;
 
-        if rem < cnt {
-            self.write_ptr(src.as_ptr(), rem)
-        } else {
-            self.write_ptr(src.as_ptr(), cnt)
-        }
-    }
+        self.buf.mem.bytes_mut()[pos..pos+cnt]
+            .copy_from_slice(&src[0..cnt]);
 
-    #[inline]
-    fn write_ptr(&mut self, src: *const u8, len: u32) -> usize {
-        unsafe {
-            ptr::copy_nonoverlapping(
-                src,
-                self.buf.mem.ptr().offset(self.buf.pos as isize),
-                len as usize);
+        self.buf.pos += cnt as u32;
 
-            self.buf.pos += len;
-            len as usize
-        }
+        cnt
     }
 
     pub fn bytes<'a>(&'a self) -> &'a [u8] {

src/buf/ring.rs

@@ -1,5 +1,5 @@
 use {alloc, Buf, MutBuf};
-use std::{cmp, fmt, ptr};
+use std::{cmp, fmt};
 
 enum Mark {
     NoMark,
@@ -137,35 +137,6 @@ impl RingBuf {
     }
 }
 
-impl Clone for RingBuf {
-    fn clone(&self) -> RingBuf {
-        use std::cmp;
-
-        let mut ret = RingBuf::new(self.cap);
-
-        ret.pos = self.pos;
-        ret.len = self.len;
-
-        unsafe {
-            let to = self.pos + self.len;
-
-            if to > self.cap {
-                ptr::copy(self.ptr.ptr() as *const u8, ret.ptr.ptr(), to % self.cap);
-            }
-
-            ptr::copy(
-                self.ptr.ptr().offset(self.pos as isize) as *const u8,
-                ret.ptr.ptr().offset(self.pos as isize),
-                cmp::min(self.len, self.cap - self.pos));
-        }
-
-        ret
-    }
-
-    // TODO: an improved version of clone_from is possible that potentially
-    // re-uses the buffer
-}
-
 impl fmt::Debug for RingBuf {
     fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
         write!(fmt, "RingBuf[.. {}]", self.len)

src/lib.rs

@@ -1,6 +1,8 @@
 #![crate_name = "bytes"]
 #![deny(warnings)]
 
+extern crate stable_heap;
+
 pub mod alloc;
 pub mod buf;
 pub mod str;

src/str/seq.rs

@@ -81,11 +81,7 @@ impl ops::Index<usize> for SeqByteStr {
 
     fn index(&self, index: usize) -> &u8 {
         assert!(index < self.len());
-
-        unsafe {
-            &*self.mem.ptr()
-                .offset(index as isize + self.pos as isize)
-        }
+        self.mem.bytes().index(index + self.pos as usize)
     }
 }