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};

use crate::align_up;

#[derive(Debug, PartialEq, Eq)]
pub struct ArchivedSizes {
    body_size: usize,
    central_padding: usize,
    resolver_size: usize,
}

impl ArchivedSizes {
    pub fn total_size(&self) -> usize {
        self.body_size + self.central_padding + self.resolver_size
    }
}

/// 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<(ArchivedSizes, 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);

    let resolver = value.serialize(serializer_wrapped)?;
    let body_end = serializer_wrapped.writer.pos();

    serializer_wrapped.align_for::<T::Archived>()?;
    let resolver_start = serializer_wrapped.writer.pos();

    unsafe { serializer_wrapped.resolve_aligned(value, resolver)? };
    let resolver_end = serializer_wrapped.writer.pos();

    let (_, allocator, _) = serializer.into_raw_parts();
    Ok((
        ArchivedSizes {
            resolver_size: resolver_end.checked_sub(resolver_start).unwrap(),
            body_size: body_end.checked_sub(start).unwrap(),
            central_padding: resolver_start.checked_sub(body_end).unwrap(),
        },
        allocator,
    ))
}

#[derive(Debug, Clone, Copy)]
pub struct SizePredictor {
    pub body_alignment: usize,
    pub resolver_alignment: usize,
    pub body_size: usize,       // body size when start = 0
    pub central_padding: usize, // central padding when start = 0
    pub resolver_size: usize,   // resolver size when start = 0
}

impl SizePredictor {
    /// Minimum alignment for the body to start with no leading padding
    pub fn start_alignment(&self) -> usize {
        self.body_alignment
    }

    pub fn predict_sizes(&self, start: usize) -> (usize, ArchivedSizes) {
        // body
        let body_off = align_up(start, self.body_alignment);
        let body_end = body_off + self.body_size;

        // resolver
        let resolver_off = align_up(body_end, self.resolver_alignment);
        let final_end = resolver_off + self.resolver_size;

        let body_padding = body_off - start;
        let central_padding = resolver_off - body_end;

        let sizes = ArchivedSizes {
            resolver_size: self.resolver_size,
            body_size: body_end - start,
            central_padding,
        };

        assert_eq!(
            sizes.total_size(),
            final_end - start,
            "\nself={self:#?}\nstart={start}\nbody_off={body_off}\nbody_end={body_end}\nresolver_off={resolver_off}\nfinal_end={final_end}\nbody_padding={body_padding}\ncentral_padding={central_padding}\nsizes={sizes:#?}"
        );

        (body_padding, sizes)
    }

    /// `aligned_sizes` is from start=0
    /// `misaligned_sizes` is from start=1
    /// `archived_align` is `align_of(T::Archived)`
    fn from_measurements(
        aligned_sizes: ArchivedSizes,
        misaligned_sizes: ArchivedSizes,
        archived_align: usize,
    ) -> Self {
        // Derive body alignment by the delta rule:
        let body_alignment = misaligned_sizes.body_size + 1 - aligned_sizes.body_size;

        assert!(
            body_alignment > 0,
            "body_alignment=0\naligned_sizes={aligned_sizes:#?}\nmisaligned_sizes={misaligned_sizes:#?}"
        );

        let resolver_alignment = archived_align;

        let body_size = aligned_sizes.body_size;
        let central_padding = aligned_sizes.central_padding;
        let resolver_size = aligned_sizes.resolver_size;

        assert_eq!(aligned_sizes.resolver_size, misaligned_sizes.resolver_size);

        SizePredictor {
            body_alignment,
            resolver_alignment,
            body_size,
            central_padding,
            resolver_size,
        }
    }
}

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

    let res = SizePredictor::from_measurements(
        aligned_sizes,
        misaligned_sizes,
        std::mem::align_of::<T::Archived>(),
    );

    Ok((res, allocator))
}