realloc 0.1.1

//! A growable collection analagous to `std::vec::Vec`.

use core::mem::MaybeUninit;
use core::ops::{Deref, DerefMut};

use crate::{Alloc, Allocator, AllocatorExt, Error, Fallible, Layout, Strategy};

mod iter;
pub use iter::IntoIter;

/// A contiguous growable array type, written as `Vec<T>`, short for ‘vector’.
///
/// Currently does not have feature parity with `std::vec::Vec`.
pub struct Vec<'alloc, T, S: Strategy = Fallible> {
    state: VecState<'alloc, T>,
    len: usize,

    _strategy: S,
}

enum VecState<'alloc, T> {
    Empty(&'alloc dyn Allocator),
    Alloc { alloc: Alloc<'alloc, T>, cap: usize },
}

impl<'alloc, T> VecState<'alloc, T> {
    fn take(&mut self) -> Option<(Alloc<'alloc, T>, usize)> {
        let allocator = self.allocator();
        match core::mem::replace(self, Self::Empty(allocator)) {
            Self::Empty(_) => None,
            Self::Alloc { alloc, cap } => Some((alloc, cap)),
        }
    }

    fn is_allocated(&self) -> bool {
        matches!(self, Self::Alloc { .. })
    }

    fn force_alloc(&mut self) -> Result<(&mut Alloc<'alloc, T>, usize), Error> {
        Ok(match self {
            Self::Empty(allocator) => {
                let alloc = allocator.alloc::<[T; 4]>()?.cast();
                *self = Self::Alloc { alloc, cap: 4 };
                let Self::Alloc { alloc, .. } = self else {
                    unreachable!()
                };
                (alloc, 4)
            }
            Self::Alloc { alloc, cap } => (alloc, *cap),
        })
    }

    fn allocator(&self) -> &'alloc dyn Allocator {
        match self {
            &Self::Empty(allocator) => allocator,
            Self::Alloc { alloc, .. } => alloc.allocator(),
        }
    }

    fn buf(&self) -> &[MaybeUninit<T>] {
        match self {
            Self::Empty(_) => &[],
            Self::Alloc { alloc, cap } => {
                let slice = core::ptr::slice_from_raw_parts(alloc.as_ptr().cast(), *cap);

                // SAFETY: `alloc` is `cap` elements long, and `MaybeUninit<T>`
                // is always valid
                unsafe { &*slice }
            }
        }
    }

    fn buf_mut(&mut self) -> &mut [MaybeUninit<T>] {
        match self {
            Self::Empty(_) => &mut [],
            Self::Alloc { alloc, cap } => {
                let slice = core::ptr::slice_from_raw_parts_mut(alloc.as_ptr().cast(), *cap);

                // SAFETY: `alloc` is `cap` elements long, and `MaybeUninit<T>`
                // is always valid
                unsafe { &mut *slice }
            }
        }
    }
}

#[cfg(feature = "global")]
impl<T, S: Strategy> Vec<'static, T, S> {
    /// Create a new vector in the global allocator.
    ///
    /// Fails if there is no global allocator.
    ///
    /// Only available on feature `global`.
    pub fn new() -> S::Result<Self, Error> {
        let allocator = S::create(crate::global::get().ok_or(Error::NoGlobalAllocator));
        S::map(allocator, |allocator| Self::new_in(allocator))
    }
}

impl<'alloc, T, S: Strategy> Vec<'alloc, T, S> {
    /// Create a new vector in the given allocator.
    pub const fn new_in(allocator: &'alloc dyn Allocator) -> Self {
        Self {
            state: VecState::Empty(allocator),
            len: 0,

            _strategy: S::UNIT,
        }
    }

    /// Returns the number of elements currently contained in this vector.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Returns `true` if the vector contains no elements.
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Returns the number of elements this vector can hold without
    /// reallocating.
    pub fn capacity(&self) -> usize {
        self.state.buf().len()
    }

    /// Adds a new element to the end of this vector.
    pub fn push(&mut self, element: T) -> S::Result<(), Error> {
        if let Some(next) = self.next_open() {
            next.write(element);
            self.len += 1;

            S::ok(())
        } else {
            S::and_then(self.grow(), |()| {
                self.next_open().unwrap().write(element);
                self.len += 1;

                S::ok(())
            })
        }
    }

    /// Removes and returns the last element of this vector, if any.
    pub fn pop(&mut self) -> Option<T> {
        if self.len > 0 {
            self.len -= 1;
            let element = &self.state.buf_mut()[self.len];
            // SAFETY: the element just came from the initialized section of the
            // buf
            Some(unsafe { element.assume_init_read() })
        } else {
            None
        }
    }

    /// Drops all elements in this vector.
    ///
    /// This does not deallocate the underlying memory.
    pub fn clear(&mut self) {
        if self.len == 0 {
            return;
        }

        // We set `len = 0` early, in case a `Drop` panics and causes the
        // mid-clear state to be observed
        let len = self.len;
        self.len = 0;

        for element in &mut self.state.buf_mut()[..len] {
            // SAFETY:
            // - `element` is (presumably; we are not responsible for the user)
            // valid for dropping, and is never touched again.
            unsafe { element.assume_init_drop() };
        }

        self.len = 0;
    }

