musli-zerocopy 0.1.1

use core::alloc::{Layout, LayoutError};
use core::fmt;
use core::mem::{ManuallyDrop, MaybeUninit, size_of};
use core::ptr::NonNull;
use core::slice;

use crate::buf;
use crate::error::CoerceError;
use crate::pointer::{Pointee, Size};
use crate::traits::ZeroCopy;

/// A value which might or might not have been initialized.
///
/// This differs from the standard library [`MaybeUninit`] in that its methods
/// does not inherit the alignment of the inner value so it can correctly refer
/// to elements of `T` in unaligned memory. Which [`OwnedBuf`] might refer to.
///
/// # Examples
///
/// Writing to a pre-allocation location in an [`OwnedBuf`].
///
/// [`OwnedBuf`]: crate::buf::OwnedBuf
///
/// ```
/// use musli_zerocopy::{OwnedBuf, Ref, ZeroCopy};
/// use musli_zerocopy::mem::PackedMaybeUninit;
///
/// #[derive(ZeroCopy)]
/// #[repr(C)]
/// struct Custom { string: Ref<str> }
///
/// let mut buf = OwnedBuf::new();
///
/// let reference: Ref<PackedMaybeUninit<Custom>> = buf.store_uninit::<Custom>()?;
///
/// let string = buf.store_unsized("Hello World!")?;
///
/// buf.load_uninit_mut(reference)?.write(&Custom { string });
///
/// let reference = reference.assume_init();
///
/// assert_eq!(reference.offset(), 0);
///
/// let custom = buf.load(reference)?;
/// assert_eq!(buf.load(custom.string)?, "Hello World!");
/// # Ok::<_, musli_zerocopy::Error>(())
/// ```
#[repr(C, packed)]
pub union PackedMaybeUninit<T> {
    uninit: (),
    value: ManuallyDrop<T>,
}

impl<T> PackedMaybeUninit<T> {
    /// Creates a new `MaybeUninit<T>` in an uninitialized state.
    ///
    /// Note that dropping a `MaybeUninit<T>` will never call `T`'s drop code.
    /// It is your responsibility to make sure `T` gets dropped if it got
    /// initialized.
    ///
    /// See the [type-level documentation][MaybeUninit] for some examples.
    ///
    /// # Example
    ///
    /// ```
    /// use musli_zerocopy::mem::PackedMaybeUninit;
    ///
    /// let mut v: PackedMaybeUninit<u32> = PackedMaybeUninit::uninit();
    /// ```
    pub const fn uninit() -> Self {
        PackedMaybeUninit { uninit: () }
    }

    /// Write a value to the current location being pointed to.
    ///
    /// Note that we cannot return a reference to the written value, because it
    /// might not be aligned.
    ///
    /// We can however return the underlying bytes that were written because of
    /// this type, since they are now initialized.
    ///
    /// See the [type-level documentation][MaybeUninit] for some examples.
    ///
    /// # Example
    ///
    /// Writing to an uninitialized location on the stack:
    ///
    /// ```
    /// use musli_zerocopy::mem::PackedMaybeUninit;
    ///
    /// let mut v: PackedMaybeUninit<u32> = PackedMaybeUninit::uninit();
    /// assert_eq!(v.write(&10u32.to_le()), &[10, 0, 0, 0]);
    /// ```
    #[inline]
    pub fn write(&mut self, value: &T) -> &mut [u8]
    where
        T: ZeroCopy,
    {
        unsafe {
            let ptr = NonNull::new_unchecked(self as *mut Self as *mut MaybeUninit<u8>);
            buf::store_unaligned(ptr, value);
            slice::from_raw_parts_mut(ptr.as_ptr().cast(), size_of::<T>())
        }
    }
}

impl<T> fmt::Debug for PackedMaybeUninit<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("MaybeUninit").finish_non_exhaustive()
    }
}

impl<T> Pointee for PackedMaybeUninit<T>
where
    T: Pointee,
{
    type Metadata = T::Metadata;
    type Stored<O>
        = T::Stored<O>
    where
        O: Size;

    #[inline(always)]
    fn try_from_metadata<O>(metadata: Self::Metadata) -> Result<Self::Stored<O>, CoerceError>
    where
        O: Size,
    {
        T::try_from_metadata(metadata)
    }

    #[inline(always)]
    fn to_metadata<O>(stored: Self::Stored<O>) -> Self::Metadata
    where
        O: Size,
    {
        T::to_metadata(stored)
    }

    #[inline(always)]
    fn size(metadata: Self::Metadata) -> Option<usize> {
        T::size(metadata)
    }

    #[inline(always)]
    fn align(metadata: Self::Metadata) -> usize {
        T::align(metadata)
    }

    #[inline(always)]
    fn pointee_layout(metadata: Self::Metadata) -> Result<Layout, LayoutError> {
        T::pointee_layout(metadata)
    }
}