binius_utils/

serialization.rs

// Copyright 2024-2025 Irreducible Inc.

use bytes::{Buf, BufMut};
use hybrid_array::{Array, ArraySize};
use thiserror::Error;

/// Serialize data to a byte buffer.
pub trait SerializeBytes {
	fn serialize(&self, write_buf: impl BufMut) -> Result<(), SerializationError>;

	/// Serializes every element of `slice` back-to-back, with no length prefix.
	///
	/// The default loops, calling [`serialize`](Self::serialize) once per element.
	/// That default is always correct.
	/// Overriding this method is an optimization, never a correctness requirement.
	/// Override it only when `slice`'s byte layout already matches its serialized form.
	/// Then the whole call collapses to one bulk copy instead of N small writes.
	fn serialize_slice(slice: &[Self], write_buf: impl BufMut) -> Result<(), SerializationError>
	where
		Self: Sized,
	{
		let mut write_buf = write_buf;
		slice
			.iter()
			.try_for_each(|item| item.serialize(&mut write_buf))
	}
}

/// Deserialize data from a byte buffer.
pub trait DeserializeBytes: Sized {
	fn deserialize(read_buf: impl Buf) -> Result<Self, SerializationError>;
}

/// A type whose byte-serialized form has a fixed, compile-time-known size.
///
/// Implementors must guarantee that [`SerializeBytes::serialize`] writes exactly
/// [`BYTE_SIZE`](Self::BYTE_SIZE) bytes and that [`DeserializeBytes::deserialize`] reads exactly
/// [`BYTE_SIZE`](Self::BYTE_SIZE) bytes.
pub trait FixedSizeSerializeBytes: SerializeBytes + DeserializeBytes {
	/// The exact number of bytes written by `serialize` and read by `deserialize`.
	const BYTE_SIZE: usize;
}

macro_rules! impl_fixed_size_serialize_bytes {
	($ty:ty, $size:expr) => {
		impl FixedSizeSerializeBytes for $ty {
			const BYTE_SIZE: usize = $size;
		}
	};
}

impl_fixed_size_serialize_bytes!(u8, 1);
impl_fixed_size_serialize_bytes!(u16, 2);
impl_fixed_size_serialize_bytes!(u32, 4);
impl_fixed_size_serialize_bytes!(u64, 8);
impl_fixed_size_serialize_bytes!(u128, 16);
// `usize` is serialized as a `u32`.
impl_fixed_size_serialize_bytes!(usize, 4);
impl_fixed_size_serialize_bytes!(bool, 1);

#[derive(Error, Debug, Clone)]
pub enum SerializationError {
	#[error("Write buffer is full")]
	WriteBufferFull,
	#[error("Not enough data in read buffer to deserialize")]
	NotEnoughBytes,
	#[error("Unknown enum variant index {name}::{index}")]
	UnknownEnumVariant { name: &'static str, index: u8 },
	#[error("FromUtf8Error: {0}")]
	FromUtf8Error(#[from] std::string::FromUtf8Error),
	#[error("Invalid construction of {name}")]
	InvalidConstruction { name: &'static str },
	#[error("usize {size} is too large to serialize (max is {max})", max = u32::MAX)]
	UsizeTooLarge { size: usize },
}

impl<T: SerializeBytes + ?Sized> SerializeBytes for &T {
	fn serialize(&self, write_buf: impl BufMut) -> Result<(), SerializationError> {
		(**self).serialize(write_buf)
	}
}

impl SerializeBytes for usize {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		let value: u32 = (*self)
			.try_into()
			.map_err(|_| SerializationError::UsizeTooLarge { size: *self })?;
		SerializeBytes::serialize(&value, &mut write_buf)
	}
}

impl DeserializeBytes for usize {
	fn deserialize(mut read_buf: impl Buf) -> Result<Self, SerializationError>
	where
		Self: Sized,
	{
		let value: u32 = DeserializeBytes::deserialize(&mut read_buf)?;
		Ok(value as Self)
	}
}

impl SerializeBytes for u128 {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		assert_enough_space_for(&write_buf, std::mem::size_of::<Self>())?;
		write_buf.put_u128_le(*self);
		Ok(())
	}
}

impl DeserializeBytes for u128 {
	fn deserialize(mut read_buf: impl Buf) -> Result<Self, SerializationError>
	where
		Self: Sized,
	{
		assert_enough_data_for(&read_buf, std::mem::size_of::<Self>())?;
		Ok(read_buf.get_u128_le())
	}
}

impl SerializeBytes for u64 {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		assert_enough_space_for(&write_buf, std::mem::size_of::<Self>())?;
		write_buf.put_u64_le(*self);
		Ok(())
	}
}