    /// Creates a clone of this vector in the given allocator.
    ///
    /// If you want to clone it in the same allocator this vector was created
    /// in, see [`clone_in_same`](Self::clone_in_same).
    pub fn clone_in(&self, allocator: &'alloc dyn Allocator) -> S::Result<Self, Error> {
        if !self.state.is_allocated() {
            return S::ok(Vec {
                state: VecState::Empty(allocator),
                len: 0,

                _strategy: S::UNIT,
            });
        }

        let alloc = S::create(allocator.alloc_bytes(Layout::list::<T>(self.len)));
        S::map(alloc, |alloc| {
            let alloc = alloc.cast::<T>();
            // SAFETY: `self.buf_mut()` is at least `self.len` initialized elements long, and
            // `alloc` is exactly `self.len` elements long
            unsafe {
                self.state
                    .buf()
                    .as_ptr()
                    .copy_to_nonoverlapping(alloc.as_ptr().cast(), self.len)
            };

            Vec {
                state: VecState::Alloc {
                    alloc,
                    cap: self.len,
                },
                len: self.len,

                _strategy: S::UNIT,
            }
        })
    }

    /// Creates a clone of this vector in the same allocator it was created in.
    ///
    /// If you want to clone it in a different allocator, see
    /// [`clone_in`](Self::clone_in).
    pub fn clone_in_same(&self) -> S::Result<Self, Error> {
        self.clone_in(self.allocator())
    }

    /// Change the [failure strategy](Strategy) employed on fallible methods.
    pub fn change_strategy<NS: Strategy>(mut self) -> Vec<'alloc, T, NS> {
        let allocator = self.allocator();
        let state = core::mem::replace(&mut self.state, VecState::Empty(allocator));
        Vec {
            state,
            len: self.len,

            _strategy: NS::UNIT,
        }
    }

    /// Returns a reference to the underlying allocator.
    pub fn allocator(&self) -> &'alloc dyn Allocator {
        self.state.allocator()
    }

    fn next_open(&mut self) -> Option<&mut MaybeUninit<T>> {
        self.state.buf_mut().get_mut(self.len)
    }

    fn grow(&mut self) -> S::Result<(), Error> {
        let mut new_cap = 0;
        let buf = S::and_then(S::create(self.state.force_alloc().map(drop)), |()| {
            let (buf, cap) = self.state.take().unwrap();
            
            // don't grow if it was just allocated
            if cap == 4 {
                new_cap = 4;
                S::ok(buf.cast())
            } else {
                new_cap = cap * 2;
                let layout = Layout::list::<T>(new_cap);
                S::create(buf.grow(layout))
            }
        });

        S::map(buf, |buf| {
            self.state = VecState::Alloc {
                alloc: buf.cast(),
                cap: new_cap,
            };
        })
    }
}

impl<T, S: Strategy> Deref for Vec<'_, T, S> {
    type Target = [T];

    fn deref(&self) -> &[T] {
        let slice = &self.state.buf()[..self.len] as *const [_] as *const [T];
        // SAFETY: `self.state` has `self.len` initialized elements
        unsafe { &*slice }
    }
}

impl<T, S: Strategy> DerefMut for Vec<'_, T, S> {
    fn deref_mut(&mut self) -> &mut [T] {
        let slice = &mut self.state.buf_mut()[..self.len] as *mut [_] as *mut [T];
        // SAFETY: `self.state` has `self.len` initialized elements
        unsafe { &mut *slice }
    }
}

impl<T, S: Strategy> Clone for Vec<'_, T, S> {
    // TODO: could this be done without copy-pasting the impl from `clone_in`?
    fn clone(&self) -> Self {
        let allocator = self.allocator();

        if !self.state.is_allocated() {
            return Vec {
                state: VecState::Empty(allocator),
                len: 0,

                _strategy: S::UNIT,
            };
        }

        let alloc = allocator
            .alloc_bytes(Layout::list::<T>(self.len))
            .expect(crate::FAILURE_MESSAGE)
            .cast::<T>();
        // SAFETY: `self.buf_mut()` is at least `self.len` initialized elements long, and
        // `alloc` is exactly `self.len` elements long
        unsafe {
            self.state
                .buf()
                .as_ptr()
                .copy_to_nonoverlapping(alloc.as_ptr().cast(), self.len)
        };

        Vec {
            state: VecState::Alloc {
                alloc,
                cap: self.len,
            },
            len: self.len,

            _strategy: S::UNIT,
        }
    }
}

impl<'alloc, T, S: Strategy> IntoIterator for Vec<'alloc, T, S> {
    type Item = T;
    type IntoIter = iter::IntoIter<'alloc, T>;

    fn into_iter(self) -> Self::IntoIter {
        IntoIter::new(self)
    }
}

impl<T, S: Strategy> Drop for Vec<'_, T, S> {
    fn drop(&mut self) {
        self.clear();
        if let Some((alloc, _)) = self.state.take() {
            alloc.dealloc();
        }
    }
}