#![forbid(missing_docs)]

//! Bytemuck, version 2!

use super::*;
use core::slice;

/// Performs a panic as a separate function.
///
/// "Outlining" is the opposite of inlining. By "outlining" the panic formatting
/// machinery, any function that *calls* this outlined panic stays lighter when
/// our caller is inlined.
#[cold]
#[cfg_attr(feature = "track_caller", track_caller)]
#[inline(never)]
pub fn outline_panic(msg: &str) -> ! {
  panic!("{msg}");
}

/// For types that don't contain any padding bytes.
///
/// ## Safety
/// * The type must not have any padding bytes.
/// * The type must not contain an [UnsafeCell]. Though a type with an
///   UnsafeCell might indeed not have any padding, the purpose of this trait is
///   to cast between reference types, and UnsafeCell alters the metadata of
///   references.
///
/// ## Blanket Impl
/// * This type has a blanket impl for all [Pod] types. If your type is `Pod`,
///   then just impl that trait and you'll get the `NoPadding` impl
///   automatically.
#[marker]
pub unsafe trait NoPadding {}
unsafe impl<T> NoPadding for T where T: Pod {}
unsafe impl<T> NoPadding for [T] where T: NoPadding {}
unsafe impl<T, const N: usize> NoPadding for [T; N] where T: NoPadding {}
unsafe impl NoPadding for bool {}
unsafe impl NoPadding for char {}
unsafe impl NoPadding for str {}

/// Re-view the value as an immutable byte slice.
#[inline]
#[must_use]
pub fn bytes_of<T: ?Sized + NoPadding>(t: &T) -> &[u8] {
  unsafe { slice::from_raw_parts((t as *const T).cast::<u8>(), size_of_val(t)) }
}

/// For types that are valid when all bytes are zero.
///
/// ## Safety
/// * The type must be valid to create at any time by simple using all zeroed
///   bytes.
///
/// ## Blanket Impl
/// * This type has a blanket impl for all [AnyBitPattern] types. If your type
///   is `AnyBitPattern`, then just impl that trait and you'll get the
///   `Zeroable` impl automatically.
#[marker]
pub unsafe trait Zeroable: Copy {}
unsafe impl<T> Zeroable for T where T: AnyBitPattern {}
unsafe impl<T, const N: usize> Zeroable for [T; N] where T: Zeroable {}
unsafe impl Zeroable for bool {}
unsafe impl Zeroable for char {}

/// Create a value with all bits zero.
#[inline]
#[must_use]
pub fn safe_zeroed<Z: Zeroable>() -> Z {
  unsafe { zeroed() }
}

/// Types that are allowed to be any bit pattern.
///
/// These types **might** have padding bytes in them, so you can't simply look
/// at all the bytes. However, you can turn a random pile of bytes *into* this
/// type (and then some of the source bytes might end up in the padding of the
/// target type, which is fine).
///
/// ## Safety
/// * All bytes of the type must be allowed to be any bit pattern.
///
/// ## Blanket Impl
/// * This type has a blanket impl for all [Pod] types. If your type is `Pod`,
///   then just impl that trait and you'll get the `AnyBitPattern` impl
///   automatically.
#[marker]
pub unsafe trait AnyBitPattern: Zeroable + Copy {}
unsafe impl<T> AnyBitPattern for T where T: Pod {}
unsafe impl<T, const N: usize> AnyBitPattern for [T; N] where T: AnyBitPattern {}

/// Attempt to turn a shared byte slice into a shared reference to a single
/// value.
///
/// ## Panics
/// * If the slice isn't the same size as `T`
/// * If the slice isn't sufficiently aligned.
#[inline]
#[must_use]
#[cfg_attr(feature = "track_caller", track_caller)]
pub fn as_type<T: AnyBitPattern>(bytes: &[u8]) -> &T {
  if size_of_val(bytes) != size_of::<T>() {
    outline_panic("SizeMismatch");
  } else if align_of::<T>() > 1
    && (bytes.as_ptr() as usize) % align_of::<T>() != 0
  {
    outline_panic("TargetAlignmentGreaterAndInputNotAligned");
  } else {
    unsafe { &*bytes.as_ptr().cast::<T>() }
  }
}

/// Takes a `T` from the front of the slice.
///
/// This does an unaligned read, so the alignment of `T` is not a concern.
///
/// ## Panics
/// * If the slice provided is shorter than `size_of::<T>()`.
#[inline]
#[must_use]
#[cfg_attr(feature = "track_caller", track_caller)]
pub fn take_val<T: AnyBitPattern>(bytes: &[u8]) -> (T, &[u8]) {
  if size_of_val(bytes) < size_of::<T>() {
    outline_panic("InsufficientBytes");
  } else {
    let t = unsafe { bytes.as_ptr().cast::<T>().read_unaligned() };
    let spare = unsafe { bytes.get_unchecked(size_of::<T>()..) };
    (t, spare)
  }
}

