Skip to main content

platform_mem/
alloc.rs

1use {
2    crate::{
3        Error::{AllocError, CapacityOverflow},
4        RawMem, RawPlace, Result, utils,
5    },
6    allocator_api2::alloc::Allocator,
7    std::{
8        alloc::Layout,
9        fmt::{self, Debug, Formatter},
10        mem::{self, MaybeUninit},
11        ptr,
12    },
13};
14
15/// Allocator-backed memory storage.
16///
17/// Wraps any [`Allocator`] to provide a [`RawMem`] implementation.
18/// Memory is freed and elements are dropped when the `Alloc` is dropped.
19pub struct Alloc<T, A: Allocator> {
20    buf: RawPlace<T>,
21    alloc: A,
22}
23
24impl<T, A: Allocator> Alloc<T, A> {
25    /// Construct a new empty `Alloc<T, A>`.
26    /// It will not allocate until [growing][RawMem::grow].
27    /// ```
28    /// # use platform_mem::Global;
29    /// // It's able to be static
30    /// static ALLOC: Global<()> = Global::new();
31    /// ```
32    pub const fn new(alloc: A) -> Self {
33        Self { buf: RawPlace::dangling(), alloc }
34    }
35}
36
37impl<T, A: Allocator> RawMem for Alloc<T, A> {
38    type Item = T;
39
40    fn allocated(&self) -> &[Self::Item] {
41        unsafe { self.buf.as_slice() }
42    }
43
44    fn allocated_mut(&mut self) -> &mut [Self::Item] {
45        unsafe { self.buf.as_slice_mut() }
46    }
47
48    unsafe fn grow(
49        &mut self,
50        addition: usize,
51        fill: impl FnOnce(usize, (&mut [T], &mut [MaybeUninit<T>])),
52    ) -> Result<&mut [T]> {
53        unsafe {
54            let cap = self.buf.cap().checked_add(addition).ok_or(CapacityOverflow)?;
55            let new_layout = Layout::array::<T>(cap).map_err(|_| CapacityOverflow)?;
56
57            let ptr = if let Some((ptr, old_layout)) = self.buf.current_memory() {
58                self.alloc.grow(ptr, old_layout, new_layout)
59            } else {
60                self.alloc.allocate(new_layout)
61            }
62            .map_err(|_| AllocError { layout: new_layout, non_exhaustive: () })?
63            .cast();
64
65            // allocator always provide uninit memory
66            Ok(self.buf.handle_fill((ptr, cap), 0, fill))
67        }
68    }
69
70    fn shrink(&mut self, cap: usize) -> Result<()> {
71        let cap = self.buf.cap().checked_sub(cap).expect("Tried to shrink to a larger capacity");
72
73        let Some((ptr, layout)) = self.buf.current_memory() else {
74            return Ok(());
75        };
76        self.buf.shrink_to(cap);
77
78        let ptr = unsafe {
79            // `Layout::array` cannot overflow here because it would have
80            // overflowed earlier when capacity was larger.
81            let new_size = mem::size_of::<T>().unchecked_mul(cap);
82            let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
83            self.alloc
84                .shrink(ptr, layout, new_layout)
85                .map_err(|_| AllocError { layout: new_layout, non_exhaustive: () })?
86        };
87
88        #[allow(clippy::unit_arg)] // it is allows shortest return `Ok(())`
89        Ok({
90            self.buf.set_ptr(ptr);
91        })
92    }
93}
94
95impl<T, A: Allocator> Drop for Alloc<T, A> {
96    fn drop(&mut self) {
97        unsafe {
98            if let Some((ptr, layout)) = self.buf.current_memory() {
99                ptr::drop_in_place(self.buf.as_slice_mut());
100                self.alloc.deallocate(ptr, layout);
101            }
102        }
103    }
104}
105
106impl<T, A: Allocator + Debug> Debug for Alloc<T, A> {
107    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
108        utils::debug_mem(f, &self.buf, "Alloc")?.field("alloc", &self.alloc).finish()
109    }
110}