Skip to main content

o3/alloc/
mod.rs

1use std::{mem, ptr};
2
3pub use error::Error;
4pub use gap::Gap;
5
6use crate::alloc;
7
8pub mod array;
9mod error;
10mod gap;
11pub mod raw;
12pub mod tail;
13
14/// The allocator could not satisfy a validated layout.
15#[derive(Clone, Copy, Debug, PartialEq, Eq)]
16pub struct Exhausted;
17
18/// Capacity or layout arithmetic exceeded a representable bound.
19#[derive(Clone, Copy, Debug, PartialEq, Eq)]
20pub struct Overflow;
21
22/// Reserves one insertion before constructing its value.
23pub fn try_vacant<T>(values: &mut Vec<T>) -> Result<impl FnOnce(T) + '_, Error> {
24    use raw::VecExt as _;
25
26    VecExt::try_reserve(values, 1)?;
27    // SAFETY: successful reservation admits the exact retained next slot.
28    Ok(unsafe { values.vacancy() })
29}
30
31/// Fallible allocation for one boxed value.
32pub trait BoxExt<T>: Sized {
33    fn try_new(value: T) -> Result<Self, Exhausted>;
34}
35
36impl<T> BoxExt<T> for Box<T> {
37    fn try_new(value: T) -> Result<Self, Exhausted> {
38        Box::try_new(value).map_err(|_| Exhausted)
39    }
40}
41
42/// Fallible allocation for an uninitialized boxed slice.
43pub trait BoxUninitExt<T>: Sized {
44    fn try_uninit(len: usize) -> Result<Self, Error>;
45}
46
47impl<T> BoxUninitExt<T> for Box<[mem::MaybeUninit<T>]> {
48    fn try_uninit(len: usize) -> Result<Self, Error> {
49        let plan = crate::alloc::raw::Layout::<T>::try_layout(len)?;
50        let data = crate::alloc::raw::LayoutExt::try_allocate::<mem::MaybeUninit<T>>(plan)?;
51        let entries = ptr::slice_from_raw_parts_mut(data.as_ptr(), len);
52        // SAFETY: the allocation has this slice's exact layout; MaybeUninit
53        // carries the uninitialized state until the caller writes each entry.
54        Ok(unsafe { Box::from_raw(entries) })
55    }
56}
57
58/// Fallible allocation for an uninitialized boxed array of exact length.
59pub trait BoxArrayExt<T, const N: usize>: Sized {
60    fn try_uninit() -> Result<Self, Exhausted>;
61}
62
63impl<T, const N: usize> BoxArrayExt<T, N> for Box<[mem::MaybeUninit<T>; N]> {
64    fn try_uninit() -> Result<Self, Exhausted> {
65        let entries = Self::try_new_uninit().map_err(|_| Exhausted)?;
66        // SAFETY: every array element is MaybeUninit<T> and permits uninitialized storage.
67        Ok(unsafe { entries.assume_init() })
68    }
69}
70
71/// Fallible allocation and initialization for one boxed slice.
72pub trait BoxSliceExt<T>: Sized {
73    fn try_copy(source: &[T]) -> Result<Self, Exhausted>
74    where
75        T: Copy;
76
77    fn try_repeat(len: usize, value: T) -> Result<Self, Error>
78    where
79        T: Copy;
80
81    /// Allocates first, then constructs each default value.
82    fn try_default(len: usize) -> Result<Self, Error>
83    where
84        T: Default;
85
86    /// Allocates first, then calls the generator once per element.
87    fn try_init(len: usize, init: impl FnMut() -> T) -> Result<Self, Error>;
88
89    /// Allocates first, then initializes each index in ascending order.
90    fn try_from_fn(len: usize, init: impl FnMut(usize) -> T) -> Result<Self, Error>;
91}
92
93impl<T> BoxSliceExt<T> for Box<[T]> {
94    fn try_copy(source: &[T]) -> Result<Self, Exhausted>
95    where
96        T: Copy,
97    {
98        let mut entries = array::Plan::from_slice(source).try_allocate(())?;
99        if !source.is_empty() {
100            entries.write_copy_of_slice(source);
101        }
102        // SAFETY: the complete owned extent was initialized by the exact copy.
103        Ok(unsafe { entries.assume_init() }
104            .into_vec()
105            .into_boxed_slice())
106    }
107
108    fn try_repeat(len: usize, value: T) -> Result<Self, Error>
109    where
110        T: Copy,
111    {
112        Self::try_init(len, || value)
113    }
114
115    fn try_default(len: usize) -> Result<Self, Error>
116    where
117        T: Default,
118    {
119        Self::try_init(len, T::default)
120    }
121
122    fn try_init(len: usize, mut init: impl FnMut() -> T) -> Result<Self, Error> {
123        Self::try_from_fn(len, |_| init())
124    }
125
126    fn try_from_fn(len: usize, init: impl FnMut(usize) -> T) -> Result<Self, Error> {
127        let mut entries = Box::<[mem::MaybeUninit<T>]>::try_uninit(len)?;
128        crate::alloc::raw::init(&mut entries, init);
129        // SAFETY: the complete owned slice was initialized in index order.
130        Ok(unsafe { entries.assume_init() })
131    }
132}
133
134/// Fallible allocation and reservation for a vector.
135pub trait VecExt<T>: Sized {
136    fn try_new(cap: usize) -> Result<Self, Error>;
137
138    fn try_grow(&mut self, extra: usize) -> Result<(), Error>;
139
140    /// Reserves additional elements with geometric capacity growth.
141    fn try_reserve(&mut self, extra: usize) -> Result<(), Error>;
142}
143
144/// Fallible ownership of an already validated UTF-8 string.
145pub trait StrExt {
146    fn try_to_owned(&self) -> Result<String, Exhausted>;
147}
148
149impl StrExt for str {
150    fn try_to_owned(&self) -> Result<String, Exhausted> {
151        let bytes = Box::<[u8]>::try_copy(self.as_bytes())?.into_vec();
152        // SAFETY: the allocation contains an exact copy of this valid UTF-8 string.
153        Ok(unsafe { String::from_utf8_unchecked(bytes) })
154    }
155}
156
157impl<T> VecExt<T> for Vec<T> {
158    fn try_new(cap: usize) -> Result<Self, Error> {
159        Ok(array::Plan::try_new(cap)?.try_vec()?)
160    }
161
162    fn try_reserve(&mut self, extra: usize) -> Result<(), Error> {
163        alloc::raw::Layout::<T>::try_grow::<true>(self, extra)
164    }
165
166    fn try_grow(&mut self, extra: usize) -> Result<(), Error> {
167        alloc::raw::Layout::<T>::try_grow::<false>(self, extra)
168    }
169}