rc_vec.rs

use std::{ops, cmp, iter, mem, ptr, slice};
use std::marker::PhantomData;
use std::mem::MaybeUninit;
use std::rc::Rc;

/// A copy-on-write `Vec`.
pub struct RcVec<T> {
    buf: Rc<[MaybeUninit<T>]>,
    len: usize,

    _marker: PhantomData<[T]>,
}

impl<T> Default for RcVec<T> {
    fn default() -> Self {
        RcVec { buf: Rc::new([]), len: 0, _marker: PhantomData }
    }
}

impl<T> Clone for RcVec<T> {
    fn clone(&self) -> RcVec<T> {
        RcVec { buf: Rc::clone(&self.buf), len: self.len, _marker: PhantomData }
    }
}

impl<T> Drop for RcVec<T> {
    fn drop(&mut self) {
        // If this is the last reference to the vec, drop its elements.
        if Rc::strong_count(&self.buf) == 1 {
            unsafe {
                let buf = slice::from_raw_parts_mut(self.as_mut_ptr(), self.len);
                ptr::drop_in_place(buf);
            }
        }
    }
}

impl<T> ops::Deref for RcVec<T> {
    type Target = [T];

    fn deref(&self) -> &[T] {
        unsafe { slice::from_raw_parts(self.as_ptr(), self.len) }
    }
}

impl<T: Clone> ops::DerefMut for RcVec<T> {
    fn deref_mut(&mut self) -> &mut [T] {
        self.make_mut();

        // Safety: The strong count of `self.buf` is now 1.
        unsafe { slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
    }
}

impl<T> RcVec<T> {
    pub fn as_ptr(&self) -> *const T {
        // Safety: This immediately "forgets" the duplicate `Rc`.
        let ptr = unsafe { Rc::into_raw(ptr::read(&self.buf)) };
        ptr as *const T
    }

    fn as_mut_ptr(&mut self) -> *mut T {
        // Safety: This immediately "forgets" the duplicate `Rc`.
        let ptr = unsafe { Rc::into_raw(ptr::read(&self.buf)) };
        ptr as *mut T
    }

    pub fn len(&self) -> usize { self.len }
}

impl<T: Clone> RcVec<T> {
    pub fn reserve(&mut self, extra: usize) {
        if self.buf.len().wrapping_sub(self.len) >= extra {
            return;
        }

        let buf_len = self.len.checked_add(extra).expect("capacity overflow");
        let buf_len = cmp::max(buf_len, 2 * self.buf.len());
        let size = buf_len.checked_mul(mem::size_of::<T>()).expect("capacity overflow");
        if mem::size_of::<usize>() < 8 && size > isize::MAX as usize {
            panic!("capacity overflow");
        }

        // If this is the only reference to the vec, move its elements.
        // Otherwise, clone them.
        if Rc::strong_count(&self.buf) == 1 {
            // Safety: The moved from values are immediately freed without being dropped.
            let iter = self.buf.iter().take(self.len).map(|value| unsafe { value.read() });
            self.buf = Self::fill_with_capacity(iter, buf_len);
        } else {
            self.buf = Self::fill_with_capacity(self.iter().cloned(), buf_len);
        }
    }

    pub fn make_mut(&mut self) -> &mut [T] {
        if Rc::strong_count(&self.buf) != 1 {
            self.buf = Self::fill_with_capacity(self.iter().cloned(), self.buf.len());
        }

        // Safety: The strong count of `self.buf` is now 1, and we hold `&mut self`.
        unsafe { slice::from_raw_parts_mut(self.as_mut_ptr(), self.len) }
    }

    fn fill_with_capacity(iter: impl Iterator<Item = T>, cap: usize) -> Rc<[MaybeUninit<T>]> {
        iter
            .map(MaybeUninit::new)
            .chain(iter::repeat_with(MaybeUninit::uninit))
            .take(cap)
            .collect()
    }

    pub fn push(&mut self, value: T) {
        if self.len == self.buf.len() {
            self.reserve(1);
        }

        // Safety: If there was excess capacity, then no other references have formed any shared
        // `&[T]`s to that memory. Otherwise, the strong count of `self.buf` is now 1 and there are
        // no other references.
        unsafe {
            let end = self.as_mut_ptr().add(self.len);
            ptr::write(end, value);
            self.len += 1;
        }
    }

    pub fn remove(&mut self, index: usize) -> T {
        let len = self.len();
        assert!(index < len);

        let ret;
        self.make_mut();

        // Safety: The strong count of `self.buf` is now 1, and `index` is in bounds.
        unsafe {
            let ptr = self.as_mut_ptr().add(index);
            ret = ptr::read(ptr);
            ptr::copy(ptr.offset(1), ptr, len - index - 1);
        }

        self.len -= 1;
        ret
    }
}