/// For types that are "plain old data".
///
/// ## Safety
/// * There may be additional rules not written down here.
/// * Don't mess this up.
/// * I mean it.
#[marker]
pub unsafe trait Pod: NoPadding + AnyBitPattern + Copy {}
unsafe impl<T, const N: usize> Pod for [T; N] where T: Pod {}
//
unsafe impl Pod for u8 {}
unsafe impl Pod for Option<NonZeroU8> {}
unsafe impl Pod for u16 {}
unsafe impl Pod for Option<NonZeroU16> {}
unsafe impl Pod for u32 {}
unsafe impl Pod for Option<NonZeroU32> {}
unsafe impl Pod for u64 {}
unsafe impl Pod for Option<NonZeroU64> {}
unsafe impl Pod for usize {}
unsafe impl Pod for Option<NonZeroUsize> {}
//
unsafe impl Pod for i8 {}
unsafe impl Pod for Option<NonZeroI8> {}
unsafe impl Pod for i16 {}
unsafe impl Pod for Option<NonZeroI16> {}
unsafe impl Pod for i32 {}
unsafe impl Pod for Option<NonZeroI32> {}
unsafe impl Pod for i64 {}
unsafe impl Pod for Option<NonZeroI64> {}
unsafe impl Pod for isize {}
unsafe impl Pod for Option<NonZeroIsize> {}
//
unsafe impl Pod for f32 {}
unsafe impl Pod for f64 {}

/// Cast by value.
///
/// ## Panics
/// * If the two types aren't the same size.
#[inline]
#[must_use]
#[cfg_attr(feature = "track_caller", track_caller)]
pub fn cast<A: NoPadding, B: AnyBitPattern>(a: A) -> B {
  if size_of::<A>() == size_of::<B>() {
    unsafe { (&a as *const A).cast::<B>().read_unaligned() }
  } else {
    outline_panic("size mismatch");
  }
}

/// Cast by shared slice.
///
/// This will adjust the size of the slice if the source and target types aren't
/// the same.
///
/// ## Panics
/// * If the target type is higher aligned and the source slice isn't
///   sufficiently aligned.
/// * If either type is zero sized and the other is not.
/// * If the output slice would have a fractional number of elements.
#[inline]
#[must_use]
#[cfg_attr(feature = "track_caller", track_caller)]
pub fn cast_slice<A: NoPadding, B: AnyBitPattern>(a: &[A]) -> &[B] {
  if align_of::<B>() > align_of::<A>()
    && (a.as_ptr() as usize) % align_of::<B>() != 0
  {
    // if we aren't sufficiently aligned, death.
    outline_panic("TargetAlignmentGreaterAndInputNotAligned");
  } else if size_of::<B>() == size_of::<A>() {
    // if the two types are the same size we can exit "early" to prevent having
    // to do a division.
    unsafe { slice::from_raw_parts(a.as_ptr().cast::<B>(), a.len()) }
  } else if size_of::<A>() == 0 || size_of::<B>() == 0 {
    // if just one type is zero sized we can't do it (and if they're both zero
    // sized then the above case already handled it).
    outline_panic("SizeMismatch");
  } else if size_of_val(a) % size_of::<B>() == 0 {
    // In this case the length will change. Also, let's really hope that LLVM
    // sees the mod and the div near each other and combines the two operations.
    let new_len = size_of_val(a) / size_of::<B>();
    unsafe { slice::from_raw_parts(a.as_ptr().cast::<B>(), new_len) }
  } else {
    // if we'd have partial bytes left over after the cast we disallow it.
    outline_panic("OutputSliceWouldHaveSlop");
  }
}

/// Cast by unique slice.
///
/// This will adjust the size of the slice if the source and target types aren't
/// the same.
///
/// ## Panics
/// * If the target type is higher aligned and the source slice isn't
///   sufficiently aligned.
/// * If either type is zero sized and the other is not.
/// * If the output slice would have a fractional number of elements.
#[inline]
#[must_use]
#[cfg_attr(feature = "track_caller", track_caller)]
pub fn cast_slice_mut<A: Pod, B: Pod>(a: &mut [A]) -> &mut [B] {
  if align_of::<B>() > align_of::<A>()
    && (a.as_ptr() as usize) % align_of::<B>() != 0
  {
    // if we aren't sufficiently aligned, death.
    outline_panic("TargetAlignmentGreaterAndInputNotAligned");
  } else if size_of::<B>() == size_of::<A>() {
    // if the two types are the same size we can exit "early" to prevent having
    // to do a division.
    unsafe { slice::from_raw_parts_mut(a.as_mut_ptr().cast::<B>(), a.len()) }
  } else if size_of::<A>() == 0 || size_of::<B>() == 0 {
    // if just one type is zero sized we can't do it (and if they're both zero
    // sized then the above case already handled it).
    outline_panic("SizeMismatch");
  } else if size_of_val(a) % size_of::<B>() == 0 {
    // In this case the length will change. Also, let's really hope that LLVM
    // sees the mod and the div near each other and combines the two operations.
    let new_len = size_of_val(a) / size_of::<B>();
    unsafe { slice::from_raw_parts_mut(a.as_mut_ptr().cast::<B>(), new_len) }
  } else {
    // if we'd have partial bytes left over after the cast we disallow it.
    outline_panic("OutputSliceWouldHaveSlop");
  }
}
