tl/inline/

vec.rs

use core::fmt::{Debug, Formatter};
use core::mem::MaybeUninit;
use core::ops::Index;
use core::ptr;

use crate::ParseError;

/// A wrapper around a `Vec<T>` that lives on the stack if it is small enough.
#[derive(Debug, Clone)]
pub struct InlineVec<T, const N: usize>(InlineVecInner<T, N>);

impl<T, const N: usize> InlineVec<T, N> {
    /// Creates a new InlineVec
    #[inline]
    pub(crate) fn new() -> Self {
        Self(InlineVecInner::new())
    }

    /// Returns the number of elements in the vector
    #[inline]
    pub fn len(&self) -> usize {
        self.0.len()
    }

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

    /// Checks whether this vector is allocated on the heap
    #[inline]
    pub fn is_heap_allocated(&self) -> bool {
        self.0.is_heap_allocated()
    }

    /// If `self` is inlined, this returns the underlying raw parts that make up this `InlineVec`.
    ///
    /// Only the first `.1` elements are initialized.
    ///
    /// ```compile_fail
    /// fn mutate<T, const N: usize>(values: &mut tl::inline::vec::InlineVec<T, N>) {
    ///     let _ = values.inline_parts_mut(); // an unsafe call is required
    /// }
    /// ```
    ///
    /// # Safety
    /// Before the borrow ends, the initialized prefix must contain valid, uniquely
    /// owned values of the element type. Replacing a slot with uninitialized data,
    /// duplicating an owner or dropping it without replacing it violates this contract.
    /// Prefer `get_mut` or typed mutable access instead.
    #[inline]
    pub unsafe fn inline_parts_mut(&mut self) -> Option<(&mut [MaybeUninit<T>; N], usize)> {
        self.0.inline_parts_mut()
    }

    /// Copies `self` into a new `Vec<T>`
    #[inline]
    #[cfg(feature = "std")]
    pub fn to_vec(&self) -> Vec<T>
    where
        T: Clone,
    {
        self.0.to_vec()
    }

    /// Inserts a new element into the vector
    #[inline]
    pub fn push(&mut self, value: T) -> Result<(), ParseError> {
        self.0.push(value)
    }

    /// Returns a reference to the value at the given index
    #[inline]
    pub fn get(&self, index: usize) -> Option<&T> {
        self.0.get(index)
    }

    /// Returns a mutable reference to the value at the given index
    #[inline]
    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
        self.0.get_mut(index)
    }

    /// Returns the last element in the vector.
    #[inline]
    pub fn last(&self) -> Option<&T> {
        self.as_slice().last()
    }

    /// Removes the last element and returns it, or `None` if it is empty.
    #[inline]
    pub fn pop(&mut self) -> Option<T> {
        if self.is_empty() {
            None
        } else {
            Some(self.remove(self.len() - 1))
        }
    }

    /// Removes an element at a given index
    ///
    /// # Panics
    /// Just like `Vec::remove`, this method will panic if the index is out of bounds.
    #[inline]
    pub fn remove(&mut self, index: usize) -> T {
        self.0.remove(index)
    }

    /// Returns an iterator over the elements of this vector
    #[inline]
    pub fn iter(&self) -> InlineVecIter<'_, T, N> {
        self.0.iter()
    }

    /// Returns a slice to the contents of this vector
    #[inline]
    pub fn as_slice(&self) -> &[T] {
        self.0.as_slice()
    }

    /// Returns a mutable slice to the contents of this vector.
    #[inline]
    pub fn as_mut_slice(&mut self) -> &mut [T] {
        self.0.as_mut_slice()
    }
}

enum InlineVecInner<T, const N: usize> {
    Inline {
        len: usize,
        data: [MaybeUninit<T>; N],
    },
    #[cfg(feature = "std")]
    Heap(Vec<T>),
}

impl<T, const N: usize> Debug for InlineVecInner<T, N>
where
    T: Debug,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), core::fmt::Error> {
        write!(f, "InlineVec<{} items>", self.len())
    }
}

