Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Vec::try_with_capacity #120504

Merged
merged 3 commits into from
Mar 9, 2024
Merged

Vec::try_with_capacity #120504

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift click to select a range
f27a22c
try_with_capacity for RawVec
kornelski Jan 30, 2024
78fb977
try_with_capacity for Vec, VecDeque, String
kornelski Jan 30, 2024
784e6a1
Move capacity_overflow function to make ui tests change less
kornelski Jan 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Next Next commit
try_with_capacity for RawVec
  • Loading branch information
@kornelski
kornelski committed Mar 1, 2024
commit f27a22c24a51ba734e4eb25c499320d819ebe236
36 changes: 21 additions & 15 deletions library/alloc/src/raw_vec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 17,6 @@ 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,
Expand Down Expand Up @@ -93,6 92,8 @@ impl<T> RawVec<T, Global> {
/// zero-sized. Note that if `T` is zero-sized this means you will
/// *not* get a `RawVec` with the requested capacity.
///
/// Non-fallible version of `try_with_capacity`
///
/// # Panics
///
/// Panics if the requested capacity exceeds `isize::MAX` bytes.
Expand All @@ -104,7 105,7 @@ impl<T> RawVec<T, Global> {
#[must_use]
#[inline]
pub fn with_capacity(capacity: usize) -> Self {
Self::with_capacity_in(capacity, Global)
handle_reserve(Self::try_allocate_in(capacity, AllocInit::Uninitialized, Global))
}

/// Like `with_capacity`, but guarantees the buffer is zeroed.
Expand Down Expand Up @@ -142,15 143,15 @@ impl<T, A: Allocator> RawVec<T, A> {
#[cfg(not(no_global_oom_handling))]
#[inline]
pub fn with_capacity_in(capacity: usize, alloc: A) -> Self {
Self::allocate_in(capacity, AllocInit::Uninitialized, alloc)
handle_reserve(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))]
#[inline]
pub fn with_capacity_zeroed_in(capacity: usize, alloc: A) -> Self {
Self::allocate_in(capacity, AllocInit::Zeroed, alloc)
handle_reserve(Self::try_allocate_in(capacity, AllocInit::Zeroed, alloc))
}

/// Converts the entire buffer into `Box<[MaybeUninit<T>]>` with the specified `len`.
Expand Down Expand Up @@ -179,35 180,40 @@ impl<T, A: Allocator> RawVec<T, A> {
}
}

#[cfg(not(no_global_oom_handling))]
fn allocate_in(capacity: usize, init: AllocInit, alloc: A) -> Self {
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 T::IS_ZST || capacity == 0 {
Self::new_in(alloc)
Ok(Self::new_in(alloc))
} else {
// We avoid `unwrap_or_else` here because it bloats the amount of
// LLVM IR generated.
let layout = match Layout::array::<T>(capacity) {
Ok(layout) => layout,
Err(_) => capacity_overflow(),
Err(_) => return Err(CapacityOverflow.into()),
};
match alloc_guard(layout.size()) {
Ok(_) => {}
Err(_) => capacity_overflow(),

if let Err(err) = alloc_guard(layout.size()) {
return Err(err);
}

let result = match init {
AllocInit::Uninitialized => alloc.allocate(layout),
AllocInit::Zeroed => alloc.allocate_zeroed(layout),
};
let ptr = match result {
Ok(ptr) => ptr,
Err(_) => handle_alloc_error(layout),
Err(_) => return Err(AllocError { layout, non_exhaustive: () }.into()),
};

// 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>()`.
Self { ptr: Unique::from(ptr.cast()), cap: unsafe { Cap(capacity) }, alloc }
Ok(Self { ptr: Unique::from(ptr.cast()), cap: unsafe { Cap(capacity) }, alloc })
}
}

Expand Down Expand Up @@ -537,11 543,11 @@ unsafe impl<#[may_dangle] T, A: Allocator> Drop for RawVec<T, A> {
// Central function for reserve error handling.
#[cfg(not(no_global_oom_handling))]
#[inline]
fn handle_reserve(result: Result<(), TryReserveError>) {
fn handle_reserve<T>(result: Result<T, TryReserveError>) -> T {
match result.map_err(|e| e.kind()) {
Ok(res) => res,
Err(CapacityOverflow) => capacity_overflow(),
Err(AllocError { layout, .. }) => handle_alloc_error(layout),
Ok(()) => { /* yay */ }
}
}

Expand Down
7 changes: 4 additions & 3 deletions library/alloc/src/raw_vec/tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -105,13 105,14 @@ fn zst() {
let v: RawVec<ZST> = RawVec::with_capacity_in(100, Global);
zst_sanity(&v);

let v: RawVec<ZST> = RawVec::allocate_in(0, AllocInit::Uninitialized, Global);
let v: RawVec<ZST> = RawVec::try_allocate_in(0, AllocInit::Uninitialized, Global).unwrap();
zst_sanity(&v);

let v: RawVec<ZST> = RawVec::allocate_in(100, AllocInit::Uninitialized, Global);
let v: RawVec<ZST> = RawVec::try_allocate_in(100, AllocInit::Uninitialized, Global).unwrap();
zst_sanity(&v);

let mut v: RawVec<ZST> = RawVec::allocate_in(usize::MAX, AllocInit::Uninitialized, Global);
let mut v: RawVec<ZST> =
RawVec::try_allocate_in(usize::MAX, AllocInit::Uninitialized, Global).unwrap();
zst_sanity(&v);

// Check all these operations work as expected with zero-sized elements.
Expand Down