nibblecode 0.1.0

use core::alloc::Layout;
use core::marker::{PhantomData, PhantomPinned};
use core::mem::{ManuallyDrop, MaybeUninit, offset_of};
use core::ptr::{self, Alignment, copy_nonoverlapping};
use core::str::from_utf8;
use core::{slice, str};

use crate::list::ArchivedList;
use crate::primitive::{ArchivedUsize, FixedUsize};
use crate::string::{ArchivedString, DecodedString, INLINE_CAPACITY, OUT_OF_LINE_CAPACITY};
use crate::tuple::{
	ArchivedTuple1, ArchivedTuple2, ArchivedTuple3, ArchivedTuple4, ArchivedTuple5, ArchivedTuple6,
	ArchivedTuple7, ArchivedTuple8, ArchivedTuple9, ArchivedTuple10, ArchivedTuple11,
	ArchivedTuple12, ArchivedTuple13,
};
use crate::util::{align_offset, max_alignment, offset_archived};
use crate::{Serialize, SerializeError, VerifyError};

mod option;
mod primitive;
mod result;

macro_rules! impl_tuple {
    ($name:ident, $($type:ident $index:tt),+) => {
		impl<$($type: Serialize),+> Serialize for ($($type,)+) {
			type Archived = $name<$($type::Archived,)+>;
			const ALIGN: Alignment = max_alignment(
				[Alignment::of::<Self::Archived>(), $($type::ALIGN),+]
			);
			const COPY_OPTIMIZATION: bool = $(
				offset_of!(Self, $index) == offset_of!(Self::Archived, $index)
					&& $type::COPY_OPTIMIZATION
			)&&+;

			unsafe fn serialize(
				&self,
				out: *mut MaybeUninit<Self::Archived>,
				mut heap: *mut MaybeUninit<u8>,
			) -> usize {
				let heap_start = heap;
			unsafe {
					$(
						heap = heap.add(
							self.$index.serialize(
								(&raw mut (*out.cast::<Self::Archived>()).$index).cast(),
								heap
							)
						);
					)+
					heap.offset_from_unsigned(heap_start)
				}
			}

			fn serialized_size(&self, mut offset: usize) -> Result<usize, SerializeError> {
				let offset_start = offset;
				$(offset += self.$index.serialized_size(offset)?;)+
				Ok(offset - offset_start)
			}

			#[inline(always)]
			unsafe fn verify(
				this: *const Self::Archived,
				buffer_end: *const u8
			) -> Result<(), VerifyError> {
				unsafe {
					$($type::verify(&raw const (*this).$index, buffer_end)?;)+
				}
				Ok(())
			}
		}
    }
}

impl_tuple!(ArchivedTuple1, T0 0);
impl_tuple!(ArchivedTuple2, T0 0, T1 1);
impl_tuple!(ArchivedTuple3, T0 0, T1 1, T2 2);
impl_tuple!(ArchivedTuple4, T0 0, T1 1, T2 2, T3 3);
impl_tuple!(ArchivedTuple5, T0 0, T1 1, T2 2, T3 3, T4 4);
impl_tuple!(ArchivedTuple6, T0 0, T1 1, T2 2, T3 3, T4 4, T5 5);
impl_tuple!(ArchivedTuple7, T0 0, T1 1, T2 2, T3 3, T4 4, T5 5, T6 6);
impl_tuple!(ArchivedTuple8, T0 0, T1 1, T2 2, T3 3, T4 4, T5 5, T6 6, T7 7);
impl_tuple!(
	ArchivedTuple9, T0 0, T1 1, T2 2, T3 3, T4 4, T5 5, T6 6, T7 7, T8 8
);
impl_tuple!(
	ArchivedTuple10, T0 0, T1 1, T2 2, T3 3, T4 4, T5 5, T6 6, T7 7, T8 8, T9 9
);
impl_tuple!(
	ArchivedTuple11, T0 0, T1 1, T2 2, T3 3, T4 4, T5 5, T6 6, T7 7, T8 8, T9 9,
	T10 10
);
impl_tuple!(
	ArchivedTuple12, T0 0, T1 1, T2 2, T3 3, T4 4, T5 5, T6 6, T7 7, T8 8, T9 9,
	T10 10, T11 11
);
impl_tuple!(
	ArchivedTuple13, T0 0, T1 1, T2 2, T3 3, T4 4, T5 5, T6 6, T7 7, T8 8, T9 9,
	T10 10, T11 11, T12 12
);