impl<T, const N: usize> Clone for InlineVecInner<T, N>
where
    T: Clone,
{
    fn clone(&self) -> Self {
        match self {
            #[cfg(feature = "std")]
            Self::Heap(m) => Self::Heap(m.clone()),
            Self::Inline { len, data } => {
                let mut cloned = Self::new();
                let (initialized, target) = match &mut cloned {
                    Self::Inline { len, data } => (len, data),
                    #[cfg(feature = "std")]
                    Self::Heap(_) => unreachable!("new storage is inline"),
                };
                for element in data.iter().take(*len) {
                    // SAFETY: only the initialized prefix is read. The new owner
                    // updates its prefix after each successful clone, so a panic
                    // drops all previously cloned elements exactly once.
                    let value = unsafe { element.assume_init_ref() }.clone();
                    target[*initialized].write(value);
                    *initialized += 1;
                }
                cloned
            }
        }
    }
}

impl<T, const N: usize> InlineVecInner<T, N> {
    #[inline]
    pub(crate) fn new() -> Self {
        Self::Inline {
            len: 0,
            data: super::uninit_array(),
        }
    }

    pub fn as_slice(&self) -> &[T] {
        match self {
            Self::Inline { len, data } => unsafe {
                core::slice::from_raw_parts(data.as_ptr() as *const T, *len)
            },
            #[cfg(feature = "std")]
            Self::Heap(v) => v.as_slice(),
        }
    }

    pub fn as_mut_slice(&mut self) -> &mut [T] {
        match self {
            Self::Inline { len, data } => unsafe {
                core::slice::from_raw_parts_mut(data.as_mut_ptr() as *mut T, *len)
            },
            #[cfg(feature = "std")]
            Self::Heap(v) => v.as_mut_slice(),
        }
    }

    #[inline]
    pub fn inline_parts_mut(&mut self) -> Option<(&mut [MaybeUninit<T>; N], usize)> {
        match self {
            #[cfg(feature = "std")]
            Self::Heap(_) => None,
            Self::Inline { len, data } => Some((data, *len)),
        }
    }

    #[cfg(feature = "std")]
    pub fn to_vec(&self) -> Vec<T>
    where
        T: Clone,
    {
        match &self {
            InlineVecInner::Heap(m) => m.to_vec(),
            InlineVecInner::Inline { len, data } => {
                let mut new_data = Vec::with_capacity(*len);

                let iter = data.iter().take(*len);

                for element in iter {
                    new_data.push(unsafe { T::clone(&*element.as_ptr()) });
                }

                new_data
            }
        }
    }

    #[inline]
    pub fn iter(&self) -> InlineVecIter<'_, T, N> {
        InlineVecIter { idx: 0, vec: self }
    }

    #[inline]
    pub fn len(&self) -> usize {
        match self {
            Self::Inline { len, .. } => *len,
            #[cfg(feature = "std")]
            Self::Heap(vec) => vec.len(),
        }
    }

    pub fn get(&self, idx: usize) -> Option<&T> {
        match self {
            Self::Inline { data, len } => {
                if idx < *len {
                    Some(unsafe { &*data.get_unchecked(idx).as_ptr() })
                } else {
                    None
                }
            }
            #[cfg(feature = "std")]
            Self::Heap(vec) => vec.get(idx),
        }
    }

    pub fn get_mut(&mut self, idx: usize) -> Option<&mut T> {
        match self {
            Self::Inline { data, len } => {
                if idx < *len {
                    Some(unsafe { &mut *data.get_unchecked_mut(idx).as_mut_ptr() })
                } else {
                    None
                }
            }
            #[cfg(feature = "std")]
            Self::Heap(vec) => vec.get_mut(idx),
        }
    }

    pub fn remove(&mut self, idx: usize) -> T {
        match self {
            Self::Inline { data, len } => {
                assert!(idx < *len);

                // at this point we know idx is in bounds
                // carefully replace the value with MaybeUninit::uninit(), so it can be returned
                let element = unsafe {
                    core::mem::replace(data.get_unchecked_mut(idx), MaybeUninit::uninit())
                };

                for i in idx + 1..*len {
                    // TODO(y21): data.swap_unchecked() worth it?
                    data.swap(i, i - 1);
                }

                *len -= 1;

                // we've made sure that idx is in bounds and if idx is in bounds, then `T` must be initialized
                unsafe { element.assume_init() }
            }
            #[cfg(feature = "std")]
            Self::Heap(h) => h.remove(idx),
        }
    }

    pub fn push(&mut self, value: T) -> Result<(), ParseError> {
        let (array, len) = match self {
            Self::Inline { data, len } => (data, len),
            #[cfg(feature = "std")]
            Self::Heap(vec) => {
                vec.push(value);
                return Ok(());
            }
        };

        if *len >= N {
            #[cfg(not(feature = "std"))]
            {
                return Err(ParseError::ChildCapacityExceeded);
            }

            #[cfg(feature = "std")]
            {
                let mut vec = Vec::with_capacity(*len + 1);

                // move old elements to heap
                for element in array.iter_mut().take(*len) {
                    let element = core::mem::replace(element, MaybeUninit::uninit());

                    vec.push(unsafe { element.assume_init() });
                }

                // push the new element
                vec.push(value);
                let new_heap = InlineVecInner::Heap(vec);

                // do not call the destructor!
                unsafe { ptr::write(self, new_heap) };
            }
        } else {
            array[*len].write(value);
            *len += 1;
        }

        Ok(())
    }

    #[inline]
    pub fn is_heap_allocated(&self) -> bool {
        #[cfg(feature = "std")]
        {
            matches!(self, Self::Heap(_))
        }
        #[cfg(not(feature = "std"))]
        {
            false
        }
    }
}

