diff --git a/library/alloc/src/collections/vec_deque/mod.rs b/library/alloc/src/collections/vec_deque/mod.rs index c35bab5ef6657..41adc2e79dc74 100644 --- a/library/alloc/src/collections/vec_deque/mod.rs +++ b/library/alloc/src/collections/vec_deque/mod.rs @@ -559,6 +559,30 @@ impl VecDeque { pub fn with_capacity(capacity: usize) -> VecDeque { Self::with_capacity_in(capacity, Global) } + + /// Creates an empty deque with space for at least `capacity` elements. + /// + /// # Errors + /// + /// Returns an error if the capacity exceeds `isize::MAX` _bytes_, + /// or if the allocator reports allocation failure. + /// + /// # Examples + /// + /// ``` + /// # #![feature(try_with_capacity)] + /// # #[allow(unused)] + /// # fn example() -> Result<(), std::collections::TryReserveError> { + /// use std::collections::VecDeque; + /// + /// let deque: VecDeque = VecDeque::try_with_capacity(10)?; + /// # Ok(()) } + /// ``` + #[inline] + #[unstable(feature = "try_with_capacity", issue = "91913")] + pub fn try_with_capacity(capacity: usize) -> Result, TryReserveError> { + Ok(VecDeque { head: 0, len: 0, buf: RawVec::try_with_capacity_in(capacity, Global)? }) + } } impl VecDeque { diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs index 28695ade5bf55..ca504b05a9689 100644 --- a/library/alloc/src/lib.rs +++ b/library/alloc/src/lib.rs @@ -163,6 +163,7 @@ #![feature(trusted_len)] #![feature(trusted_random_access)] #![feature(try_trait_v2)] +#![feature(try_with_capacity)] #![feature(tuple_trait)] #![feature(unchecked_math)] #![feature(unicode_internals)] diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs index 4d9694d4bebad..ace73c0fdaa5e 100644 --- a/library/alloc/src/raw_vec.rs +++ b/library/alloc/src/raw_vec.rs @@ -17,10 +17,19 @@ use crate::collections::TryReserveErrorKind::*; #[cfg(test)] mod tests; +// One central function responsible for reporting capacity overflows. This'll +// ensure that the code generation related to these panics is minimal as there's +// only one location which panics rather than a bunch throughout the module. #[cfg(not(no_global_oom_handling))] +#[cfg_attr(not(feature = "panic_immediate_abort"), inline(never))] +fn capacity_overflow() -> ! { + panic!("capacity overflow"); +} + enum AllocInit { /// The contents of the new memory are uninitialized. Uninitialized, + #[cfg(not(no_global_oom_handling))] /// The new memory is guaranteed to be zeroed. Zeroed, } @@ -93,6 +102,8 @@ impl RawVec { /// 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. @@ -104,7 +115,7 @@ impl RawVec { #[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. @@ -142,7 +153,14 @@ impl RawVec { #[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 `try_with_capacity`, but parameterized over the choice of + /// allocator for the returned `RawVec`. + #[inline] + pub fn try_with_capacity_in(capacity: usize, alloc: A) -> Result { + Self::try_allocate_in(capacity, AllocInit::Uninitialized, alloc) } /// Like `with_capacity_zeroed`, but parameterized over the choice @@ -150,7 +168,7 @@ impl 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]>` with the specified `len`. @@ -179,35 +197,41 @@ impl RawVec { } } - #[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 { // 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::(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), + #[cfg(not(no_global_oom_handling))] 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::()`. - Self { ptr: Unique::from(ptr.cast()), cap: unsafe { Cap(capacity) }, alloc } + Ok(Self { ptr: Unique::from(ptr.cast()), cap: unsafe { Cap(capacity) }, alloc }) } } @@ -537,11 +561,11 @@ unsafe impl<#[may_dangle] T, A: Allocator> Drop for RawVec { // Central function for reserve error handling. #[cfg(not(no_global_oom_handling))] #[inline] -fn handle_reserve(result: Result<(), TryReserveError>) { +fn handle_reserve(result: Result) -> T { match result.map_err(|e| e.kind()) { + Ok(res) => res, Err(CapacityOverflow) => capacity_overflow(), Err(AllocError { layout, .. }) => handle_alloc_error(layout), - Ok(()) => { /* yay */ } } } @@ -561,12 +585,3 @@ fn alloc_guard(alloc_size: usize) -> Result<(), TryReserveError> { Ok(()) } } - -// One central function responsible for reporting capacity overflows. This'll -// ensure that the code generation related to these panics is minimal as there's -// only one location which panics rather than a bunch throughout the module. -#[cfg(not(no_global_oom_handling))] -#[cfg_attr(not(feature = "panic_immediate_abort"), inline(never))] -fn capacity_overflow() -> ! { - panic!("capacity overflow"); -} diff --git a/library/alloc/src/raw_vec/tests.rs b/library/alloc/src/raw_vec/tests.rs index f8cada01c0309..4194be530612d 100644 --- a/library/alloc/src/raw_vec/tests.rs +++ b/library/alloc/src/raw_vec/tests.rs @@ -105,13 +105,14 @@ fn zst() { let v: RawVec = RawVec::with_capacity_in(100, Global); zst_sanity(&v); - let v: RawVec = RawVec::allocate_in(0, AllocInit::Uninitialized, Global); + let v: RawVec = RawVec::try_allocate_in(0, AllocInit::Uninitialized, Global).unwrap(); zst_sanity(&v); - let v: RawVec = RawVec::allocate_in(100, AllocInit::Uninitialized, Global); + let v: RawVec = RawVec::try_allocate_in(100, AllocInit::Uninitialized, Global).unwrap(); zst_sanity(&v); - let mut v: RawVec = RawVec::allocate_in(usize::MAX, AllocInit::Uninitialized, Global); + let mut v: RawVec = + RawVec::try_allocate_in(usize::MAX, AllocInit::Uninitialized, Global).unwrap(); zst_sanity(&v); // Check all these operations work as expected with zero-sized elements. diff --git a/library/alloc/src/string.rs b/library/alloc/src/string.rs index 98ded7f6cdf5b..c4dcff1b1c49c 100644 --- a/library/alloc/src/string.rs +++ b/library/alloc/src/string.rs @@ -492,6 +492,19 @@ impl String { String { vec: Vec::with_capacity(capacity) } } + /// Creates a new empty `String` with at least the specified capacity. + /// + /// # Errors + /// + /// Returns [`Err`] if the capacity exceeds `isize::MAX` bytes, + /// or if the memory allocator reports failure. + /// + #[inline] + #[unstable(feature = "try_with_capacity", issue = "91913")] + pub fn try_with_capacity(capacity: usize) -> Result { + Ok(String { vec: Vec::try_with_capacity(capacity)? }) + } + // HACK(japaric): with cfg(test) the inherent `[T]::to_vec` method, which is // required for this method definition, is not available. Since we don't // require this method for testing purposes, I'll just stub it diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index c0e934b3b1fcb..f2f42e63d6b0b 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -481,6 +481,22 @@ impl Vec { Self::with_capacity_in(capacity, Global) } + /// Constructs a new, empty `Vec` with at least the specified capacity. + /// + /// The vector will be able to hold at least `capacity` elements without + /// reallocating. This method is allowed to allocate for more elements than + /// `capacity`. If `capacity` is 0, the vector will not allocate. + /// + /// # Errors + /// + /// Returns an error if the capacity exceeds `isize::MAX` _bytes_, + /// or if the allocator reports allocation failure. + #[inline] + #[unstable(feature = "try_with_capacity", issue = "91913")] + pub fn try_with_capacity(capacity: usize) -> Result { + Self::try_with_capacity_in(capacity, Global) + } + /// Creates a `Vec` directly from a pointer, a length, and a capacity. /// /// # Safety @@ -672,6 +688,24 @@ impl Vec { Vec { buf: RawVec::with_capacity_in(capacity, alloc), len: 0 } } + /// Constructs a new, empty `Vec` with at least the specified capacity + /// with the provided allocator. + /// + /// The vector will be able to hold at least `capacity` elements without + /// reallocating. This method is allowed to allocate for more elements than + /// `capacity`. If `capacity` is 0, the vector will not allocate. + /// + /// # Errors + /// + /// Returns an error if the capacity exceeds `isize::MAX` _bytes_, + /// or if the allocator reports allocation failure. + #[inline] + #[unstable(feature = "allocator_api", issue = "32838")] + // #[unstable(feature = "try_with_capacity", issue = "91913")] + pub fn try_with_capacity_in(capacity: usize, alloc: A) -> Result { + Ok(Vec { buf: RawVec::try_with_capacity_in(capacity, alloc)?, len: 0 }) + } + /// Creates a `Vec` directly from a pointer, a length, a capacity, /// and an allocator. /// diff --git a/library/alloc/tests/lib.rs b/library/alloc/tests/lib.rs index ed928994ad697..e8496989bcf12 100644 --- a/library/alloc/tests/lib.rs +++ b/library/alloc/tests/lib.rs @@ -20,6 +20,7 @@ #![feature(pattern)] #![feature(trusted_len)] #![feature(try_reserve_kind)] +#![feature(try_with_capacity)] #![feature(unboxed_closures)] #![feature(associated_type_bounds)] #![feature(binary_heap_into_iter_sorted)] diff --git a/library/alloc/tests/string.rs b/library/alloc/tests/string.rs index 711e4eef2e724..e20ceae87b0d5 100644 --- a/library/alloc/tests/string.rs +++ b/library/alloc/tests/string.rs @@ -723,6 +723,17 @@ fn test_reserve_exact() { assert!(s.capacity() >= 33) } +#[test] +#[cfg_attr(miri, ignore)] // Miri does not support signalling OOM +#[cfg_attr(target_os = "android", ignore)] // Android used in CI has a broken dlmalloc +fn test_try_with_capacity() { + let string = String::try_with_capacity(1000).unwrap(); + assert_eq!(0, string.len()); + assert!(string.capacity() >= 1000 && string.capacity() <= isize::MAX as usize); + + assert!(String::try_with_capacity(usize::MAX).is_err()); +} + #[test] #[cfg_attr(miri, ignore)] // Miri does not support signalling OOM #[cfg_attr(target_os = "android", ignore)] // Android used in CI has a broken dlmalloc diff --git a/library/alloc/tests/vec.rs b/library/alloc/tests/vec.rs index 15ee4d6520523..aa95b4e977081 100644 --- a/library/alloc/tests/vec.rs +++ b/library/alloc/tests/vec.rs @@ -1694,6 +1694,18 @@ fn test_reserve_exact() { assert!(v.capacity() >= 33) } +#[test] +#[cfg_attr(miri, ignore)] // Miri does not support signalling OOM +#[cfg_attr(target_os = "android", ignore)] // Android used in CI has a broken dlmalloc +fn test_try_with_capacity() { + let mut vec: Vec = Vec::try_with_capacity(5).unwrap(); + assert_eq!(0, vec.len()); + assert!(vec.capacity() >= 5 && vec.capacity() <= isize::MAX as usize / 4); + assert!(vec.spare_capacity_mut().len() >= 5); + + assert!(Vec::::try_with_capacity(isize::MAX as usize + 1).is_err()); +} + #[test] #[cfg_attr(miri, ignore)] // Miri does not support signalling OOM #[cfg_attr(target_os = "android", ignore)] // Android used in CI has a broken dlmalloc diff --git a/library/alloc/tests/vec_deque.rs b/library/alloc/tests/vec_deque.rs index eda2f8bb812b5..cea5de4dd5984 100644 --- a/library/alloc/tests/vec_deque.rs +++ b/library/alloc/tests/vec_deque.rs @@ -1182,6 +1182,17 @@ fn test_reserve_exact_2() { assert!(v.capacity() >= 33) } +#[test] +#[cfg_attr(miri, ignore)] // Miri does not support signalling OOM +#[cfg_attr(target_os = "android", ignore)] // Android used in CI has a broken dlmalloc +fn test_try_with_capacity() { + let vec: VecDeque = VecDeque::try_with_capacity(5).unwrap(); + assert_eq!(0, vec.len()); + assert!(vec.capacity() >= 5 && vec.capacity() <= isize::MAX as usize / 4); + + assert!(VecDeque::::try_with_capacity(isize::MAX as usize + 1).is_err()); +} + #[test] #[cfg_attr(miri, ignore)] // Miri does not support signalling OOM #[cfg_attr(target_os = "android", ignore)] // Android used in CI has a broken dlmalloc diff --git a/tests/codegen/vec-with-capacity.rs b/tests/codegen/vec-with-capacity.rs new file mode 100644 index 0000000000000..47051f2eef891 --- /dev/null +++ b/tests/codegen/vec-with-capacity.rs @@ -0,0 +1,35 @@ +//@ compile-flags: -O +//@ ignore-debug +// (with debug assertions turned on, `assert_unchecked` generates a real assertion) + +#![crate_type = "lib"] +#![feature(try_with_capacity)] + +// CHECK-LABEL: @with_capacity_does_not_grow1 +#[no_mangle] +pub fn with_capacity_does_not_grow1() -> Vec { + let v = Vec::with_capacity(1234); + // CHECK: call {{.*}}__rust_alloc( + // CHECK-NOT: call {{.*}}__rust_realloc + // CHECK-NOT: call {{.*}}capacity_overflow + // CHECK-NOT: call {{.*}}finish_grow + // CHECK-NOT: call {{.*}}reserve + // CHECK-NOT: memcpy + // CHECK-NOT: memset + v +} + +// CHECK-LABEL: @try_with_capacity_does_not_grow2 +#[no_mangle] +pub fn try_with_capacity_does_not_grow2() -> Option>> { + let v = Vec::try_with_capacity(1234).ok()?; + // CHECK: call {{.*}}__rust_alloc( + // CHECK-NOT: call {{.*}}__rust_realloc + // CHECK-NOT: call {{.*}}capacity_overflow + // CHECK-NOT: call {{.*}}finish_grow + // CHECK-NOT: call {{.*}}handle_alloc_error + // CHECK-NOT: call {{.*}}reserve + // CHECK-NOT: memcpy + // CHECK-NOT: memset + Some(v) +} diff --git a/tests/ui/hygiene/panic-location.run.stderr b/tests/ui/hygiene/panic-location.run.stderr index 5c552411da7f3..ec0ce18c3dfa7 100644 --- a/tests/ui/hygiene/panic-location.run.stderr +++ b/tests/ui/hygiene/panic-location.run.stderr @@ -1,3 +1,3 @@ -thread 'main' panicked at library/alloc/src/raw_vec.rs:571:5: +thread 'main' panicked at library/alloc/src/raw_vec.rs:26:5: capacity overflow note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace diff --git a/tests/ui/suggestions/deref-path-method.stderr b/tests/ui/suggestions/deref-path-method.stderr index a2b68fa966fcb..b27d9aef06614 100644 --- a/tests/ui/suggestions/deref-path-method.stderr +++ b/tests/ui/suggestions/deref-path-method.stderr @@ -7,9 +7,9 @@ LL | Vec::contains(&vec, &0); note: if you're trying to build a new `Vec<_, _>` consider using one of the following associated functions: Vec::::new Vec::::with_capacity + Vec::::try_with_capacity Vec::::from_raw_parts - Vec::::new_in - and 2 others + and 4 others --> $SRC_DIR/alloc/src/vec/mod.rs:LL:COL help: the function `contains` is implemented on `[_]` | diff --git a/tests/ui/ufcs/bad-builder.stderr b/tests/ui/ufcs/bad-builder.stderr index e1c5e45b3ebb7..9cfeb7a5d09d6 100644 --- a/tests/ui/ufcs/bad-builder.stderr +++ b/tests/ui/ufcs/bad-builder.stderr @@ -7,9 +7,9 @@ LL | Vec::::mew() note: if you're trying to build a new `Vec` consider using one of the following associated functions: Vec::::new Vec::::with_capacity + Vec::::try_with_capacity Vec::::from_raw_parts - Vec::::new_in - and 2 others + and 4 others --> $SRC_DIR/alloc/src/vec/mod.rs:LL:COL help: there is an associated function `new` with a similar name |