From f27a22c24a51ba734e4eb25c499320d819ebe236 Mon Sep 17 00:00:00 2001 From: Kornel Date: Tue, 30 Jan 2024 15:58:08 +0000 Subject: [PATCH 1/3] try_with_capacity for RawVec --- library/alloc/src/raw_vec.rs | 36 +++++++++++++++++------------- library/alloc/src/raw_vec/tests.rs | 7 +++--- 2 files changed, 25 insertions(+), 18 deletions(-) diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs index dd8d6f6c7e634..2d86fa377f16d 100644 --- a/library/alloc/src/raw_vec.rs +++ b/library/alloc/src/raw_vec.rs @@ -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, @@ -93,6 +92,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 +105,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 +143,7 @@ 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 `with_capacity_zeroed`, but parameterized over the choice @@ -150,7 +151,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 +180,40 @@ 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), 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 +543,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 */ } } } 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. From 78fb977d6b600865b7887245d24f6dca22a0099a Mon Sep 17 00:00:00 2001 From: Kornel Date: Tue, 30 Jan 2024 16:08:57 +0000 Subject: [PATCH 2/3] try_with_capacity for Vec, VecDeque, String #91913 --- .../alloc/src/collections/vec_deque/mod.rs | 24 +++++++++++++ library/alloc/src/lib.rs | 1 + library/alloc/src/raw_vec.rs | 9 +++++ library/alloc/src/string.rs | 13 +++++++ library/alloc/src/vec/mod.rs | 34 ++++++++++++++++++ library/alloc/tests/lib.rs | 1 + library/alloc/tests/string.rs | 11 ++++++ library/alloc/tests/vec.rs | 12 +++++++ library/alloc/tests/vec_deque.rs | 11 ++++++ tests/codegen/vec-with-capacity.rs | 35 +++++++++++++++++++ tests/ui/suggestions/deref-path-method.stderr | 4 +-- tests/ui/ufcs/bad-builder.stderr | 4 +-- 12 files changed, 155 insertions(+), 4 deletions(-) create mode 100644 tests/codegen/vec-with-capacity.rs 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 2d86fa377f16d..c5cf12209d90a 100644 --- a/library/alloc/src/raw_vec.rs +++ b/library/alloc/src/raw_vec.rs @@ -20,6 +20,7 @@ mod tests; 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, } @@ -146,6 +147,13 @@ impl RawVec { 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 /// of allocator for the returned `RawVec`. #[cfg(not(no_global_oom_handling))] @@ -203,6 +211,7 @@ impl RawVec { 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 { 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 7bd19875584a3..4b8b095c752f7 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 c4e89a58a05ac..e4e1a02fd8325 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/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 | From 784e6a1e080e5ba18e5c246e744e2d20525d1c3d Mon Sep 17 00:00:00 2001 From: Kornel Date: Wed, 31 Jan 2024 15:23:52 +0000 Subject: [PATCH 3/3] Move capacity_overflow function to make ui tests change less Code changes in raw_vec require blessing UI tests every time --- library/alloc/src/raw_vec.rs | 18 +++++++++--------- tests/ui/hygiene/panic-location.run.stderr | 2 +- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs index c5cf12209d90a..5e37de18c954f 100644 --- a/library/alloc/src/raw_vec.rs +++ b/library/alloc/src/raw_vec.rs @@ -17,6 +17,15 @@ 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, @@ -576,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/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