// Arrays

impl<T: Serialize, const N: usize> Serialize for [T; N] {
	type Archived = [T::Archived; N];
	const ALIGN: Alignment = T::ALIGN;
	const COPY_OPTIMIZATION: bool = T::COPY_OPTIMIZATION;

	unsafe fn serialize(
		&self,
		out: *mut MaybeUninit<[T::Archived; N]>,
		mut heap: *mut MaybeUninit<u8>,
	) -> usize {
		let start_heap = heap;

		unsafe {
			for (item, item_out) in self
				.iter()
				.zip(&mut *out.cast::<[MaybeUninit<T::Archived>; N]>())
			{
				heap = heap.add(item.serialize(ptr::from_mut(item_out), heap));
			}

			heap.offset_from_unsigned(start_heap)
		}
	}

	fn serialized_size(&self, offset: usize) -> Result<usize, SerializeError> {
		self.iter()
			.try_fold(offset, |acc, x| {
				x.serialized_size(acc).map(|size| acc + size)
			})
			.map(|final_offset| final_offset - offset)
	}

	#[inline(always)]
	unsafe fn verify(
		this: *const [T::Archived; N],
		buffer_end: *const u8,
	) -> Result<(), VerifyError> {
		unsafe {
			for item in &*this {
				T::verify(ptr::from_ref(item), buffer_end)?;
			}
		}

		Ok(())
	}
}

// Slices

impl<T: Serialize> Serialize for [T] {
	type Archived = ArchivedList<T::Archived>;
	const ALIGN: Alignment = max_alignment([T::ALIGN, Alignment::of::<ArchivedUsize>()]);

	unsafe fn serialize(
		&self,
		out: *mut MaybeUninit<ArchivedList<T::Archived>>,
		heap_start: *mut MaybeUninit<u8>,
	) -> usize {
		let align = align_offset::<T::Archived>(heap_start as usize);

		unsafe {
			// Pointer to the list element we're currently on. Aligned to the align of the items.
			let mut ptr = heap_start.add(align);

			*out = MaybeUninit::new(ArchivedList::new(
				self.len(),
				ptr.offset_from_unsigned(out.cast()),
			));

			if T::COPY_OPTIMIZATION {
				copy_nonoverlapping((self as *const [T]).cast::<T>(), ptr.cast(), self.len());

				align + size_of::<T::Archived>() * self.len()
			} else {
				// Pointer to the rest of the heap
				let mut heap = ptr.add(size_of::<T::Archived>() * self.len());

				for item in self {
					heap = heap.add(item.serialize(ptr.cast(), heap));
					ptr = ptr.add(size_of::<T::Archived>());
				}

				// Return the number of bytes that we and our children wrote
				heap.offset_from_unsigned(heap_start)
			}
		}
	}

	fn serialized_size(&self, offset: usize) -> Result<usize, SerializeError> {
		if (FixedUsize::MAX as usize) < self.len() {
			Err(SerializeError::ListTooLong)
		} else {
			let aligned_offset = offset + align_offset::<T::Archived>(offset);

			if (FixedUsize::MAX as usize) < aligned_offset {
				Err(SerializeError::OverflowedPointer)
			} else {
				self.iter()
					.try_fold(
						aligned_offset + size_of::<T::Archived>() * self.len(),
						|acc, x| x.serialized_size(acc).map(|size| acc + size),
					)
					.map(|final_offset| final_offset - offset)
			}
		}
	}