impl DeserializeBytes for u64 {
	fn deserialize(mut read_buf: impl Buf) -> Result<Self, SerializationError>
	where
		Self: Sized,
	{
		assert_enough_data_for(&read_buf, std::mem::size_of::<Self>())?;
		Ok(read_buf.get_u64_le())
	}
}

impl SerializeBytes for u32 {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		assert_enough_space_for(&write_buf, std::mem::size_of::<Self>())?;
		write_buf.put_u32_le(*self);
		Ok(())
	}
}

impl DeserializeBytes for u32 {
	fn deserialize(mut read_buf: impl Buf) -> Result<Self, SerializationError>
	where
		Self: Sized,
	{
		assert_enough_data_for(&read_buf, std::mem::size_of::<Self>())?;
		Ok(read_buf.get_u32_le())
	}
}

impl SerializeBytes for u16 {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		assert_enough_space_for(&write_buf, std::mem::size_of::<Self>())?;
		write_buf.put_u16_le(*self);
		Ok(())
	}
}

impl DeserializeBytes for u16 {
	fn deserialize(mut read_buf: impl Buf) -> Result<Self, SerializationError>
	where
		Self: Sized,
	{
		assert_enough_data_for(&read_buf, std::mem::size_of::<Self>())?;
		Ok(read_buf.get_u16_le())
	}
}

impl SerializeBytes for u8 {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		assert_enough_space_for(&write_buf, std::mem::size_of::<Self>())?;
		write_buf.put_u8(*self);
		Ok(())
	}
}

impl DeserializeBytes for u8 {
	fn deserialize(mut read_buf: impl Buf) -> Result<Self, SerializationError>
	where
		Self: Sized,
	{
		assert_enough_data_for(&read_buf, std::mem::size_of::<Self>())?;
		Ok(read_buf.get_u8())
	}
}

impl SerializeBytes for bool {
	fn serialize(&self, write_buf: impl BufMut) -> Result<(), SerializationError> {
		u8::serialize(&(*self as u8), write_buf)
	}
}

impl DeserializeBytes for bool {
	fn deserialize(read_buf: impl Buf) -> Result<Self, SerializationError>
	where
		Self: Sized,
	{
		Ok(u8::deserialize(read_buf)? != 0)
	}
}

impl<T> SerializeBytes for std::marker::PhantomData<T> {
	fn serialize(&self, _write_buf: impl BufMut) -> Result<(), SerializationError> {
		Ok(())
	}
}

impl<T> DeserializeBytes for std::marker::PhantomData<T> {
	fn deserialize(_read_buf: impl Buf) -> Result<Self, SerializationError>
	where
		Self: Sized,
	{
		Ok(Self)
	}
}

impl SerializeBytes for &str {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		let bytes = self.as_bytes();
		SerializeBytes::serialize(&bytes.len(), &mut write_buf)?;
		assert_enough_space_for(&write_buf, bytes.len())?;
		write_buf.put_slice(bytes);
		Ok(())
	}
}

impl SerializeBytes for String {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		SerializeBytes::serialize(&self.as_str(), &mut write_buf)
	}
}

impl DeserializeBytes for String {
	fn deserialize(mut read_buf: impl Buf) -> Result<Self, SerializationError>
	where
		Self: Sized,
	{
		let len = DeserializeBytes::deserialize(&mut read_buf)?;
		assert_enough_data_for(&read_buf, len)?;
		Ok(Self::from_utf8(read_buf.copy_to_bytes(len).to_vec())?)
	}
}

impl<T: SerializeBytes> SerializeBytes for [T] {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		SerializeBytes::serialize(&self.len(), &mut write_buf)?;
		T::serialize_slice(self, write_buf)
	}
}