impl<T, const N: usize> Index<usize> for InlineVec<T, N> {
    type Output = T;

    fn index(&self, idx: usize) -> &Self::Output {
        self.0.get(idx).expect("index out of bounds")
    }
}

/// An iterator over the elements stored in an [`InlineVec`]
pub struct InlineVecIter<'a, T, const N: usize> {
    vec: &'a InlineVecInner<T, N>,
    idx: usize,
}

impl<'a, T, const N: usize> Iterator for InlineVecIter<'a, T, N> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        self.idx += 1;
        self.vec.get(self.idx - 1)
    }
}

impl<T, const N: usize> Drop for InlineVecInner<T, N> {
    fn drop(&mut self) {
        if let Some((data, len)) = self.inline_parts_mut() {
            for element in data.iter_mut().take(len) {
                unsafe { ptr::drop_in_place(element.as_mut_ptr()) };
            }
        }
    }
}

#[cfg(all(test, feature = "std"))]
mod tests {
    #![allow(unused_must_use)]

    use super::*;

    #[test]
    fn inlinevec_to_vec_stack() {
        let mut x = InlineVec::<usize, 4>::new();

        for i in 0..4 {
            x.push(i * 2);
        }

        assert!(!x.is_heap_allocated());
        assert_eq!(x.len(), 4);

        let xx = x.to_vec();
        assert_eq!(xx.as_slice(), &[0, 2, 4, 6]);

        x.push(42);
        assert!(x.is_heap_allocated());
        assert_eq!(x.as_slice(), &[0, 2, 4, 6, 42]);
        assert_eq!(x.get(4), Some(&42));

        let xx = x.to_vec();
        assert_eq!(xx.as_slice(), &[0, 2, 4, 6, 42]);
    }

    #[test]
    fn inlinevec_to_vec_heap() {
        let mut x = InlineVec::<String, 4>::new();

        for i in 0..4u8 {
            x.push(i.to_string());
        }

        assert!(!x.is_heap_allocated());
        assert_eq!(x.len(), 4);

        let xx = x.to_vec();
        assert_eq!(xx.as_slice(), &["0", "1", "2", "3"]);

        x.push("1337".into());
        assert!(x.is_heap_allocated());
        assert_eq!(x.as_slice(), &["0", "1", "2", "3", "1337"]);
        assert_eq!(x.get(4).map(|x| &**x), Some("1337"));

        let xx = x.to_vec();
        assert_eq!(xx.as_slice(), &["0", "1", "2", "3", "1337"]);
    }

    #[test]
    fn inlinevec_drop_stack() {
        let mut x = InlineVec::<String, 4>::new();

        for i in 0..3u8 {
            x.push(i.to_string());
        }

        assert_eq!(x.as_slice(), &["0", "1", "2"]);
        assert!(!x.is_heap_allocated());
    }