	#[inline(always)]
	unsafe fn verify(
		this: *const ArchivedList<T::Archived>,
		buffer_end: *const u8,
	) -> Result<(), VerifyError> {
		unsafe {
			let len = (*this).len.to_native() as usize;

			// Make sure the length is valid for this type
			let Ok(layout) = Layout::array::<T::Archived>(len) else {
				return Err(VerifyError::ListTooLong);
			};

			let mut items = offset_archived(
				this.cast(),
				(*this).offset.to_native() as usize,
				layout.size(),
				buffer_end,
				Alignment::of::<T::Archived>(),
			)?;

			for _ in 0..len {
				T::verify(items.cast(), buffer_end)?;
				items = items.add(size_of::<T::Archived>());
			}
		}

		Ok(())
	}
}

// `str`

/// Quickly check if an inline string is ascii.
fn is_inline_ascii(bytes: &[u8; INLINE_CAPACITY]) -> bool {
	for &c in bytes {
		match c {
			// ASCII
			..0x7F => (),
			// Not ASCII
			0x7F..0xFF => return false,
			// End of string
			0xFF => return true,
		}
	}

	true
}

impl Serialize for str {
	type Archived = ArchivedString;
	const ALIGN: Alignment = Alignment::of::<ArchivedUsize>();

	unsafe fn serialize(
		&self,
		out: *mut MaybeUninit<ArchivedString>,
		heap: *mut MaybeUninit<u8>,
	) -> usize {
		unsafe {
			if self.len() <= INLINE_CAPACITY {
				// SAFETY: The caller has guaranteed that `out` points to a
				// dereferenceable location.
				let out_bytes = &mut (*out.cast::<ArchivedString>()).inline;

				out_bytes[..self.len()].copy_from_slice(self.as_bytes());

				if self.len() != INLINE_CAPACITY {
					out_bytes[self.len()] = 0xff;
				}

				0
			} else {
				let l = self.len();
				// Little-endian: insert 10 as the 7th and 8th bits
				#[cfg(not(feature = "big_endian"))]
				let l = (l & 0b0011_1111) | 0b1000_0000 | ((l & 0b1100_0000) << 2);
				// Big-endian: set the top two bits to 10
				#[cfg(feature = "big_endian")]
				let l = l & (FixedUsize::MAX >> 2) | (1 << FixedUsize::BITS - 1);

				*out = MaybeUninit::new(ArchivedString {
					out_of_line: ManuallyDrop::new(ArchivedList::new(
						l,
						heap.offset_from_unsigned(out.cast()),
					)),
				});

				copy_nonoverlapping(<*const str>::cast(self), heap, self.len());

				self.len()
			}
		}
	}

	fn serialized_size(&self, _: usize) -> Result<usize, SerializeError> {
		if OUT_OF_LINE_CAPACITY < self.len() {
			Err(SerializeError::StringTooLong)
		} else {
			Ok(if self.len() <= INLINE_CAPACITY {
				0
			} else {
				self.len()
			})
		}
	}

	#[inline(always)]
	unsafe fn verify(
		this: *const ArchivedString,
		buffer_end: *const u8,
	) -> Result<(), VerifyError> {
		let bytes = unsafe {
			match (*this).decode() {
				DecodedString::Inline => {
					if is_inline_ascii(&(*this).inline) {
						return Ok(());
					}
					if let Some(extra_bytes) = (*this).inline.iter().position(|&x| x == 0xff) {
						&(&(*this).inline)[..extra_bytes]
					} else {
						&(*this).inline
					}
				}
				DecodedString::OutOfLine { len, offset } => {
					if (isize::MAX as usize) < len {
						return Err(VerifyError::StringTooLong);
					}

					let bytes = slice::from_raw_parts(
						offset_archived(this.cast(), offset, len, buffer_end, Alignment::MIN)?,
						len,
					);

					if bytes.is_ascii() {
						return Ok(());
					}

					bytes
				}
			}
		};

		if let Err(utf8_error) = from_utf8(bytes) {
			Err(VerifyError::Utf8Error(utf8_error))
		} else {
			Ok(())
		}
	}
}

// PhantomData

