storages 0.1.1

Storage backends for collections
use super::{Buffer, UnmanagedBuffer};
use alloc::alloc::Global;
use core::{
    alloc::{AllocError, AllocRef, Layout},
    marker::{PhantomData, Unsize},
    mem,
    ops::CoerceUnsized,
    ptr::NonNull,
    slice,
};

#[derive(Copy, Clone)]
enum Init {
    Unspecified,
    Zeroed,
}

pub struct AllocatedBuffer<T: ?Sized, A: ?Sized = Global> {
    ptr: NonNull<T>,
    _owned: PhantomData<T>,
    _marker: PhantomData<fn(*const A)>,
}

impl<T: ?Sized, A: ?Sized> AllocatedBuffer<T, A> {
    pub unsafe fn from_raw(ptr: NonNull<T>) -> Self {
        Self {
            ptr,
            _owned: PhantomData,
            _marker: PhantomData,
        }
    }
}

impl<T> AllocatedBuffer<T> {
    pub fn new() -> Result<Self, AllocError> {
        Self::new_in(&Global)
    }

    pub fn new_zeroed() -> Result<Self, AllocError> {
        Self::new_zeroed_in(&Global)
    }
}

impl<T, A: ?Sized + AllocRef> AllocatedBuffer<T, A> {
    fn allocate_in(allocator: &A, init: Init) -> Result<Self, AllocError> {
        let layout = Layout::new::<T>();
        let ptr = match init {
            Init::Unspecified => allocator.alloc(layout)?,
            Init::Zeroed => allocator.alloc_zeroed(layout)?,
        };
        unsafe { Ok(Self::from_raw(ptr.as_non_null_ptr().cast())) }
    }

    pub fn new_in(allocator: &A) -> Result<Self, AllocError> {
        Self::allocate_in(allocator, Init::Unspecified)
    }

    pub fn new_zeroed_in(allocator: &A) -> Result<Self, AllocError> {
        Self::allocate_in(allocator, Init::Zeroed)
    }
}

impl<T, A: ?Sized + AllocRef> AllocatedBuffer<[T], A> {
    fn capacity_from_bytes(bytes: usize) -> usize {
        debug_assert_ne!(mem::size_of::<T>(), 0);
        bytes / mem::size_of::<T>()
    }

    #[allow(clippy::map_err_ignore)]
    fn allocate_slice(allocator: &A, len: usize, init: Init) -> Result<Self, AllocError> {
        let ptr = if mem::size_of::<T>() == 0 {
            NonNull::slice_from_raw_parts(NonNull::dangling(), 0)
        } else {
            let layout = Layout::array::<T>(len).map_err(|_| AllocError)?;
            alloc_guard(layout.size()).map_err(|_| AllocError)?;
            let ptr = match init {
                Init::Unspecified => allocator.alloc(layout)?,
                Init::Zeroed => allocator.alloc_zeroed(layout)?,
            };

            NonNull::slice_from_raw_parts(
                ptr.as_non_null_ptr().cast(),
                Self::capacity_from_bytes(ptr.len()),
            )
        };
        unsafe { Ok(Self::from_raw(ptr)) }
    }

    pub fn new_slice(allocator: &A, len: usize) -> Result<Self, AllocError> {
        Self::allocate_slice(allocator, len, Init::Unspecified)
    }

    pub fn new_slice_zeroed(allocator: &A, len: usize) -> Result<Self, AllocError> {
        Self::allocate_slice(allocator, len, Init::Zeroed)
    }
}

impl<T: ?Sized, A: ?Sized + AllocRef> Buffer<T> for AllocatedBuffer<T, A> {
    type ExternalData = A;

    fn as_ptr(&self, _data: &Self::ExternalData) -> *const T {
        self.ptr.as_ptr()
    }

    fn as_mut_ptr(&mut self, _data: &Self::ExternalData) -> *mut T {
        self.ptr.as_ptr()
    }
}

impl<T, A: ?Sized + AllocRef> Buffer<mem::MaybeUninit<T>> for AllocatedBuffer<T, A> {
    type ExternalData = A;

    fn as_ptr(&self, _data: &Self::ExternalData) -> *const mem::MaybeUninit<T> {
        self.ptr.as_ptr().cast()
    }

    fn as_mut_ptr(&mut self, _data: &Self::ExternalData) -> *mut mem::MaybeUninit<T> {
        self.ptr.as_ptr().cast()
    }
}

impl<T, A: ?Sized + AllocRef> Buffer<[mem::MaybeUninit<T>]> for AllocatedBuffer<[T], A> {
    type ExternalData = A;

    fn as_ptr(&self, _data: &Self::ExternalData) -> *const [mem::MaybeUninit<T>] {
        unsafe { slice::from_raw_parts(self.ptr.cast().as_ptr(), self.ptr.len()) }
    }

    fn as_mut_ptr(&mut self, _data: &Self::ExternalData) -> *mut [mem::MaybeUninit<T>] {
        unsafe { slice::from_raw_parts_mut(self.ptr.cast().as_ptr(), self.ptr.len()) }
    }
}

impl<T: ?Sized, A: AllocRef> UnmanagedBuffer<T> for AllocatedBuffer<T, A> {
    unsafe fn free_unchecked(&mut self, allocator: &Self::ExternalData) {
        let size = mem::size_of_val(self.ptr.as_ref());
        let align = mem::align_of_val(self.ptr.as_ref());
        let layout = Layout::from_size_align_unchecked(size, align);
        allocator.dealloc(self.ptr.cast(), layout);
    }
}

#[inline]
const fn alloc_guard(alloc_size: usize) -> Result<(), AllocError> {
    if usize::BITS < 64 && alloc_size > isize::MAX as usize {
        Err(AllocError)
    } else {
        Ok(())
    }
}

impl<T: ?Sized + Unsize<U>, U: ?Sized, A: AllocRef> CoerceUnsized<AllocatedBuffer<U, A>>
    for AllocatedBuffer<T, A>
{
}