Skip to main content

rkyv/ser/allocator/
core.rs

1use core::{
2    alloc::Layout,
3    error::Error,
4    fmt,
5    marker::PhantomData,
6    mem::MaybeUninit,
7    ptr::{slice_from_raw_parts_mut, NonNull},
8};
9
10use rancor::{fail, Source};
11
12use crate::ser::Allocator;
13
14#[derive(Debug)]
15struct OutOfSpaceError {
16    layout: Layout,
17}
18
19impl fmt::Display for OutOfSpaceError {
20    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
21        write!(
22            f,
23            "not enough space to allocate request of size {} and align {}",
24            self.layout.size(),
25            self.layout.align()
26        )
27    }
28}
29
30impl Error for OutOfSpaceError {}
31
32/// An allocator that sub-allocates a fixed-size memory space.
33#[derive(Debug)]
34pub struct SubAllocator<'a> {
35    bytes: NonNull<u8>,
36    used: usize,
37    size: usize,
38    _phantom: PhantomData<&'a mut [MaybeUninit<u8>]>,
39}
40
41impl<'a> SubAllocator<'a> {
42    /// Creates an empty suballocator.
43    pub fn empty() -> Self {
44        Self {
45            bytes: NonNull::dangling(),
46            used: 0,
47            size: 0,
48            _phantom: PhantomData,
49        }
50    }
51
52    /// Creates a new sub-allocator from the given byte slice.
53    pub fn new(bytes: &'a mut [MaybeUninit<u8>]) -> Self {
54        Self {
55            bytes: unsafe { NonNull::new_unchecked(bytes.as_mut_ptr().cast()) },
56            used: 0,
57            size: bytes.len(),
58            _phantom: PhantomData,
59        }
60    }
61}
62
63unsafe impl<E> Allocator<E> for SubAllocator<'_>
64where
65    E: Source,
66{
67    unsafe fn push_alloc(
68        &mut self,
69        layout: Layout,
70    ) -> Result<NonNull<[u8]>, E> {
71        let pos = self.bytes.as_ptr() as usize + self.used;
72        let pad = 0usize.wrapping_sub(pos) % layout.align();
73        if pad + layout.size() <= self.size - self.used {
74            self.used += pad;
75        } else {
76            fail!(OutOfSpaceError { layout });
77        }
78
79        // SAFETY: `self.used` is always less than the length of the allocated
80        // block that `self.bytes` points to.
81        let ptr = unsafe { self.bytes.as_ptr().add(self.used) };
82        let slice_ptr = slice_from_raw_parts_mut(ptr, layout.size());
83        // SAFETY: `slice_ptr` is guaranteed not to be null because it is
84        // offset from `self.bytes` which is always non-null.
85        let result = unsafe { NonNull::new_unchecked(slice_ptr) };
86        self.used += layout.size();
87        Ok(result)
88    }
89
90    unsafe fn pop_alloc(
91        &mut self,
92        ptr: NonNull<u8>,
93        _: Layout,
94    ) -> Result<(), E> {
95        let bytes = self.bytes.as_ptr();
96        self.used = ptr.as_ptr() as usize - bytes as usize;
97
98        Ok(())
99    }
100}