smallvec 2.0.0-beta.5

'Small vector' optimization: store up to a small number of items on the stack
Documentation
#[cfg(feature = "defmt")]
use defmt::{
    Format as DeFormat,
    Formatter as DeFormatter,
    write as dewrite
};
use {
    super::{
        error::SmallVecError,
        locatedlength::LocatedLength
    },
    core::{
        alloc::{
            Allocator,
            Layout
        },
        fmt::{
            Debug,
            Formatter,
            Result as Format
        },
        mem::{
            ManuallyDrop,
            MaybeUninit
        },
        ptr::{
            NonNull,
            copy_nonoverlapping
        }
    }
};

#[repr(C)]
pub union RawSmallVec<Item, const INLINE: usize> {
    pub inline: ManuallyDrop<[MaybeUninit<Item>; INLINE]>,
    pub heap: NonNull<[MaybeUninit<Item>]>
}

// The active member is tracked by `LocatedLength`, not by the union itself, so
// neither member can be read safely here.
impl<Item, const INLINE: usize> Debug for RawSmallVec<Item, INLINE> {
    fn fmt(&self, f: &mut Formatter<'_>) -> Format {
        f.debug_struct("RawSmallVec").finish_non_exhaustive()
    }
}

#[cfg(feature = "defmt")]
impl<Item, const INLINE: usize> DeFormat for RawSmallVec<Item, INLINE> {
    fn format(&self, fmt: DeFormatter<'_>) {
        dewrite!(fmt, "RawSmallVec {{ .. }}");
    }
}

impl<Item, const INLINE: usize> RawSmallVec<Item, INLINE> {
    #[inline]
    pub const fn new() -> Self {
        Self {
            inline: ManuallyDrop::new([const { MaybeUninit::uninit() }; INLINE])
        }
    }

    #[inline]
    pub const fn as_ptr_inline(&self) -> *const Item {
        // SAFETY: it is safe because we aren't reading the value, just getting
        // a reference to it. reading it would be UB potentially, but
        // for that downstream unsafe is required
        #[allow(unused_unsafe, reason = "Unsafe in MSRV")]
        (unsafe { &raw const self.inline }).cast()
    }

    #[inline]
    pub const fn as_mut_ptr_inline(&mut self) -> *mut Item {
        // SAFETY: same as above
        #[allow(unused_unsafe, reason = "Unsafe in MSRV")]
        (unsafe { &raw mut self.inline }).cast()
    }

    /// # Safety
    ///
    /// `spilled` must be true if and only if `self.heap` is the active member.
    #[inline(always)]
    pub const unsafe fn as_ptr(&self, spilled: bool) -> *const Item {
        if spilled {
            unsafe { self.heap.as_ptr().cast_const().cast() }
        } else {
            self.as_ptr_inline()
        }
    }

    /// # Safety
    ///
    /// `spilled` must be true if and only if `self.heap` is the active member.
    #[inline(always)]
    pub const unsafe fn as_mut_ptr(&mut self, spilled: bool) -> *mut Item {
        if spilled {
            unsafe { self.heap.as_ptr().cast() }
        } else {
            self.as_mut_ptr_inline()
        }
    }

    /// # Safety
    ///
    /// `spilled` must be true if and only if `self.heap` is the active member.
    #[inline(always)]
    pub const unsafe fn capacity(&self, spilled: bool) -> usize {
        if spilled {
            unsafe { self.heap.len() }
        } else if size_of::<Item>() == 0 {
            usize::MAX
        } else {
            INLINE
        }
    }

    /// # Safety
    ///
    /// `new_capacity` must be non zero, and greater or equal to the length.
    /// Item must not be a ZST.
    ///
    /// the allocator must be the same one the data was allocated with
    pub unsafe fn try_grow_raw<Heap: Allocator>(
        &mut self,
        length: LocatedLength<Item>,
        new_capacity: usize,
        allocator: &Heap
    ) -> Result<(), SmallVecError> {
        let (length, was_spilled) = length.parts();
        debug_assert!(size_of::<Item>() != 0);
        debug_assert!(new_capacity > 0 && new_capacity >= length);

        // SAFETY: the tag tells which member is active
        let pointer = unsafe { self.as_mut_ptr(was_spilled) };

        let new_layout =
            Layout::array::<Item>(new_capacity).map_err(|_| SmallVecError::CapacityOverflow)?;
        if new_layout.size() > isize::MAX as usize {
            return Err(SmallVecError::CapacityOverflow);
        }

        let new_ptr = if !was_spilled {
            // get a fresh allocation
            // `new_layout` has nonzero size.
            let new_ptr = Allocator::allocate(allocator, new_layout)
                .map_err(|_| SmallVecError::AllocationError(new_layout))?
                .cast();
            unsafe { copy_nonoverlapping(pointer, new_ptr.as_ptr(), length) };
            new_ptr
        } else {
            // use grow

            // this can't overflow since we already constructed an equivalent
            // layout during the previous allocation
            let old_layout = unsafe {
                Layout::from_size_align_unchecked(
                    self.heap.len() * size_of::<Item>(),
                    align_of::<Item>()
                )
            };

            // SAFETY: pointer was allocated with this allocator
            // old_layout is the same as the layout used to allocate the
            // previous memory block new_layout.size() is greater
            // than zero does not overflow when rounded up to
            // alignment. since it was constructed
            // with Layout::array
            unsafe {
                (if self.heap.len() < new_capacity {
                    <Heap as Allocator>::grow
                } else {
                    <Heap as Allocator>::shrink
                })(
                    allocator,
                    NonNull::new(pointer as *mut u8).unwrap(),
                    old_layout,
                    new_layout
                )
            }
            .map_err(|_| SmallVecError::AllocationError(new_layout))?
            .cast()
        };
        self.heap = NonNull::slice_from_raw_parts(new_ptr.cast(), new_capacity);
        Ok(())
    }
}

impl<Item, const INLINE: usize> Default for RawSmallVec<Item, INLINE> {
    fn default() -> Self { Self::new() }
}