backed_archive.rs

use std::{marker::PhantomData, pin::Pin};

use bytecheck::CheckBytes;
use rkyv::{
    archived_value, archived_value_mut, check_archived_value,
    ser::{serializers::AllocSerializer, Serializer},
    validation::validators::DefaultValidator,
    Archive, Serialize,
};

pub type DefaultSerializer = AllocSerializer<4096>;

pub enum SerializeError<T> {
    /// The value failed to be serialized correctly.
    FailedToSerialize(String),

    /// The backing store was not big enough to fit the serialized version of the value.
    ///
    /// The original value that was given is returned, along with the minimum size that the backing
    /// store must be sized to hold the serialized value.  Providing a backing store that is larger
    /// than the given value is acceptable, but not necessary.
    BackingStoreTooSmall(T, usize),
}

pub struct DeserializeError;

/// Backed wrapper for any type that implements `Archive`.
///
/// For any backing store that can provide references to an underlying byte slice of suitable size,
/// we can deserialize and serialize a type that is [archivable]. `BackedArchive` provides specific
/// entrypoints to either deserialize the given type from the backing store, as well as serialize a
/// provided value to the backing store.
///
/// Once wrapped, the archived type can be accessed either immutably or mutably.  This provides a
/// simple mechanism to use a variety of backing stores, such as `Vec<u8>` or a memory-mapped
/// region.  This can be used to avoid serializing to intermediate buffers when possible.
pub struct BackedArchive<B, T> {
    backing: B,
    _archive: PhantomData<T>,
}

impl<B, T> BackedArchive<B, T>
where
    B: AsRef<[u8]> + AsMut<[u8]>,
    T: Archive,
    T::Archived: Unpin,
{
    /// Deserializes the archived value from the backing store and wraps it.
    ///
    /// # Errors
    /// If the data in the backing store is not valid for `T`, an error variant will be returned.
    pub fn from_backing<C>(backing: B) -> Result<BackedArchive<B, T>, DeserializeError>
    where
        for<'a> T::Archived: CheckBytes<DefaultValidator<'a>>,
    {
        // Validate that the input is, well, valid.
        let buf = backing.as_ref();
        let _ = check_archived_value::<T>(buf, 0).map_err(|_| DeserializeError)?;

        // Now that we know the buffer fits T, we're good to go!
        Ok(Self {
            backing,
            _archive: PhantomData,
        })
    }

    /// Serializes the provided value to the backing store and wraps it.
    ///
    /// # Errors
    /// If there is an error during serializing of the value, an error variant will be returned that
    /// describes the error.  If the backing store is too small to hold the serialized version of
    /// the value, an error variant will be returned defining the minimum size the backing store
    /// must be, as well containing the value that failed to get serialized.
    pub fn from_value<const N: usize>(
        mut backing: B,
        value: T,
    ) -> Result<BackedArchive<B, T>, SerializeError<T>>
    where
        T: Serialize<DefaultSerializer>,
    {
        // Serialize our value so we can shove it into the backing.
        let mut serializer = DefaultSerializer::default();
        let _ = serializer
            .serialize_value(&value)
            .map_err(|e| SerializeError::FailedToSerialize(e.to_string()))?;

        let src_buf = serializer.into_serializer().into_inner();

        // Now we have to write the serialized version to the backing store.  For obvious reasons,
        // the backing store needs to be able to hold the entire serialized representation, so we
        // check for that.  As well, instead of using `archived_root_mut`, we use
        // `archived_value_mut`, because this lets us relax need the backing store to be sized
        // _identically_ to the serialized size.
        let dst_buf = backing.as_mut();
        if dst_buf.len() < src_buf.len() {
            return Err(SerializeError::BackingStoreTooSmall(value, src_buf.len()));
        }

        dst_buf[..src_buf.len()].copy_from_slice(&src_buf);

        Ok(Self {
            backing,
            _archive: PhantomData,
        })
    }

    /// Runs a closure with immutable access to the archived value.
    pub fn with<F, V>(&self, f: F) -> V
    where
        F: FnOnce(&T::Archived) -> V,
    {
        let archive = unsafe { archived_value::<T>(self.backing.as_ref(), 0) };
        f(archive)
    }

    /// Runs a closure with mutable access to the archived value.
    pub fn with_mut<F, V>(&mut self, f: F) -> V
    where
        F: FnOnce(&mut T::Archived) -> V,
    {
        let pinned = Pin::new(self.backing.as_mut());
        let archive = unsafe { archived_value_mut::<T>(pinned, 0) };
        f(archive.get_mut())
    }
}