Combine Rc and Vec allocation into one allocation

Author: dtolnay

src/rcvec.rs

@@ -1,120 +1,238 @@
+use std::alloc::{self, Layout};
+use std::cell::Cell;
+use std::cmp;
 use std::mem;
-use std::rc::Rc;
+use std::process;
+use std::ptr;
 use std::slice;
-use std::vec;
+
+struct Header {
+    cap: usize,
+    len: usize,
+    refcount: Cell<usize>,
+}
 
 pub(crate) struct RcVec<T> {
-    inner: Rc<Vec<T>>,
+    first: *mut T,
 }
 
 pub(crate) struct RcVecBuilder<T> {
-    inner: Vec<T>,
+    first: *mut T,
 }
 
 pub(crate) struct RcVecMut<'a, T> {
-    inner: &'a mut Vec<T>,
+    first: &'a mut *mut T,
 }
 
 #[derive(Clone)]
 pub(crate) struct RcVecIntoIter<T> {
-    inner: vec::IntoIter<T>,
+    first: *mut T,
+    cur: *mut T,
+}
+
+unsafe fn header<'a, T>(first: *mut T) -> &'a Header {
+    &*(first as *const Header).offset(-1)
+}
+
+unsafe fn header_mut<'a, T>(first: *mut T) -> &'a mut Header {
+    &mut *(first as *mut Header).offset(-1)
+}
+
+unsafe fn allocation_layout<T>(first: *mut T) -> (*mut u8, Layout) {
+    // Alignment of the allocation.
+    let align = cmp::max(mem::align_of::<T>(), mem::align_of::<Header>());
+
+    // Round up Header's size to a multiple of allocation's alignment.
+    let mut header_size = mem::size_of::<Header>();
+    header_size = (header_size + align - 1) / align * align;
+
+    let allocation = (first as *mut u8).sub(header_size);
+    let total_size = header_size + header(first).cap * mem::size_of::<T>();
+    let layout = Layout::from_size_align_unchecked(total_size, align);
+    (allocation, layout)
 }
 
 impl<T> RcVec<T> {
     pub fn is_empty(&self) -> bool {
-        self.inner.is_empty()
+        self.len() == 0
     }
 
     pub fn len(&self) -> usize {
-        self.inner.len()
+        self.header().len
     }
 
     pub fn iter(&self) -> slice::Iter<T> {
-        self.inner.iter()
+        let len = self.len();
+        unsafe { slice::from_raw_parts(self.first, len) }.iter()
     }
 
     pub fn make_mut(&mut self) -> RcVecMut<T>
     where
         T: Clone,
     {
+        let header = self.header();
+        let refcount = header.refcount.get();
+        if refcount > 1 {
+            let mut builder = RcVecBuilder::new();
+            builder.extend(self.iter().cloned());
+            // NOTE: assumes that none of the clones on the previous line
+            // resulted in self's refcount changing. The public API of RcVec
+            // does not expose any way to form cyclic structures.
+            header.refcount.set(refcount - 1);
+            self.first = builder.first;
+            mem::forget(builder);
+        }
         RcVecMut {
-            inner: Rc::make_mut(&mut self.inner),
+            first: &mut self.first,
         }
     }
 
     pub fn get_mut(&mut self) -> Option<RcVecMut<T>> {
-        let inner = Rc::get_mut(&mut self.inner)?;
-        Some(RcVecMut { inner })
+        if self.header().refcount.get() > 1 {
+            None
+        } else {
+            Some(RcVecMut {
+                first: &mut self.first,
+            })
+        }
     }
 
-    pub fn make_owned(mut self) -> RcVecBuilder<T>
+    pub fn make_owned(self) -> RcVecBuilder<T>
     where
         T: Clone,
     {
-        let vec = if let Some(owned) = Rc::get_mut(&mut self.inner) {
-            mem::replace(owned, Vec::new())
+        let header = self.header();
+        let refcount = header.refcount.get();
+        let first;
+        if refcount > 1 {
+            let mut builder = RcVecBuilder::new();
+            builder.extend(self.iter().cloned());
+            header.refcount.set(refcount - 1);
+            first = builder.first;
+            mem::forget(builder);
         } else {
-            Vec::clone(&self.inner)
-        };
-        RcVecBuilder { inner: vec }
+            first = self.first;
+        }
+        mem::forget(self);
+        RcVecBuilder { first }
+    }
+
+    fn header(&self) -> &Header {
+        unsafe { header(self.first) }
     }
 }
 
 impl<T> RcVecBuilder<T> {
     pub fn new() -> Self {
-        RcVecBuilder { inner: Vec::new() }
+        Self::with_capacity(8)
     }
 
     pub fn with_capacity(cap: usize) -> Self {
-        RcVecBuilder {
-            inner: Vec::with_capacity(cap),
+        assert!(mem::size_of::<T>() > 0);
+
+        // Alignment of the allocation.
+        let align = cmp::max(mem::align_of::<T>(), mem::align_of::<Header>());
+
+        // Round up Header's size to a multiple of allocation's alignment.
+        let mut header_size = mem::size_of::<Header>();
+        header_size = (header_size + align - 1) / align * align;
+
+        let total_size = mem::size_of::<T>()
+            .saturating_mul(cap)
+            .saturating_add(header_size);
+        let layout = match Layout::from_size_align(total_size, align) {
+            Ok(layout) => layout,
+            Err(_) => process::abort(),
+        };
+
+        unsafe {
+            let ptr = alloc::alloc(layout);
+            let first = ptr.add(header_size) as *mut T;
+            (first as *mut Header).offset(-1).write(Header {
+                cap,
+                len: 0,
+                refcount: Cell::new(1),
+            });
+            RcVecBuilder { first }
         }
     }
 
     pub fn push(&mut self, element: T) {
-        self.inner.push(element);
+        self.as_mut().push(element);
     }
 
     pub fn extend(&mut self, iter: impl IntoIterator<Item = T>) {
-        self.inner.extend(iter);
+        self.as_mut().extend(iter);
     }
 
     pub fn as_mut(&mut self) -> RcVecMut<T> {
         RcVecMut {
-            inner: &mut self.inner,
+            first: &mut self.first,
         }
     }
 
     pub fn build(self) -> RcVec<T> {
-        RcVec {
-            inner: Rc::new(self.inner),
-        }
+        let first = self.first;
+        mem::forget(self);
+        RcVec { first }
     }
 }
 
 impl<'a, T> RcVecMut<'a, T> {
     pub fn push(&mut self, element: T) {
-        self.inner.push(element);
+        self.reserve(1);
+        let header = unsafe { header_mut(*self.first) };
+        unsafe { ptr::write(self.first.add(header.len), element) };
+        header.len += 1;
     }
 
-    pub fn extend(&mut self, iter: impl IntoIterator<Item = T>) {
-        self.inner.extend(iter);
+    pub fn extend(&mut self, elements: impl IntoIterator<Item = T>) {
+        let iter = elements.into_iter();
+        self.reserve(iter.size_hint().0);
+        iter.for_each(|element| self.push(element));
     }
 
     pub fn pop(&mut self) -> Option<T> {
-        self.inner.pop()
+        unsafe {
+            let header = header_mut(*self.first);
+            header.len = header.len.checked_sub(1)?;
+            Some(ptr::read(self.first.add(header.len)))
+        }
     }
 
     pub fn as_mut(&mut self) -> RcVecMut<T> {
-        RcVecMut { inner: self.inner }
+        RcVecMut { first: self.first }
+    }
+
+    fn reserve(&mut self, additional: usize) {
+        let header = unsafe { header(*self.first) };
+        let required_capacity = header.len.saturating_add(additional);
+        if header.cap >= required_capacity {
+            return;
+        }
+
+        let min_new_capacity = cmp::max(8, header.cap.saturating_mul(2));
+        let new_capacity = cmp::max(required_capacity, min_new_capacity);
+
+        // Move ownership of the old allocation.
+        let small = RcVecBuilder { first: *self.first };
+
+        // Move elements into new allocation.
+        let mut big = RcVecBuilder::with_capacity(new_capacity);
+        big.extend(small.into_iter());
+
+        // Self becomes big.
+        *self.first = big.first;
+        mem::forget(big);
     }
 }
 
 impl<T> Clone for RcVec<T> {
     fn clone(&self) -> Self {
-        RcVec {
-            inner: Rc::clone(&self.inner),
-        }
+        let header = self.header();
+        let count = header.refcount.get().checked_add(1).unwrap();
+        header.refcount.set(count);
+        RcVec { first: self.first }
     }
 }
 
@@ -123,20 +241,68 @@ impl<T> IntoIterator for RcVecBuilder<T> {
     type IntoIter = RcVecIntoIter<T>;
 
     fn into_iter(self) -> Self::IntoIter {
-        RcVecIntoIter {
-            inner: self.inner.into_iter(),
-        }
+        let first = self.first;
+        mem::forget(self);
+        RcVecIntoIter { first, cur: first }
     }
 }
 
 impl<T> Iterator for RcVecIntoIter<T> {
     type Item = T;
 
     fn next(&mut self) -> Option<Self::Item> {
-        self.inner.next()
+        unsafe {
+            let len = header(self.first).len;
+            if self.cur < self.first.add(len) {
+                let element = ptr::read(self.cur);
+                self.cur = self.cur.add(1);
+                Some(element)
+            } else {
+                None
+            }
+        }
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        self.inner.size_hint()
+        let end = unsafe { self.first.add(header(self.first).len) };
+        let remaining = (end as usize - self.cur as usize) / mem::size_of::<T>();
+        (remaining, Some(remaining))
+    }
+}
+
+impl<T> Drop for RcVec<T> {
+    fn drop(&mut self) {
+        let header = self.header();
+        let refcount = header.refcount.get();
+        if refcount > 1 {
+            header.refcount.set(refcount - 1);
+        } else {
+            drop(RcVecBuilder { first: self.first });
+        }
+    }
+}
+
+impl<T> Drop for RcVecBuilder<T> {
+    fn drop(&mut self) {
+        unsafe {
+            let len = header(self.first).len;
+            let slice = slice::from_raw_parts_mut(self.first, len);
+            ptr::drop_in_place(slice);
+            let (ptr, layout) = allocation_layout(self.first);
+            alloc::dealloc(ptr, layout);
+        }
+    }
+}
+
+impl<T> Drop for RcVecIntoIter<T> {
+    fn drop(&mut self) {
+        unsafe {
+            let len = header(self.first).len;
+            let offset = (self.cur as usize - self.first as usize) / mem::size_of::<T>();
+            let slice = slice::from_raw_parts_mut(self.cur, len - offset);
+            ptr::drop_in_place(slice);
+            let (ptr, layout) = allocation_layout(self.first);
+            alloc::dealloc(ptr, layout);
+        }
     }
 }