Simplify allocation strategy for now

Not having `unsafe_no_drop_flag` caused some weirdness with optimizing buffers
and bytes. For now, remeove it.

src/alloc/heap.rs

@@ -1,111 +1,9 @@
-use alloc::{Mem, MemRef};
+use alloc::{MemRef};
-use stable_heap as heap;
+use std::sync::Arc;
-use std::{mem, ptr, isize, usize};
-use std::sync::atomic::{self, AtomicUsize, Ordering};
 
-const MAX_ALLOC_SIZE: usize = usize::MAX;
+pub unsafe fn allocate(len: usize) -> MemRef {
-const MAX_REFCOUNT: usize = (isize::MAX) as usize;
+    let mut v = Vec::with_capacity(len);
+    v.set_len(len);
 
-/// Tracks a heap allocation and stores the atomic ref counter
+    MemRef::new(Arc::new(v))
-struct Allocation {
-    refs: AtomicUsize,
-}
-
-pub fn allocate(len: usize) -> MemRef {
-    // Make sure that the allocation is within the permitted range
-    if len > MAX_ALLOC_SIZE {
-        return MemRef::none();
-    }
-
-    unsafe {
-        let mut ptr = heap::allocate(alloc_len(len), align());
-        let mut off = 0;
-
-        ptr::write(ptr as *mut Allocation, Allocation::new());
-
-        off += mem::size_of::<Allocation>();
-        ptr::write(ptr.offset(off as isize) as *mut &Mem, &*(ptr as *const Allocation));
-
-        off += mem::size_of::<&Mem>();
-        ptr::write(ptr.offset(off as isize) as *mut usize, len);
-
-        ptr = ptr.offset(mem::size_of::<Allocation>() as isize);
-
-        MemRef::new(ptr)
-    }
-}
-
-fn deallocate(ptr: *mut u8) {
-    unsafe {
-        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 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);
-        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,122 +1,58 @@
 mod heap;
-mod pool;
 
-pub use self::pool::Pool;
+use std::sync::Arc;
-
-use std::{mem, ptr};
-
-/// Ref-counted segment of memory
-pub trait Mem: Send + Sync {
-    /// Increment the ref count
-    fn ref_inc(&self);
-
-    /// Decrement the ref count
-    fn ref_dec(&self);
-}
 
 pub struct MemRef {
-    // Pointer to the memory
+    mem: Arc<Vec<u8>>,
-    // Layout:
-    // - &Mem
-    // - usize (len)
-    // - u8... bytes
-    ptr: *mut u8,
 }
 
 /// Allocate a segment of memory and return a `MemRef`.
-pub fn heap(len: usize) -> MemRef {
+pub unsafe fn heap(len: usize) -> MemRef {
     heap::allocate(len)
 }
 
 impl MemRef {
     #[inline]
-    pub unsafe fn new(ptr: *mut u8) -> MemRef {
+    pub unsafe fn new(mem: Arc<Vec<u8>>) -> MemRef {
-        MemRef { ptr: ptr }
+        MemRef { mem: mem }
-    }
-
-    #[inline]
-    pub fn none() -> MemRef {
-        MemRef { ptr: ptr::null_mut() }
-    }
-
-    #[inline]
-    pub fn is_none(&self) -> bool {
-        self.ptr.is_null()
     }
 
     #[inline]
     pub fn len(&self) -> usize {
-        unsafe { *self.len_ptr() }
+        self.mem.len()
     }
 
     #[inline]
     pub unsafe fn bytes(&self) -> &[u8] {
-        use std::slice;
+        &*self.mem
-        slice::from_raw_parts(self.bytes_ptr(), self.len())
     }
 
     #[inline]
     pub unsafe fn bytes_slice(&self, start: usize, end: usize) -> &[u8] {
         use std::slice;
-        let ptr = self.bytes_ptr().offset(start as isize);
+        let ptr = self.mem.as_ptr().offset(start as isize);
         slice::from_raw_parts(ptr, end - start)
     }
 
     #[inline]
     pub unsafe fn mut_bytes(&mut self) -> &mut [u8] {
         use std::slice;
-        slice::from_raw_parts_mut(self.bytes_ptr(), self.len())
+        let len = self.mem.len();
+        slice::from_raw_parts_mut(self.mem.as_ptr() as *mut u8, len)
     }
 
     /// Unsafe, unchecked access to the bytes
     #[inline]
     pub unsafe fn mut_bytes_slice(&mut self, start: usize, end: usize) -> &mut [u8] {
         use std::slice;
-        let ptr = self.bytes_ptr().offset(start as isize);
+        let ptr = self.mem.as_ptr().offset(start as isize);
-        slice::from_raw_parts_mut(ptr, end - start)
+        slice::from_raw_parts_mut(ptr as *mut u8, end - start)
-    }
-
-    #[inline]
-    fn mem(&self) -> &Mem {
-        unsafe {
-            *(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 {
-        if self.is_none() {
+        MemRef { mem: self.mem.clone() }
-            return MemRef::none();
-        }
-
-        self.mem().ref_inc();
-        MemRef { ptr: self.ptr }
     }
 }
-
-impl Drop for MemRef {
-    fn drop(&mut self) {
-        if self.is_none() {
-            return;
-        }
-
-        self.mem().ref_dec();
-    }
-}
-
-unsafe impl Send for MemRef { }
-unsafe impl Sync for MemRef { }

src/buf/append.rs

@@ -16,32 +16,14 @@ pub struct AppendBuf {
 
 impl AppendBuf {
     pub fn with_capacity(mut capacity: u32) -> AppendBuf {
-        // Handle 0 capacity case
-        if capacity == 0 {
-            return AppendBuf::none();
-        }
-
         // Round the capacity to the closest power of 2
         capacity = capacity.next_power_of_two();
 
-        // Allocate the memory
+        unsafe {
-        let mem = alloc::heap(capacity as usize);
+            // Allocate the memory
+            let mem = alloc::heap(capacity as usize);
 
-        // If the allocation failed, return a blank buf
+            AppendBuf::from_mem_ref(mem, capacity, 0)
-        if mem.is_none() {
-            return AppendBuf::none();
-        }
-
-        unsafe { AppendBuf::from_mem_ref(mem, capacity, 0) }
-    }
-
-    /// Returns an AppendBuf with no capacity
-    pub fn none() -> AppendBuf {
-        AppendBuf {
-            mem: alloc::MemRef::none(),
-            rd: Cell::new(0),
-            wr: 0,
-            cap: 0,
         }
     }
 
@@ -73,6 +55,7 @@ impl AppendBuf {
     pub fn shift(&self, n: usize) -> Bytes {
         let ret = self.slice(0, n);
         self.rd.set(self.rd.get() + ret.len() as u32);
+        assert!(self.rd.get() <= self.wr, "buffer overflow");
         ret
     }
 
@@ -82,13 +65,15 @@ impl AppendBuf {
     }
 
     pub fn slice(&self, begin: usize, end: usize) -> Bytes {
-        let rd = self.rd.get() as usize;
+        // TODO: Fix overflow potential
-        let wr = self.wr as usize;
 
-        assert!(begin <= end && end <= wr - rd, "invalid range");
+        let rd = self.rd.get();
+        let wr = self.wr;
 
-        let begin = (begin + rd) as u32;
+        let begin = begin as u32 + rd;
-        let end = (end + rd) as u32;
+        let end = end as u32 + rd;
+
+        assert!(begin <= end && end <= wr, "invalid range");
 
         unsafe { Bytes::from_mem_ref(self.mem.clone(), begin, end - begin) }
     }
@@ -111,7 +96,7 @@ impl MutBuf for AppendBuf {
         self.wr += cnt as u32;
 
         if self.wr > self.cap {
-            self.wr = self.cap;
+            panic!("buffer overflow");
         }
     }
 

src/buf/block.rs

@@ -1,7 +1,7 @@
 #![allow(warnings)]
 
 use {Buf, MutBuf, AppendBuf, Bytes};
-use alloc::{self, Pool};
+use alloc::{self, /* Pool */};
 use std::{cmp, ptr, slice};
 use std::io::Cursor;
 use std::rc::Rc;
@@ -19,9 +19,9 @@ pub struct BlockBuf {
     new_block: NewBlock,
 }
 
-pub enum NewBlock {
+enum NewBlock {
     Heap(usize),
-    Pool(Rc<Pool>),
+    // Pool(Rc<Pool>),
 }
 
 pub struct BlockBufCursor<'a> {
@@ -37,9 +37,11 @@ pub struct BlockBufCursor<'a> {
 //
 impl BlockBuf {
     /// Create BlockBuf
-    pub fn new(max_blocks: usize, new_block: NewBlock) -> BlockBuf {
+    pub fn new(max_blocks: usize, block_size: usize) -> BlockBuf {
         assert!(max_blocks > 1, "at least 2 blocks required");
 
+        let new_block = NewBlock::Heap(block_size);
+
         BlockBuf {
             len: 0,
             cap: max_blocks * new_block.block_size(),
@@ -51,7 +53,7 @@ impl BlockBuf {
     /// Returns the number of buffered bytes
     #[inline]
     pub fn len(&self) -> usize {
-        debug_assert!(self.len == self.blocks.iter().map(|b| b.len()).fold(0, |a, b| a+b));
+        debug_assert_eq!(self.len, self.blocks.iter().map(|b| b.len()).fold(0, |a, b| a+b));
         self.len
     }
 
@@ -83,9 +85,39 @@ impl BlockBuf {
     /// # Panics
     ///
     /// Panics if `n` is greater than the number of buffered bytes.
-    pub fn shift(&mut self, mut n: usize) -> Bytes {
+    #[inline]
+    pub fn shift(&mut self, n: usize) -> Bytes {
         trace!("BlockBuf::shift; n={}", n);
 
+        // Fast path
+        match self.blocks.len() {
+            0 => {
+                assert!(n == 0, "buffer overflow");
+                Bytes::empty()
+            }
+            1 => {
+                let (ret, pop) = {
+                    let block = self.blocks.front().expect("unexpected state");
+
+                    let ret = block.shift(n);
+                    self.len -= n;
+
+                    (ret, self.len == 0 && !MutBuf::has_remaining(block))
+                };
+
+                if pop {
+                    let _ = self.blocks.pop_front();
+                }
+
+                ret
+            }
+            _ => {
+                self.shift_multi(n)
+            }
+        }
+    }
+
+    fn shift_multi(&mut self, mut n: usize) -> Bytes {
         let mut ret: Option<Bytes> = None;
 
         while n > 0 {
@@ -96,11 +128,15 @@ impl BlockBuf {
             let (segment, pop) = {
                 let block = self.blocks.front().expect("unexpected state");
 
-                let segment_n = cmp::min(n, block.len());
+
+                let block_len = block.len();
+                let segment_n = cmp::min(n, block_len);
                 n -= segment_n;
                 self.len -= segment_n;
 
-                (block.shift(segment_n), !MutBuf::has_remaining(block))
+                let pop = block_len == segment_n && !MutBuf::has_remaining(block);
+
+                (block.shift(segment_n), pop)
             };
 
             if pop {
@@ -108,13 +144,15 @@ impl BlockBuf {
             }
 
             ret = Some(match ret.take() {
-                Some(curr) => curr.concat(&segment),
+                Some(curr) => {
+                    curr.concat(&segment)
+                }
                 None => segment,
             });
 
         }
 
-        ret.unwrap_or(Bytes::empty())
+        ret.unwrap_or_else(|| Bytes::empty())
     }
 
     /// Drop the first `n` buffered bytes
@@ -146,6 +184,10 @@ impl BlockBuf {
         }
     }
 
+    pub fn is_compact(&mut self) -> bool {
+        self.blocks.len() <= 1
+    }
+
     /// Moves all buffered bytes into a single block.
     ///
     /// # Panics
@@ -208,6 +250,19 @@ impl BlockBuf {
     fn have_buffered_data(&self) -> bool {
         self.len() > 0
     }
+
+    #[inline]
+    fn needs_alloc(&self) -> bool {
+        if let Some(buf) = self.blocks.back() {
+            // `unallocated_blocks` is checked here because if further blocks
+            // cannot be allocated, an empty slice should be returned.
+            if MutBuf::has_remaining(buf) {
+                return false;
+            }
+        }
+
+        true
+    }
 }
 
 impl MutBuf for BlockBuf {
@@ -236,18 +291,9 @@ impl MutBuf for BlockBuf {
         }
     }
 
+    #[inline]
     unsafe fn mut_bytes(&mut self) -> &mut [u8] {
-        let mut need_alloc = true;
+        if self.needs_alloc() {
-
-        if let Some(buf) = self.blocks.back() {
-            // `unallocated_blocks` is checked here because if further blocks
-            // cannot be allocated, an empty slice should be returned.
-            if MutBuf::has_remaining(buf) {
-                need_alloc = false
-            }
-        }
-
-        if need_alloc {
             if self.blocks.len() != self.blocks.capacity() {
                 self.allocate_block()
             }
@@ -261,7 +307,7 @@ impl MutBuf for BlockBuf {
 
 impl Default for BlockBuf {
     fn default() -> BlockBuf {
-        BlockBuf::new(16, NewBlock::Heap(8_192))
+        BlockBuf::new(16, 8_192)
     }
 }
 
@@ -307,7 +353,7 @@ impl NewBlock {
     fn block_size(&self) -> usize {
         match *self {
             NewBlock::Heap(size) => size,
-            NewBlock::Pool(ref pool) => pool.buffer_len(),
+            // NewBlock::Pool(ref pool) => pool.buffer_len(),
         }
     }
 
@@ -315,7 +361,7 @@ impl NewBlock {
     fn new_block(&self) -> Option<AppendBuf> {
         match *self {
             NewBlock::Heap(size) => Some(AppendBuf::with_capacity(size as u32)),
-            NewBlock::Pool(ref pool) => pool.new_append_buf(),
+            // NewBlock::Pool(ref pool) => pool.new_append_buf(),
         }
     }
 }

src/buf/byte.rs

@@ -32,16 +32,6 @@ impl ByteBuf {
         MutByteBuf { buf: ByteBuf::new(capacity as u32) }
     }
 
-    pub fn none() -> ByteBuf {
-        ByteBuf {
-            mem: alloc::MemRef::none(),
-            cap: 0,
-            pos: 0,
-            lim: 0,
-            mark: None,
-        }
-    }
-
     pub unsafe fn from_mem_ref(mem: alloc::MemRef, cap: u32, pos: u32, lim: u32) -> ByteBuf {
         debug_assert!(pos <= lim && lim <= cap, "invalid arguments; cap={}; pos={}; lim={}", cap, pos, lim);
 
@@ -55,28 +45,20 @@ impl ByteBuf {
     }
 
     fn new(mut capacity: u32) -> ByteBuf {
-        // Handle 0 capacity case
-        if capacity == 0 {
-            return ByteBuf::none();
-        }
-
         // Round the capacity to the closest power of 2
         capacity = capacity.next_power_of_two();
 
-        // Allocate the memory
+        unsafe {
-        let mem = alloc::heap(capacity as usize);
+            // Allocate the memory
+            let mem = alloc::heap(capacity as usize);
 
-        // If the allocation failed, return a blank buf
+            ByteBuf {
-        if mem.is_none() {
+                mem: mem,
-            return ByteBuf::none();
+                cap: capacity,
-        }
+                pos: 0,
-
+                lim: capacity,
-        ByteBuf {
+                mark: None,
-            mem: mem,
+            }
-            cap: capacity,
-            pos: 0,
-            lim: capacity,
-            mark: None,
         }
     }
 

src/buf/ring.rs

@@ -23,29 +23,20 @@ pub struct RingBuf {
 impl RingBuf {
     /// Allocates a new `RingBuf` with the specified capacity.
     pub fn new(mut capacity: usize) -> RingBuf {
-        // Handle the 0 length buffer case
+        // Round to the next power of 2 for better alignment
-        if capacity == 0 {
+        capacity = capacity.next_power_of_two();
-            return RingBuf {
+
-                ptr: alloc::MemRef::none(),
+        unsafe {
-                cap: 0,
+            let mem = alloc::heap(capacity as usize);
+
+            RingBuf {
+                ptr: mem,
+                cap: capacity,
                 pos: 0,
                 len: 0,
                 mark: Mark::NoMark,
             }
         }
-
-        // Round to the next power of 2 for better alignment
-        capacity = capacity.next_power_of_two();
-
-        let mem = alloc::heap(capacity as usize);
-
-        RingBuf {
-            ptr: mem,
-            cap: capacity,
-            pos: 0,
-            len: 0,
-            mark: Mark::NoMark,
-        }
     }
 
     /// Returns `true` if the buf cannot accept any further writes.

src/bytes/mod.rs

@@ -9,6 +9,7 @@ use self::small::Small;
 use self::rope::{Rope, RopeBuf};
 use std::{cmp, fmt, ops};
 use std::io::Cursor;
+use std::sync::Arc;
 
 #[derive(Clone)]
 pub struct Bytes {
@@ -19,7 +20,7 @@ pub struct Bytes {
 enum Kind {
     Seq(Seq),
     Small(Small),
-    Rope(Rope),
+    Rope(Arc<Rope>),
 }
 
 pub struct BytesBuf<'a> {
@@ -41,8 +42,9 @@ impl Bytes {
     ///
     /// This function is unsafe as there are no guarantees that the given
     /// arguments are valid.
+    #[inline]
     pub unsafe fn from_mem_ref(mem: alloc::MemRef, pos: u32, len: u32) -> Bytes {
-        Small::from_slice(&mem.bytes()[pos as usize .. pos as usize + len as usize])
+        Small::from_slice(&mem.bytes_slice(pos as usize, pos as usize + len as usize))
             .map(|b| Bytes { kind: Kind::Small(b) })
             .unwrap_or_else(|| {
                 let seq = Seq::from_mem_ref(mem, pos, len);
@@ -110,7 +112,7 @@ impl Bytes {
         }
     }
 
-    fn into_rope(self) -> Result<Rope, Bytes> {
+    fn into_rope(self) -> Result<Arc<Rope>, Bytes> {
         match self.kind {
             Kind::Rope(r) => Ok(r),
             _ => Err(self),

src/bytes/rope.rs

@@ -151,7 +151,7 @@ impl Rope {
                     // the same as the given left tree.
                     let new_right = concat_bytes(&left.right, &right, len);
 
-                    return Rope::new(left.left, new_right).into_bytes();
+                    return Rope::new(left.left.clone(), new_right).into_bytes();
                 }
 
                 if left.left.depth() > left.right.depth() && left.depth > right.depth() {
@@ -161,12 +161,12 @@ impl Rope {
                     // the the node on the RHS.  This is yet another optimization
                     // for building the string by repeatedly concatenating on the
                     // right.
-                    let new_right = Rope::new(left.right, right);
+                    let new_right = Rope::new(left.right.clone(), right);
 
-                    return Rope::new(left.left, new_right).into_bytes();
+                    return Rope::new(left.left.clone(), new_right).into_bytes();
                 }
 
-                left.into_bytes()
+                Bytes { kind: super::Kind::Rope(left) }
             }
             Err(left) => left,
         };
@@ -234,7 +234,7 @@ impl Rope {
 
     fn into_bytes(self) -> Bytes {
         use super::Kind;
-        Bytes { kind: Kind::Rope(self) }
+        Bytes { kind: Kind::Rope(Arc::new(self)) }
     }
 }
 
@@ -302,7 +302,7 @@ impl From<Bytes> for Node {
         match src.kind {
             Kind::Seq(b) => Node::Seq(b),
             Kind::Small(b) => Node::Small(b),
-            Kind::Rope(b) => Node::Rope(Arc::new(b)),
+            Kind::Rope(b) => Node::Rope(b),
         }
     }
 }
@@ -619,15 +619,15 @@ impl Partial {
     }
 
     fn unwrap_rope(self) -> Rope {
-        match self {
+        let arc = match self {
             Partial::Bytes(v) => v.into_rope().ok().expect("unexpected state calling `Partial::unwrap_rope()`"),
-            Partial::Node(Node::Rope(v)) => {
+            Partial::Node(Node::Rope(v)) => v,
-                match Arc::try_unwrap(v) {
-                    Ok(v) => v,
-                    Err(v) => (*v).clone(),
-                }
-            }
             _ => panic!("unexpected state calling `Partial::unwrap_rope()`"),
+        };
+
+        match Arc::try_unwrap(arc) {
+            Ok(v) => v,
+            Err(v) => (*v).clone(),
         }
     }
 }

src/lib.rs

@@ -13,7 +13,7 @@ pub mod alloc;
 
 pub use buf::{Buf, MutBuf, Source, Sink, ReadExt, WriteExt, Fmt};
 pub use buf::append::AppendBuf;
-pub use buf::block::{BlockBuf, NewBlock, BlockBufCursor};
+pub use buf::block::{BlockBuf, BlockBufCursor};
 pub use buf::byte::{ByteBuf, MutByteBuf};
 pub use buf::ring::RingBuf;
 pub use buf::take::Take;

test/test.rs

@@ -18,7 +18,7 @@ mod test_seq;
 mod test_small;
 
 // == Pool
-mod test_pool;
+// mod test_pool;
 
 fn gen_bytes(n: usize) -> Vec<u8> {
     (0..n).map(|_| random()).collect()

test/test_append.rs

@@ -1,5 +1,4 @@
 use bytes::{Buf, MutBuf, AppendBuf};
-use bytes::alloc::Pool;
 
 #[test]
 pub fn test_initial_buf_empty() {
@@ -29,8 +28,10 @@ pub fn test_initial_buf_empty() {
     }
 }
 
+/*
 #[test]
 pub fn test_append_buf_from_pool() {
+    use bytes::alloc::Pool;
     let pool = Pool::with_capacity(2, 256);
 
     // Run in a loop a bunch in hope that if there is a memory issue, it will
@@ -58,3 +59,4 @@ pub fn test_append_buf_from_pool() {
         assert_eq!(dst, b"hello world");
     }
 }
+*/