dyn_stack/
mem.rs

1use crate::stack_req::StackReq;
2use alloc::alloc::handle_alloc_error;
3use core::alloc::Layout;
4use core::mem::ManuallyDrop;
5use core::mem::MaybeUninit;
6use core::ptr::NonNull;
7
8use crate::alloc::*;
9extern crate alloc;
10
11impl core::fmt::Display for AllocError {
12    fn fmt(&self, fmt: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
13        fmt.write_str("memory allocation failed")
14    }
15}
16
17#[cfg(any(feature = "std", feature = "core-error"))]
18impl crate::Error for AllocError {}
19
20use super::*;
21
22#[inline]
23fn to_layout(req: StackReq) -> Result<Layout, AllocError> {
24    req.layout().ok().ok_or(AllocError)
25}
26
27#[cfg(feature = "alloc")]
28#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
29impl MemBuffer {
30    /// Allocate a memory buffer with sufficient storage for the given stack requirements, using the
31    /// global allocator.
32    ///
33    /// Calls [`alloc::alloc::handle_alloc_error`] in the case of failure.
34    ///
35    /// # Example
36    /// ```
37    /// use dyn_stack::{MemStack, StackReq, MemBuffer};
38    ///
39    /// let req = StackReq::new::<i32>(3);
40    /// let mut buf = MemBuffer::new(req);
41    /// let stack = MemStack::new(&mut buf);
42    ///
43    /// // use the stack
44    /// let (arr, _) = stack.make_with::<i32>(3, |i| i as i32);
45    /// ```
46    pub fn new(req: StackReq) -> Self {
47        Self::new_in(req, Global)
48    }
49
50    /// Allocate a memory buffer with sufficient storage for the given stack requirements, using the
51    /// global allocator, or an error if the allocation did not succeed.
52    ///
53    /// # Example
54    /// ```
55    /// use dyn_stack::{MemStack, StackReq, MemBuffer};
56    ///
57    /// let req = StackReq::new::<i32>(3);
58    /// let mut buf = MemBuffer::new(req);
59    /// let stack = MemStack::new(&mut buf);
60    ///
61    /// // use the stack
62    /// let (arr, _) = stack.make_with::<i32>(3, |i| i as i32);
63    /// ```
64    pub fn try_new(req: StackReq) -> Result<Self, AllocError> {
65        Self::try_new_in(req, Global)
66    }
67
68    /// Creates a `MemBuffer` from its raw components.
69    ///
70    /// # Safety
71    ///
72    /// The arguments to this function must have been acquired from a call to
73    /// [`MemBuffer::into_raw_parts`]
74    #[inline]
75    pub unsafe fn from_raw_parts(ptr: *mut u8, len: usize, align: usize) -> Self {
76        Self {
77            ptr: NonNull::new_unchecked(ptr),
78            len,
79            align,
80            alloc: Global,
81        }
82    }
83
84    /// Decomposes a `MemBuffer` into its raw components in this order: ptr, length and
85    /// alignment.
86    #[inline]
87    pub fn into_raw_parts(self) -> (*mut u8, usize, usize) {
88        let no_drop = ManuallyDrop::new(self);
89        (no_drop.ptr.as_ptr(), no_drop.len, no_drop.align)
90    }
91}
92
93#[cfg(feature = "alloc")]
94#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
95impl PodBuffer {
96    /// Allocate a memory buffer with sufficient storage for the given stack requirements, using the
97    /// global allocator.
98    ///
99    /// Calls [`alloc::alloc::handle_alloc_error`] in the case of failure.
100    ///
101    /// # Example
102    /// ```
103    /// use dyn_stack::{PodStack, StackReq, PodBuffer};
104    ///
105    /// let req = StackReq::new::<i32>(3);
106    /// let mut buf = PodBuffer::new(req);
107    /// let stack = PodStack::new(&mut buf);
108    ///
109    /// // use the stack
110    /// let (arr, _) = stack.make_with::<i32>(3, |i| i as i32);
111    /// ```
112    pub fn new(req: StackReq) -> Self {
113        Self::new_in(req, Global)
114    }
115
116    /// Allocate a memory buffer with sufficient storage for the given stack requirements, using the
117    /// global allocator, or an error if the allocation did not succeed.
118    ///
119    /// # Example
120    /// ```
121    /// use dyn_stack::{PodStack, StackReq, PodBuffer};
122    ///
123    /// let req = StackReq::new::<i32>(3);
124    /// let mut buf = PodBuffer::new(req);
125    /// let stack = PodStack::new(&mut buf);
126    ///
127    /// // use the stack
128    /// let (arr, _) = stack.make_with::<i32>(3, |i| i as i32);
129    /// ```
130    pub fn try_new(req: StackReq) -> Result<Self, AllocError> {
131        Self::try_new_in(req, Global)
132    }
133
134    /// Creates a `PodBuffer` from its raw components.
135    ///
136    /// # Safety
137    ///
138    /// The arguments to this function must have been acquired from a call to
139    /// [`PodBuffer::into_raw_parts`]
140    #[inline]
141    pub unsafe fn from_raw_parts(ptr: *mut u8, len: usize, align: usize) -> Self {
142        Self {
143            ptr: NonNull::new_unchecked(ptr),
144            len,
145            align,
146            alloc: Global,
147        }
148    }
149
150    /// Decomposes a `PodBuffer` into its raw components in this order: ptr, length and
151    /// alignment.
152    #[inline]
153    pub fn into_raw_parts(self) -> (*mut u8, usize, usize) {
154        let no_drop = ManuallyDrop::new(self);
155        (no_drop.ptr.as_ptr(), no_drop.len, no_drop.align)
156    }
157}
158
159#[cfg(feature = "alloc")]
160/// Buffer of uninitialized bytes to serve as workspace for dynamic arrays.
161pub struct MemBuffer<A: Allocator = Global> {
162    ptr: NonNull<u8>,
163    len: usize,
164    align: usize,
165    alloc: A,
166}
167
168#[cfg(feature = "alloc")]
169/// Buffer of initialized bytes to serve as workspace for dynamic arrays.
170pub struct PodBuffer<A: Allocator = Global> {
171    ptr: NonNull<u8>,
172    len: usize,
173    align: usize,
174    alloc: A,
175}
176
177#[cfg(not(feature = "alloc"))]
178/// Buffer of uninitialized bytes to serve as workspace for dynamic arrays.
179pub struct MemBuffer<A: Allocator> {
180    ptr: NonNull<u8>,
181    len: usize,
182    align: usize,
183    alloc: A,
184}
185
186#[cfg(not(feature = "alloc"))]
187/// Buffer of initialized bytes to serve as workspace for dynamic arrays.
188pub struct PodBuffer<A: Allocator> {
189    ptr: NonNull<u8>,
190    len: usize,
191    align: usize,
192    alloc: A,
193}
194
195unsafe impl<A: Allocator + Sync> Sync for MemBuffer<A> {}
196unsafe impl<A: Allocator + Send> Send for MemBuffer<A> {}
197
198unsafe impl<A: Allocator + Sync> Sync for PodBuffer<A> {}
199unsafe impl<A: Allocator + Send> Send for PodBuffer<A> {}
200
201impl<A: Allocator> Drop for MemBuffer<A> {
202    #[inline]
203    fn drop(&mut self) {
204        // SAFETY: this was initialized with std::alloc::alloc
205        unsafe {
206            self.alloc.deallocate(
207                self.ptr,
208                Layout::from_size_align_unchecked(self.len, self.align),
209            )
210        }
211    }
212}
213
214impl<A: Allocator> Drop for PodBuffer<A> {
215    #[inline]
216    fn drop(&mut self) {
217        // SAFETY: this was initialized with std::alloc::alloc
218        unsafe {
219            self.alloc.deallocate(
220                self.ptr,
221                Layout::from_size_align_unchecked(self.len, self.align),
222            )
223        }
224    }
225}
226
227impl<A: Allocator> PodBuffer<A> {
228    /// Allocate a memory buffer with sufficient storage for the given stack requirements, using the
229    /// provided allocator.
230    ///
231    /// Calls [`alloc::alloc::handle_alloc_error`] in the case of failure.
232    ///
233    /// # Example
234    /// ```
235    /// use dyn_stack::{PodStack, StackReq, PodBuffer};
236    /// use dyn_stack::alloc::Global;
237    ///
238    /// let req = StackReq::new::<i32>(3);
239    /// let mut buf = PodBuffer::new_in(req, Global);
240    /// let stack = PodStack::new(&mut buf);
241    ///
242    /// // use the stack
243    /// let (arr, _) = stack.make_with::<i32>(3, |i| i as i32);
244    /// ```
245    pub fn new_in(req: StackReq, alloc: A) -> Self {
246        Self::try_new_in(req, alloc).unwrap_or_else(|_| handle_alloc_error(to_layout(req).unwrap()))
247    }
248
249    /// Allocate a memory buffer with sufficient storage for the given stack requirements, using the
250    /// provided allocator, or an `AllocError` in the case of failure.
251    ///
252    /// # Example
253    /// ```
254    /// use dyn_stack::{PodStack, StackReq, PodBuffer};
255    /// use dyn_stack::alloc::Global;
256    ///
257    /// let req = StackReq::new::<i32>(3);
258    /// let mut buf = PodBuffer::new_in(req, Global);
259    /// let stack = PodStack::new(&mut buf);
260    ///
261    /// // use the stack
262    /// let (arr, _) = stack.make_with::<i32>(3, |i| i as i32);
263    /// ```
264    pub fn try_new_in(req: StackReq, alloc: A) -> Result<Self, AllocError> {
265        unsafe {
266            let ptr = &mut *(alloc
267                .allocate_zeroed(to_layout(req)?)
268                .map_err(|_| AllocError)?
269                .as_ptr() as *mut [MaybeUninit<u8>]);
270            #[cfg(debug_assertions)]
271            ptr.fill(MaybeUninit::new(0xCD));
272
273            let len = ptr.len();
274            let ptr = NonNull::new_unchecked(ptr.as_mut_ptr() as *mut u8);
275            Ok(PodBuffer {
276                alloc,
277                ptr,
278                len,
279                align: req.align_bytes(),
280            })
281        }
282    }
283
284    /// Creates a `PodBuffer` from its raw components.
285    ///
286    /// # Safety
287    ///
288    /// The arguments to this function must have been acquired from a call to
289    /// [`PodBuffer::into_raw_parts`]
290    #[inline]
291    pub unsafe fn from_raw_parts_in(ptr: *mut u8, len: usize, align: usize, alloc: A) -> Self {
292        Self {
293            ptr: NonNull::new_unchecked(ptr),
294            len,
295            align,
296            alloc,
297        }
298    }
299
300    /// Decomposes a `PodBuffer` into its raw components in this order: ptr, length and
301    /// alignment.
302    #[inline]
303    pub fn into_raw_parts_with_alloc(self) -> (*mut u8, usize, usize, A) {
304        let me = ManuallyDrop::new(self);
305        (me.ptr.as_ptr(), me.len, me.align, unsafe {
306            core::ptr::read(core::ptr::addr_of!(me.alloc))
307        })
308    }
309}
310
311impl<A: Allocator> MemBuffer<A> {
312    /// Allocate a memory buffer with sufficient storage for the given stack requirements, using the
313    /// provided allocator.
314    ///
315    /// Calls [`alloc::alloc::handle_alloc_error`] in the case of failure.
316    ///
317    /// # Example
318    /// ```
319    /// use dyn_stack::{MemStack, StackReq, MemBuffer};
320    /// use dyn_stack::alloc::Global;
321    ///
322    /// let req = StackReq::new::<i32>(3);
323    /// let mut buf = MemBuffer::new_in(req, Global);
324    /// let stack = MemStack::new(&mut buf);
325    ///
326    /// // use the stack
327    /// let (arr, _) = stack.make_with::<i32>(3, |i| i as i32);
328    /// ```
329    pub fn new_in(req: StackReq, alloc: A) -> Self {
330        Self::try_new_in(req, alloc).unwrap_or_else(|_| handle_alloc_error(to_layout(req).unwrap()))
331    }
332
333    /// Allocate a memory buffer with sufficient storage for the given stack requirements, using the
334    /// provided allocator, or an `AllocError` in the case of failure.
335    ///
336    /// # Example
337    /// ```
338    /// use dyn_stack::{MemStack, StackReq, MemBuffer};
339    /// use dyn_stack::alloc::Global;
340    ///
341    /// let req = StackReq::new::<i32>(3);
342    /// let mut buf = MemBuffer::new_in(req, Global);
343    /// let stack = MemStack::new(&mut buf);
344    ///
345    /// // use the stack
346    /// let (arr, _) = stack.make_with::<i32>(3, |i| i as i32);
347    /// ```
348    pub fn try_new_in(req: StackReq, alloc: A) -> Result<Self, AllocError> {
349        unsafe {
350            let ptr = &mut *(alloc
351                .allocate(to_layout(req)?)
352                .map_err(|_| AllocError)?
353                .as_ptr() as *mut [MaybeUninit<u8>]);
354            let len = ptr.len();
355            let ptr = NonNull::new_unchecked(ptr.as_mut_ptr() as *mut u8);
356            Ok(MemBuffer {
357                alloc,
358                ptr,
359                len,
360                align: req.align_bytes(),
361            })
362        }
363    }
364
365    /// Creates a `MemBuffer` from its raw components.
366    ///
367    /// # Safety
368    ///
369    /// The arguments to this function must have been acquired from a call to
370    /// [`MemBuffer::into_raw_parts`]
371    #[inline]
372    pub unsafe fn from_raw_parts_in(ptr: *mut u8, len: usize, align: usize, alloc: A) -> Self {
373        Self {
374            ptr: NonNull::new_unchecked(ptr),
375            len,
376            align,
377            alloc,
378        }
379    }
380
381    /// Decomposes a `MemBuffer` into its raw components in this order: ptr, length and
382    /// alignment.
383    #[inline]
384    pub fn into_raw_parts_with_alloc(self) -> (*mut u8, usize, usize, A) {
385        let me = ManuallyDrop::new(self);
386        (me.ptr.as_ptr(), me.len, me.align, unsafe {
387            core::ptr::read(core::ptr::addr_of!(me.alloc))
388        })
389    }
390}
391
392impl<A: Allocator> core::ops::Deref for MemBuffer<A> {
393    type Target = [MaybeUninit<u8>];
394
395    #[inline]
396    fn deref(&self) -> &Self::Target {
397        unsafe {
398            core::slice::from_raw_parts(self.ptr.as_ptr() as *const MaybeUninit<u8>, self.len)
399        }
400    }
401}
402
403impl<A: Allocator> core::ops::DerefMut for MemBuffer<A> {
404    #[inline]
405    fn deref_mut(&mut self) -> &mut Self::Target {
406        unsafe {
407            core::slice::from_raw_parts_mut(self.ptr.as_ptr() as *mut MaybeUninit<u8>, self.len)
408        }
409    }
410}
411
412impl<A: Allocator> core::ops::Deref for PodBuffer<A> {
413    type Target = [u8];
414
415    #[inline]
416    fn deref(&self) -> &Self::Target {
417        unsafe { core::slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
418    }
419}
420
421impl<A: Allocator> core::ops::DerefMut for PodBuffer<A> {
422    #[inline]
423    fn deref_mut(&mut self) -> &mut Self::Target {
424        unsafe { core::slice::from_raw_parts_mut(self.ptr.as_ptr(), self.len) }
425    }
426}
427
428/// Error during memory allocation.
429#[derive(Debug, Clone, Copy, PartialEq, Eq)]
430pub struct AllocError;
dyn_stack/mem.rs

dyn_stack/
mem.rs
