SmallArchivedVec

lib.rs

use rkyv::rel_ptr::RelPtr;
use rkyv::ser::{ScratchSpace, Serializer};
use rkyv::with::{ArchiveWith, DeserializeWith, SerializeWith};
use rkyv::{
    out_field, Archive, Archived, Deserialize, DeserializeUnsized, Fallible, Serialize,
    SerializeUnsized,
};
use std::alloc;
use std::slice;

/// Essentially a minorly tweaked copy of ArchivedVec
pub struct SmallArchivedVec<T> {
    ptr: RelPtr<T, OffsetType>,
    len: u8,
}

/// Large enough to contain a u8
type OffsetType = Archived<i16>;

pub struct SmallVecResolver {
    pos: usize,
}

impl<T> SmallArchivedVec<T> {
    pub fn as_ptr(&self) -> *const T {
        self.ptr.as_ptr()
    }

    pub fn as_slice(&self) -> &[T] {
        unsafe { slice::from_raw_parts(self.as_ptr(), self.len()) }
    }

    pub fn len(&self) -> usize {
        self.len as usize
    }

    pub unsafe fn resolve_from_len(
        len: u8,
        pos: usize,
        resolver: SmallVecResolver,
        out: *mut Self,
    ) {
        let (fp, fo) = out_field!(out.ptr);
        RelPtr::emplace(pos + fp, resolver.pos, fo);
        let (fp, fo) = out_field!(out.len);
        u8::resolve(&len, pos + fp, (), fo);
    }

    pub fn serialize_from_slice<U: Serialize<S, Archived = T>, S: Serializer + ?Sized>(
        slice: &[U],
        serializer: &mut S,
    ) -> Result<SmallVecResolver, S::Error>
    where
        [U]: SerializeUnsized<S>,
    {
        Ok(SmallVecResolver {
            pos: slice.serialize_unsized(serializer)?,
        })
    }
}

/// `With`-wrapper to archive `Vec<T>` up to len i16::MAX more compactly
pub struct Small;

impl<T, U: Archive<Archived = T>> ArchiveWith<Vec<U>> for Small {
    type Archived = SmallArchivedVec<T>;
    type Resolver = SmallVecResolver;

    unsafe fn resolve_with(
        vec: &Vec<U>,
        pos: usize,
        resolver: Self::Resolver,
        out: *mut Self::Archived,
    ) {
        SmallArchivedVec::resolve_from_len(vec.len() as u8, pos, resolver, out);
    }
}

impl<S, T, U> SerializeWith<Vec<U>, S> for Small
where
    S: Fallible + Serializer + ScratchSpace + ?Sized,
    U: Archive<Archived = T> + Serialize<S>,
{
    fn serialize_with(vec: &Vec<U>, serializer: &mut S) -> Result<Self::Resolver, S::Error> {
        SmallArchivedVec::serialize_from_slice(vec, serializer)
    }
}

impl<D, T, U> Deserialize<Vec<U>, D> for SmallArchivedVec<T>
where
    D: Fallible + ?Sized,
    U: Archive<Archived = T>,
    T: Deserialize<U, D>,
{
    fn deserialize(&self, deserializer: &mut D) -> Result<Vec<U>, D::Error> {
        unsafe {
            let data_address = DeserializeUnsized::<[U], D>::deserialize_unsized(
                self.as_slice(),
                deserializer,
                |layout| alloc::alloc(layout),
            )?;

            let metadata =
                DeserializeUnsized::<[U], D>::deserialize_metadata(self.as_slice(), deserializer)?;

            let ptr = ptr_meta::from_raw_parts_mut(data_address, metadata);

            Ok(Box::<[U]>::from_raw(ptr).into())
        }
    }
}

impl<D, T, U> DeserializeWith<SmallArchivedVec<T>, Vec<U>, D> for Small
where
    D: Fallible + ?Sized,
    U: Archive<Archived = T>,
    T: Deserialize<U, D>,
{
    fn deserialize_with(
        archived: &SmallArchivedVec<T>,
        deserializer: &mut D,
    ) -> Result<Vec<U>, D::Error> {
        archived.deserialize(deserializer)
    }
}