    #[test]
    fn inlinehashmap_drop_heap() {
        let mut x = InlineVec::<String, 4>::new();

        for i in 0..8u8 {
            x.push(i.to_string());
        }

        assert_eq!(x.as_slice(), &["0", "1", "2", "3", "4", "5", "6", "7"]);
        assert!(x.is_heap_allocated());
    }

    #[test]
    fn inlinevec_iter() {
        let mut x = InlineVecInner::<usize, 2>::new();
        x.push(13);
        x.push(42);
        x.push(17);
        x.push(19);
        let mut iter = x.iter();
        assert_eq!(iter.next(), Some(&13));
        assert_eq!(iter.next(), Some(&42));
        assert_eq!(iter.next(), Some(&17));
        assert_eq!(iter.next(), Some(&19));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn inlinevec_remove() {
        let mut x = InlineVecInner::<usize, 4>::new();
        x.push(789);
        assert_eq!(x.len(), 1);
        assert_eq!(x.get(0), Some(&789));
        assert_eq!(x.remove(0), 789);
        assert_eq!(x.len(), 0);

        {
            let mut xc = x.clone();
            // out of bounds index must panic
            assert!(std::panic::catch_unwind(move || xc.remove(0)).is_err());
        }

        for i in 0..4 {
            x.push(i * 2);
        }

        assert!(!x.is_heap_allocated());
        assert_eq!(x.as_slice(), &[0, 2, 4, 6]);

        assert_eq!(x.remove(2), 4);
        assert_eq!(x.as_slice(), &[0, 2, 6]);

        assert_eq!(x.remove(2), 6);
        assert_eq!(x.as_slice(), &[0, 2]);

        assert_eq!(x.remove(1), 2);
        assert_eq!(x.as_slice(), &[0]);

        assert_eq!(x.remove(0), 0);
        assert_eq!(x.as_slice(), &[]);
        assert!(!x.is_heap_allocated());

        // trigger heap allocation
        for i in 0..8 {
            x.push(i * 2);
        }
        assert!(x.is_heap_allocated());
        assert_eq!(x.as_slice(), &[0, 2, 4, 6, 8, 10, 12, 14]);

        assert_eq!(x.remove(7), 14);
        assert_eq!(x.remove(0), 0);
    }

    #[test]
    fn inlinevec_remove_heap() {
        let mut x = InlineVecInner::<String, 4>::new();
        x.push("test".into());
        assert_eq!(x.len(), 1);
        assert_eq!(x.remove(0), "test");
        assert_eq!(x.len(), 0);
    }

    #[test]
    fn inlinevec() {
        let mut x = InlineVecInner::<usize, 4>::new();
        assert_eq!(x.len(), 0);
        assert_eq!(x.get(0), None);
        assert!(!x.is_heap_allocated());

        x.push(1337);
        assert_eq!(x.len(), 1);
        assert_eq!(x.get(0), Some(&1337));
        assert!(!x.is_heap_allocated());

        for v in 0..3 {
            x.push(v);
        }

        assert_eq!(x.len(), 4);

        // this call should move the vector to the heap
        x.push(42);
        assert_eq!(x.len(), 5);
        assert!(x.is_heap_allocated());

        // check that the old vector is still valid
        assert_eq!(x.get(0), Some(&1337));

        for v in 0..500 {
            x.push(v);
        }

        assert_eq!(x.len(), 505);
        assert!(x.is_heap_allocated());

        assert_eq!(x.get(1337), None);

        *x.get_mut(0).unwrap() = 444;
        assert_eq!(x.get(0), Some(&444));
        assert_eq!(x.get_mut(99999 /* out of bounds */), None);
    }

    #[test]
    fn inlinevec_as_slice() {
        let mut x = InlineVecInner::<usize, 4>::new();
        x.push(1337);
        x.push(42);
        x.push(17);
        assert_eq!(x.as_slice(), &[1337, 42, 17]);
        x.push(19);
        x.push(34);
        assert_eq!(x.as_slice(), &[1337, 42, 17, 19, 34]);
    }
}