impl<T: SerializeBytes> SerializeBytes for Vec<T> {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		SerializeBytes::serialize(self.as_slice(), &mut write_buf)
	}
}

impl<T: DeserializeBytes> DeserializeBytes for Vec<T> {
	fn deserialize(mut read_buf: impl Buf) -> Result<Self, SerializationError>
	where
		Self: Sized,
	{
		let len: usize = DeserializeBytes::deserialize(&mut read_buf)?;
		(0..len)
			.map(|_| DeserializeBytes::deserialize(&mut read_buf))
			.collect()
	}
}

impl<T: SerializeBytes> SerializeBytes for Option<T> {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		match self {
			Some(value) => {
				SerializeBytes::serialize(&true, &mut write_buf)?;
				SerializeBytes::serialize(value, &mut write_buf)?;
			}
			None => {
				SerializeBytes::serialize(&false, write_buf)?;
			}
		}
		Ok(())
	}
}

impl<T: DeserializeBytes> DeserializeBytes for Option<T> {
	fn deserialize(mut read_buf: impl Buf) -> Result<Self, SerializationError>
	where
		Self: Sized,
	{
		Ok(match bool::deserialize(&mut read_buf)? {
			true => Some(T::deserialize(&mut read_buf)?),
			false => None,
		})
	}
}

impl<U: SerializeBytes, V: SerializeBytes> SerializeBytes for (U, V) {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		U::serialize(&self.0, &mut write_buf)?;
		V::serialize(&self.1, write_buf)
	}
}

impl<U: DeserializeBytes, V: DeserializeBytes> DeserializeBytes for (U, V) {
	fn deserialize(mut read_buf: impl Buf) -> Result<Self, SerializationError>
	where
		Self: Sized,
	{
		Ok((U::deserialize(&mut read_buf)?, V::deserialize(read_buf)?))
	}
}

impl<T: SerializeBytes, const N: usize> SerializeBytes for [T; N] {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		for val in self {
			val.serialize(&mut write_buf)?;
		}
		Ok(())
	}
}

impl<T: DeserializeBytes, const N: usize> DeserializeBytes for [T; N] {
	fn deserialize(mut read_buf: impl Buf) -> Result<Self, SerializationError>
	where
		Self: Sized,
	{
		array_util::try_from_fn(|_| T::deserialize(&mut read_buf))
	}
}

impl<U: ArraySize> SerializeBytes for Array<u8, U> {
	fn serialize(&self, mut write_buf: impl BufMut) -> Result<(), SerializationError> {
		assert_enough_space_for(&write_buf, U::USIZE)?;
		write_buf.put_slice(self);
		Ok(())
	}
}

impl<U: ArraySize> DeserializeBytes for Array<u8, U> {
	fn deserialize(mut read_buf: impl Buf) -> Result<Self, SerializationError> {
		assert_enough_data_for(&read_buf, U::USIZE)?;
		let mut ret = Self::default();
		read_buf.copy_to_slice(&mut ret);
		Ok(ret)
	}
}

#[inline]
pub fn assert_enough_space_for(
	write_buf: &impl BufMut,
	size: usize,
) -> Result<(), SerializationError> {
	if write_buf.remaining_mut() < size {
		return Err(SerializationError::WriteBufferFull);
	}
	Ok(())
}

#[inline]
pub fn assert_enough_data_for(read_buf: &impl Buf, size: usize) -> Result<(), SerializationError> {
	if read_buf.remaining() < size {
		return Err(SerializationError::NotEnoughBytes);
	}
	Ok(())
}

#[cfg(test)]
mod tests {
	use hybrid_array::sizes::U32;
	use rand::prelude::*;

	use super::*;

	#[test]
	fn test_generic_array_serialize_deserialize() {
		let mut rng = StdRng::seed_from_u64(0);

		let mut data = Array::<u8, U32>::default();
		rng.fill_bytes(&mut data);

		let mut buf = Vec::new();
		data.serialize(&mut buf).unwrap();

		let data_deserialized = Array::<u8, U32>::deserialize(&mut buf.as_slice()).unwrap();
		assert_eq!(data_deserialized, data);
	}
}