diff options
author | Miguel Ojeda <ojeda@kernel.org> | 2022-11-10 19:41:24 +0300 |
---|---|---|
committer | Miguel Ojeda <ojeda@kernel.org> | 2022-12-04 03:59:15 +0300 |
commit | 51d3a25ab3a4f1dd701b17f7340e36de8600e41e (patch) | |
tree | b6634dc784239392b89db497403a7ac63ad1b956 /rust/alloc | |
parent | 25d176a4fad5841de1f85d2a4d69ce87416abc26 (diff) | |
download | linux-51d3a25ab3a4f1dd701b17f7340e36de8600e41e.tar.xz |
rust: alloc: add `RawVec::try_with_capacity_in()` constructor
Add the `RawVec::try_with_capacity_in()` constructor as the fallible
version of `RawVec::with_capacity_in()`.
The implementation follows the original.
The infallible constructor is implemented in terms of the private
`RawVec::allocate_in()` constructor, thus also add the private
`RawVec::try_allocate_in()` constructor following the other.
It will be used to implement `Vec::try_with_capacity{,_in}()` in
the next patch.
Reviewed-by: Gary Guo <gary@garyguo.net>
Signed-off-by: Miguel Ojeda <ojeda@kernel.org>
Diffstat (limited to 'rust/alloc')
-rw-r--r-- | rust/alloc/raw_vec.rs | 34 |
1 files changed, 33 insertions, 1 deletions
diff --git a/rust/alloc/raw_vec.rs b/rust/alloc/raw_vec.rs index daf5f2da7168..c342f3843972 100644 --- a/rust/alloc/raw_vec.rs +++ b/rust/alloc/raw_vec.rs @@ -20,11 +20,11 @@ use crate::collections::TryReserveErrorKind::*; #[cfg(test)] mod tests; -#[cfg(not(no_global_oom_handling))] enum AllocInit { /// The contents of the new memory are uninitialized. Uninitialized, /// The new memory is guaranteed to be zeroed. + #[allow(dead_code)] Zeroed, } @@ -133,6 +133,14 @@ impl<T, A: Allocator> RawVec<T, A> { Self::allocate_in(capacity, AllocInit::Uninitialized, alloc) } + /// Like `try_with_capacity`, but parameterized over the choice of + /// allocator for the returned `RawVec`. + #[allow(dead_code)] + #[inline] + pub fn try_with_capacity_in(capacity: usize, alloc: A) -> Result<Self, TryReserveError> { + Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc) + } + /// Like `with_capacity_zeroed`, but parameterized over the choice /// of allocator for the returned `RawVec`. #[cfg(not(no_global_oom_handling))] @@ -203,6 +211,30 @@ impl<T, A: Allocator> RawVec<T, A> { } } + fn try_allocate_in(capacity: usize, init: AllocInit, alloc: A) -> Result<Self, TryReserveError> { + // Don't allocate here because `Drop` will not deallocate when `capacity` is 0. + if mem::size_of::<T>() == 0 || capacity == 0 { + return Ok(Self::new_in(alloc)); + } + + let layout = Layout::array::<T>(capacity).map_err(|_| CapacityOverflow)?; + alloc_guard(layout.size())?; + let result = match init { + AllocInit::Uninitialized => alloc.allocate(layout), + AllocInit::Zeroed => alloc.allocate_zeroed(layout), + }; + let ptr = result.map_err(|_| AllocError { layout, non_exhaustive: () })?; + + // Allocators currently return a `NonNull<[u8]>` whose length + // matches the size requested. If that ever changes, the capacity + // here should change to `ptr.len() / mem::size_of::<T>()`. + Ok(Self { + ptr: unsafe { Unique::new_unchecked(ptr.cast().as_ptr()) }, + cap: capacity, + alloc, + }) + } + /// Reconstitutes a `RawVec` from a pointer, capacity, and allocator. /// /// # Safety |