musli/alloc/
system.rs

1use core::alloc::Layout;
2use core::cmp;
3use core::mem::{align_of, size_of};
4use core::ptr::NonNull;
5
6use rust_alloc::alloc;
7
8use super::{Allocator, RawVec};
9
10/// System buffer that can be used in combination with an [`Allocator`].
11///
12/// This uses the [`System`] allocator.
13///
14/// [`System` allocator]: https://doc.rust-lang.org/std/alloc/struct.System.html
15///
16/// # Examples
17///
18/// ```
19/// use musli::alloc::{System, Vec};
20///
21/// let alloc = System::new();
22///
23/// let mut buf1 = Vec::new_in(&alloc);
24/// let mut buf2 = Vec::new_in(&alloc);
25//
26/// assert!(buf1.write(b"Hello, "));
27/// assert!(buf2.write(b"world!"));
28///
29/// assert_eq!(buf1.as_slice(), b"Hello, ");
30/// assert_eq!(buf2.as_slice(), b"world!");
31///
32/// buf1.extend(buf2);
33/// assert_eq!(buf1.as_slice(), b"Hello, world!");
34/// ```
35#[non_exhaustive]
36pub struct System;
37
38impl System {
39    /// Construct a new allocator.
40    #[inline]
41    pub const fn new() -> Self {
42        Self
43    }
44}
45
46impl Default for System {
47    #[inline]
48    fn default() -> Self {
49        Self::new()
50    }
51}
52
53impl Allocator for System {
54    type RawVec<'this, T> = SystemBuf<T> where Self: 'this, T: 'this;
55
56    #[inline]
57    fn new_raw_vec<'a, T>(&'a self) -> Self::RawVec<'a, T>
58    where
59        T: 'a,
60    {
61        SystemBuf::DANGLING
62    }
63}
64
65/// A vector-backed allocation.
66pub struct SystemBuf<T> {
67    /// Pointer to the allocated region.
68    data: NonNull<T>,
69    /// The size in number of `T` elements in the region.
70    size: usize,
71}
72
73impl<T> RawVec<T> for SystemBuf<T> {
74    #[inline]
75    fn resize(&mut self, len: usize, additional: usize) -> bool {
76        if additional == 0 || size_of::<T>() == 0 {
77            return true;
78        }
79
80        self.reserve(len, additional)
81    }
82
83    #[inline]
84    fn as_ptr(&self) -> *const T {
85        self.data.as_ptr().cast_const().cast()
86    }
87
88    #[inline]
89    fn as_mut_ptr(&mut self) -> *mut T {
90        self.data.as_ptr().cast()
91    }
92
93    #[inline]
94    fn try_merge<B>(&mut self, _: usize, other: B, _: usize) -> Result<(), B>
95    where
96        B: RawVec<T>,
97    {
98        Err(other)
99    }
100}
101
102impl<T> SystemBuf<T> {
103    const MIN_NON_ZERO_CAP: usize = if size_of::<T>() == 1 {
104        8
105    } else if size_of::<T>() <= 1024 {
106        4
107    } else {
108        1
109    };
110
111    const DANGLING: Self = Self {
112        data: NonNull::dangling(),
113        size: 0,
114    };
115
116    /// Reallocate the region to the given capacity.
117    ///
118    /// # Safety
119    ///
120    /// The caller must ensure that the new capacity is valid per [`Layout`].
121    #[must_use = "allocating is fallible and must be checked"]
122    fn realloc(&mut self, new_layout: Layout) -> bool {
123        unsafe {
124            let data = {
125                if self.size > 0 {
126                    let old_layout = Layout::from_size_align_unchecked(
127                        self.size.wrapping_mul(size_of::<T>()),
128                        align_of::<T>(),
129                    );
130
131                    alloc::realloc(self.data.as_ptr().cast(), old_layout, new_layout.size())
132                } else {
133                    alloc::alloc(new_layout)
134                }
135            };
136
137            if data.is_null() {
138                return false;
139            }
140
141            self.data = NonNull::new_unchecked(data).cast();
142        }
143
144        true
145    }
146
147    #[must_use = "allocating is fallible and must be checked"]
148    fn reserve(&mut self, len: usize, additional: usize) -> bool {
149        debug_assert_ne!(size_of::<T>(), 0, "ZSTs should not get here");
150
151        let Some(required_cap) = len.checked_add(additional) else {
152            return false;
153        };
154
155        if self.size >= required_cap {
156            return true;
157        }
158
159        let cap = cmp::max(self.size * 2, required_cap);
160        let cap = cmp::max(Self::MIN_NON_ZERO_CAP, cap);
161
162        let Ok(new_layout) = Layout::array::<T>(cap) else {
163            return false;
164        };
165
166        if !self.realloc(new_layout) {
167            return false;
168        }
169
170        self.size = cap;
171        true
172    }
173
174    fn free(&mut self) {
175        if self.size > 0 {
176            // SAFETY: Layout assumptions are correctly encoded in the type as
177            // it was being allocated or grown.
178            unsafe {
179                let layout =
180                    Layout::from_size_align_unchecked(self.size * size_of::<T>(), align_of::<T>());
181                alloc::dealloc(self.data.as_ptr().cast(), layout);
182                self.data = NonNull::dangling();
183                self.size = 0;
184            }
185        }
186    }
187}
188
189impl<T> Drop for SystemBuf<T> {
190    fn drop(&mut self) {
191        self.free();
192    }
193}