rkyv_test/ser/serializers/
core.rs

1use crate::{
2    ser::{ScratchSpace, Serializer},
3    Fallible,
4};
5use core::{
6    alloc::Layout,
7    fmt,
8    ops::DerefMut,
9    ptr::{copy_nonoverlapping, NonNull},
10};
11
12/// The error type returned by an [`BufferSerializer`].
13#[derive(Debug)]
14pub enum BufferSerializerError {
15    /// Writing has overflowed the internal buffer.
16    Overflow {
17        /// The position of the serializer
18        pos: usize,
19        /// The number of bytes needed
20        bytes_needed: usize,
21        /// The total length of the archive
22        archive_len: usize,
23    },
24}
25
26impl fmt::Display for BufferSerializerError {
27    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
28        match self {
29            Self::Overflow {
30                pos,
31                bytes_needed,
32                archive_len,
33            } => write!(
34                f,
35                "writing has overflowed the serializer buffer: pos {}, needed {}, total length {}",
36                pos, bytes_needed, archive_len
37            ),
38        }
39    }
40}
41
42#[cfg(feature = "std")]
43const _: () = {
44    use std::error::Error;
45
46    impl Error for BufferSerializerError {}
47};
48
49/// Wraps a byte buffer and equips it with [`Serializer`].
50///
51/// Common uses include archiving in `#![no_std]` environments and archiving small objects without
52/// allocating.
53///
54/// # Examples
55/// ```
56/// use rkyv::{
57///     archived_value,
58///     ser::{Serializer, serializers::BufferSerializer},
59///     AlignedBytes,
60///     AlignedVec,
61///     Archive,
62///     Archived,
63///     Serialize,
64/// };
65///
66/// #[derive(Archive, Serialize)]
67/// enum Event {
68///     Spawn,
69///     Speak(String),
70///     Die,
71/// }
72///
73/// let mut serializer = BufferSerializer::new(AlignedBytes([0u8; 256]));
74/// let pos = serializer.serialize_value(&Event::Speak("Help me!".to_string()))
75///     .expect("failed to archive event");
76/// let buf = serializer.into_inner();
77/// let archived = unsafe { archived_value::<Event>(buf.as_ref(), pos) };
78/// if let Archived::<Event>::Speak(message) = archived {
79///     assert_eq!(message.as_str(), "Help me!");
80/// } else {
81///     panic!("archived event was of the wrong type");
82/// }
83/// ```
84#[derive(Debug)]
85pub struct BufferSerializer<T> {
86    inner: T,
87    pos: usize,
88}
89
90impl<T> BufferSerializer<T> {
91    /// Creates a new archive buffer from a byte buffer.
92    #[inline]
93    pub fn new(inner: T) -> Self {
94        Self::with_pos(inner, 0)
95    }
96
97    /// Creates a new archive buffer from a byte buffer. The buffer will start writing at the given
98    /// position, but the buffer must contain all bytes (otherwise the alignments of types may not
99    /// be correct).
100    #[inline]
101    pub fn with_pos(inner: T, pos: usize) -> Self {
102        Self { inner, pos }
103    }
104
105    /// Consumes the serializer and returns the underlying type.
106    #[inline]
107    pub fn into_inner(self) -> T {
108        self.inner
109    }
110}
111
112impl<T: Default> Default for BufferSerializer<T> {
113    #[inline]
114    fn default() -> Self {
115        Self::new(T::default())
116    }
117}
118
119impl<T> Fallible for BufferSerializer<T> {
120    type Error = BufferSerializerError;
121}
122
123impl<T: AsMut<[u8]>> Serializer for BufferSerializer<T> {
124    #[inline]
125    fn pos(&self) -> usize {
126        self.pos
127    }
128
129    fn write(&mut self, bytes: &[u8]) -> Result<(), Self::Error> {
130        let end_pos = self.pos + bytes.len();
131        let archive_len = self.inner.as_mut().len();
132        if end_pos > archive_len {
133            Err(BufferSerializerError::Overflow {
134                pos: self.pos,
135                bytes_needed: bytes.len(),
136                archive_len,
137            })
138        } else {
139            unsafe {
140                copy_nonoverlapping(
141                    bytes.as_ptr(),
142                    self.inner.as_mut().as_mut_ptr().add(self.pos),
143                    bytes.len(),
144                );
145            }
146            self.pos = end_pos;
147            Ok(())
148        }
149    }
150}
151
152/// Errors that can occur when using a fixed-size allocator.
153///
154/// Pairing a fixed-size allocator with a fallback allocator can help prevent running out of scratch
155/// space unexpectedly.
156#[derive(Debug)]
157pub enum FixedSizeScratchError {
158    /// The allocator ran out of scratch space.
159    OutOfScratch(Layout),
160    /// Scratch space was not popped in reverse order.
161    NotPoppedInReverseOrder {
162        /// The current position of the start of free memory
163        pos: usize,
164        /// The next position according to the erroneous pop
165        next_pos: usize,
166        /// The size of the memory according to the erroneous pop
167        next_size: usize,
168    },
169    /// The given allocation did not belong to the scratch allocator.
170    UnownedAllocation,
171}
172
173impl fmt::Display for FixedSizeScratchError {
174    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
175        match self {
176            Self::OutOfScratch(layout) => write!(
177                f,
178                "out of scratch: requested scratch space with size {} and align {}",
179                layout.size(),
180                layout.align()
181            ),
182            Self::NotPoppedInReverseOrder {
183                pos,
184                next_pos,
185                next_size,
186            } => write!(
187                f,
188                "scratch space was not popped in reverse order: pos {}, next pos {}, next size {}",
189                pos, next_pos, next_size
190            ),
191            Self::UnownedAllocation => write!(f, "unowned allocation"),
192        }
193    }
194}
195
196#[cfg(feature = "std")]
197impl std::error::Error for FixedSizeScratchError {}
198
199/// Scratch space that allocates within a buffer.
200#[derive(Debug)]
201pub struct BufferScratch<T> {
202    buffer: T,
203    pos: usize,
204}
205
206impl<T> BufferScratch<T> {
207    /// Creates a new buffer scratch allocator.
208    pub fn new(buffer: T) -> Self {
209        Self { buffer, pos: 0 }
210    }
211
212    /// Resets the scratch space to its initial state.
213    pub fn clear(&mut self) {
214        self.pos = 0;
215    }
216
217    /// Consumes the buffer scratch allocator, returning the underlying buffer.
218    pub fn into_inner(self) -> T {
219        self.buffer
220    }
221}
222
223impl<T: Default> Default for BufferScratch<T> {
224    fn default() -> Self {
225        Self::new(T::default())
226    }
227}
228
229impl<T> Fallible for BufferScratch<T> {
230    type Error = FixedSizeScratchError;
231}
232
233impl<T: DerefMut<Target = U>, U: AsMut<[u8]>> ScratchSpace for BufferScratch<T> {
234    #[inline]
235    unsafe fn push_scratch(&mut self, layout: Layout) -> Result<NonNull<[u8]>, Self::Error> {
236        let bytes = self.buffer.as_mut();
237
238        let start = bytes.as_ptr().add(self.pos);
239        let pad = match (start as usize) & (layout.align() - 1) {
240            0 => 0,
241            x => layout.align() - x,
242        };
243        if pad + layout.size() <= bytes.len() - self.pos {
244            self.pos += pad;
245            let result_slice = ptr_meta::from_raw_parts_mut(
246                bytes.as_mut_ptr().add(self.pos).cast(),
247                layout.size(),
248            );
249            let result = NonNull::new_unchecked(result_slice);
250            self.pos += layout.size();
251            Ok(result)
252        } else {
253            Err(FixedSizeScratchError::OutOfScratch(layout))
254        }
255    }
256
257    #[inline]
258    unsafe fn pop_scratch(&mut self, ptr: NonNull<u8>, layout: Layout) -> Result<(), Self::Error> {
259        let bytes = self.buffer.as_mut();
260
261        let ptr = ptr.as_ptr();
262        if ptr >= bytes.as_mut_ptr() && ptr < bytes.as_mut_ptr().add(bytes.len()) {
263            let next_pos = ptr.offset_from(bytes.as_ptr()) as usize;
264            if next_pos + layout.size() <= self.pos {
265                self.pos = next_pos;
266                Ok(())
267            } else {
268                Err(FixedSizeScratchError::NotPoppedInReverseOrder {
269                    pos: self.pos,
270                    next_pos,
271                    next_size: layout.size(),
272                })
273            }
274        } else {
275            Err(FixedSizeScratchError::UnownedAllocation)
276        }
277    }
278}
279
280/// Allocates scratch space with a main and backup scratch.
281#[derive(Debug)]
282pub struct FallbackScratch<M, F> {
283    main: M,
284    fallback: F,
285}
286
287impl<M, F> FallbackScratch<M, F> {
288    /// Creates fallback scratch from a main and backup scratch.
289    pub fn new(main: M, fallback: F) -> Self {
290        Self { main, fallback }
291    }
292}
293
294impl<M: Default, F: Default> Default for FallbackScratch<M, F> {
295    fn default() -> Self {
296        Self {
297            main: M::default(),
298            fallback: F::default(),
299        }
300    }
301}
302
303impl<M, F: Fallible> Fallible for FallbackScratch<M, F> {
304    type Error = F::Error;
305}
306
307impl<M: ScratchSpace, F: ScratchSpace> ScratchSpace for FallbackScratch<M, F> {
308    #[inline]
309    unsafe fn push_scratch(&mut self, layout: Layout) -> Result<NonNull<[u8]>, Self::Error> {
310        self.main
311            .push_scratch(layout)
312            .or_else(|_| self.fallback.push_scratch(layout))
313    }
314
315    #[inline]
316    unsafe fn pop_scratch(&mut self, ptr: NonNull<u8>, layout: Layout) -> Result<(), Self::Error> {
317        self.main
318            .pop_scratch(ptr, layout)
319            .or_else(|_| self.fallback.pop_scratch(ptr, layout))
320    }
321}
322
323/// A passthrough scratch space allocator that tracks scratch space usage.
324#[derive(Debug)]
325pub struct ScratchTracker<T> {
326    inner: T,
327    bytes_allocated: usize,
328    allocations: usize,
329    max_bytes_allocated: usize,
330    max_allocations: usize,
331    max_alignment: usize,
332}
333
334impl<T> ScratchTracker<T> {
335    /// Creates a new scratch tracker from the given inner scratch space.
336    pub fn new(inner: T) -> Self {
337        Self {
338            inner,
339            bytes_allocated: 0,
340            allocations: 0,
341            max_bytes_allocated: 0,
342            max_allocations: 0,
343            max_alignment: 1,
344        }
345    }
346
347    /// Returns the maximum number of bytes that were concurrently allocated during serialization.
348    pub fn max_bytes_allocated(&self) -> usize {
349        self.max_bytes_allocated
350    }
351
352    /// Returns the maximum number of concurrent allocations during serialization.
353    pub fn max_allocations(&self) -> usize {
354        self.max_allocations
355    }
356
357    /// Returns the maximum alignment of scratch space requested during serialization.
358    pub fn max_alignment(&self) -> usize {
359        self.max_alignment
360    }
361
362    /// Returns the minimum buffer size required to serialize the same data.
363    ///
364    /// This calculation takes into account packing efficiency for slab allocated scratch space. It
365    /// is not exact, and has an error bound of `max_allocations * (max_alignment - 1)` bytes. This
366    /// should be suitably small for most use cases.
367    pub fn min_buffer_size(&self) -> usize {
368        self.max_bytes_allocated + self.min_buffer_size_max_error()
369    }
370
371    /// Returns the maximum error term for the minimum buffer size calculation.
372    pub fn min_buffer_size_max_error(&self) -> usize {
373        self.max_allocations * (self.max_alignment - 1)
374    }
375}
376
377impl<T: Fallible> Fallible for ScratchTracker<T> {
378    type Error = T::Error;
379}
380
381impl<T: ScratchSpace> ScratchSpace for ScratchTracker<T> {
382    #[inline]
383    unsafe fn push_scratch(&mut self, layout: Layout) -> Result<NonNull<[u8]>, Self::Error> {
384        let result = self.inner.push_scratch(layout)?;
385
386        self.bytes_allocated += layout.size();
387        self.allocations += 1;
388        self.max_bytes_allocated = usize::max(self.bytes_allocated, self.max_bytes_allocated);
389        self.max_allocations = usize::max(self.allocations, self.max_allocations);
390        self.max_alignment = usize::max(self.max_alignment, layout.align());
391
392        Ok(result)
393    }
394
395    #[inline]
396    unsafe fn pop_scratch(&mut self, ptr: NonNull<u8>, layout: Layout) -> Result<(), Self::Error> {
397        self.inner.pop_scratch(ptr, layout)?;
398
399        self.bytes_allocated -= layout.size();
400        self.allocations -= 1;
401
402        Ok(())
403    }
404}
405
406impl<T> From<T> for ScratchTracker<T> {
407    fn from(inner: T) -> Self {
408        Self::new(inner)
409    }
410}