impl<T> Serialize for PhantomData<T> {
	type Archived = PhantomData<T>;
	const ALIGN: Alignment = Alignment::MIN;
	const COPY_OPTIMIZATION: bool = true;

	unsafe fn serialize(
		&self,
		_: *mut MaybeUninit<PhantomData<T>>,
		_: *mut MaybeUninit<u8>,
	) -> usize {
		0
	}

	fn serialized_size(&self, _: usize) -> Result<usize, SerializeError> {
		Ok(0)
	}

	#[inline(always)]
	unsafe fn verify(_: *const PhantomData<T>, _: *const u8) -> Result<(), VerifyError> {
		Ok(())
	}
}

// PhantomPinned

impl Serialize for PhantomPinned {
	type Archived = PhantomPinned;
	const ALIGN: Alignment = Alignment::MIN;
	const COPY_OPTIMIZATION: bool = true;

	unsafe fn serialize(
		&self,
		_: *mut MaybeUninit<PhantomPinned>,
		_: *mut MaybeUninit<u8>,
	) -> usize {
		0
	}

	fn serialized_size(&self, _: usize) -> Result<usize, SerializeError> {
		Ok(0)
	}

	#[inline(always)]
	unsafe fn verify(_: *const PhantomPinned, _: *const u8) -> Result<(), VerifyError> {
		Ok(())
	}
}

// `ManuallyDrop`

impl<T: Serialize> Serialize for ManuallyDrop<T> {
	type Archived = ManuallyDrop<T::Archived>;
	const ALIGN: Alignment = T::ALIGN;
	const COPY_OPTIMIZATION: bool = true;

	unsafe fn serialize(
		&self,
		out: *mut MaybeUninit<ManuallyDrop<T::Archived>>,
		heap: *mut MaybeUninit<u8>,
	) -> usize {
		unsafe { (**self).serialize(out.cast(), heap) }
	}

	fn serialized_size(&self, offset: usize) -> Result<usize, SerializeError> {
		(**self).serialized_size(offset)
	}

	#[inline(always)]
	unsafe fn verify(
		this: *const ManuallyDrop<T::Archived>,
		buffer_end: *const u8,
	) -> Result<(), VerifyError> {
		unsafe { T::verify(this.cast(), buffer_end) }
	}
}

// References

impl<T: Serialize + ?Sized> Serialize for &T {
	type Archived = T::Archived;
	const ALIGN: Alignment = T::ALIGN;
	const COPY_OPTIMIZATION: bool = T::COPY_OPTIMIZATION;

	unsafe fn serialize(
		&self,
		out: *mut MaybeUninit<T::Archived>,
		heap: *mut MaybeUninit<u8>,
	) -> usize {
		unsafe { (**self).serialize(out, heap) }
	}

	fn serialized_size(&self, offset: usize) -> Result<usize, SerializeError> {
		(**self).serialized_size(offset)
	}

	#[inline(always)]
	unsafe fn verify(this: *const T::Archived, buffer_end: *const u8) -> Result<(), VerifyError> {
		unsafe { T::verify(this, buffer_end) }
	}
}

#[cfg(test)]
mod tests {
	use core::marker::{PhantomData, PhantomPinned};
	use core::mem::ManuallyDrop;

	use crate::test::{roundtrip, roundtrip_with};
	use crate::tuple::ArchivedTuple3;

	#[test]
	fn roundtrip_tuple() {
		roundtrip_with(&(24, true, 16f32), |(a, b, c), ArchivedTuple3(d, e, f)| {
			assert_eq!(a, d);
			assert_eq!(b, e);
			assert_eq!(c, f);
		});
	}

	#[test]
	fn roundtrip_array() {
		roundtrip(&[1, 2, 3, 4, 5, 6]);
		roundtrip(&[(); 0]);
		roundtrip(&[(), (), (), ()]);
	}

	#[test]
	fn roundtrip_phantoms() {
		roundtrip(&PhantomData::<&'static u8>);
		roundtrip(&PhantomPinned);
	}

	#[test]
	fn roundtrip_manually_drop() {
		roundtrip(&ManuallyDrop::new(123i8));
	}
}