rkvy_size_counter.rs

use rkyv::{
    Serialize,
    rancor::Strategy,
    ser::{Positional, Serializer, WriterExt, allocator::ArenaHandle},
};

/// Dummy serializer which just tracks the position in the buffer
mod counter {
    use std::marker::PhantomData;

    use rkyv::{
        rancor::{Fallible, Strategy},
        ser::{Positional, Serializer, Writer, allocator::ArenaHandle},
    };

    pub struct PositionTracker<E> {
        pos: usize,
        _phantom: PhantomData<fn() -> E>,
    }

    impl<E> PositionTracker<E> {
        pub fn new(pos: usize) -> Self {
            Self {
                pos,
                _phantom: PhantomData,
            }
        }
    }
    impl<E> Fallible for PositionTracker<E> {
        type Error = E;
    }
    impl<E> Positional for PositionTracker<E> {
        fn pos(&self) -> usize {
            self.pos
        }
    }
    impl<E> Writer for PositionTracker<E> {
        fn write(&mut self, bytes: &[u8]) -> Result<(), E> {
            self.pos = self.pos.checked_add(bytes.len()).unwrap();
            Ok(())
        }
    }

    pub type Counter<'a, S, E> = Strategy<Serializer<PositionTracker<E>, ArenaHandle<'a>, S>, E>;
}
pub use counter::{Counter, PositionTracker};

/// Find the archived size and required alignment for a value written at `start`
/// The size includes any padding required to align the starting address
pub fn count_size<'a, T, S: Default, E>(
    value: &T,
    allocator: ArenaHandle<'a>,
    start: usize,
) -> Result<(usize, ArenaHandle<'a>), E>
where
    for<'b> T: Serialize<Counter<'b, S, E>>,
{
    let mut serializer = Serializer::new(PositionTracker::new(start), allocator, S::default());

    let serializer_wrapped = Strategy::<_, E>::wrap(&mut serializer);

    // Initial align
    serializer_wrapped.align_for::<T::Archived>()?;

    let resolver = value.serialize(serializer_wrapped)?;

    // Do we need to align again here?
    serializer_wrapped.align_for::<T::Archived>()?;

    unsafe { serializer_wrapped.resolve_aligned(value, resolver)? };

    let (writer, allocator, _) = serializer.into_raw_parts();
    Ok((writer.pos().checked_sub(start).unwrap(), allocator))
}

/// Find the archived size and required alignement for a value
pub fn count_size_and_align<'a, T, S: Default, E>(
    value: &T,
    allocator: ArenaHandle<'a>,
) -> Result<(usize, usize, ArenaHandle<'a>), E>
where
    for<'b> T: Serialize<Counter<'b, S, E>>,
{
    let (size_aligned, allocator) = count_size(value, allocator, 0)?;
    let (size_misaligned, allocator) = count_size(value, allocator, 1)?;

    // the misaligned result will be `align - 1` larger than the aligned result
    // as `(1 + align - 1) % align == 0`.
    let alignment = size_misaligned + 1 - size_aligned;

    Ok((size_aligned, alignment, allocator))
}

#[cfg(test)]
mod test {
    use std::fmt::Debug;

    use rand::{
        Rng,
        distr::{Distribution, StandardUniform},
    };

    use rkyv::{
        Serialize,
        ser::{allocator::ArenaHandle, sharing::Share},
    };

    use super::{Counter, count_size, count_size_and_align};

    fn test_size_counter<'a, T>(mut allocator: ArenaHandle<'a>, value: T)
    where
        for<'b> T: Serialize<Counter<'b, Share, rkyv::rancor::Error>> + Debug,
    {
        let size_aligned;
        let alignment;
        (size_aligned, alignment, allocator) = count_size_and_align(&value, allocator).unwrap();

        for start in 0..=4 * alignment {
            let offset = start % alignment;
            let expected_padding = if offset == 0 { 0 } else { alignment - offset };

            let padded_size;
            (padded_size, allocator) = count_size(&value, allocator, start).unwrap();

            assert_eq!(padded_size, size_aligned + expected_padding);
        }

        // This assert fails. So rkyv does some more aligning internally...
        assert_eq!(alignment, std::mem::align_of::<T::Archived>());
    }

    fn test_size_counter_for_primitive<T>(
        arena: &mut rkyv::ser::allocator::Arena,
        rng: &mut impl Rng,
    ) where
        StandardUniform: Distribution<T>,
        for<'a> T: Serialize<Counter<'a, Share, rkyv::rancor::Error>> + Debug,
    {
        test_size_counter(arena.acquire(), rng.random::<T>());

        for n in 0..20 {
            let values: Vec<T> = (0..n).map(|_| rng.random::<T>()).collect();
            test_size_counter(arena.acquire(), values);
        }
    }

    #[test]
    fn count_sizes() {
        let mut arena = rkyv::ser::allocator::Arena::new();

        let mut rng = rand::rng();

        test_size_counter_for_primitive::<bool>(&mut arena, &mut rng);
        test_size_counter_for_primitive::<u8>(&mut arena, &mut rng);
        test_size_counter_for_primitive::<u16>(&mut arena, &mut rng);
        test_size_counter_for_primitive::<u32>(&mut arena, &mut rng);
        test_size_counter_for_primitive::<u64>(&mut arena, &mut rng);
        test_size_counter_for_primitive::<u128>(&mut arena, &mut rng);
    }
}