From 50959c48c25839bedb138a8e830399587fa11c77 Mon Sep 17 00:00:00 2001 From: John Ericson Date: Wed, 7 Jul 2021 00:11:34 -0400 Subject: [PATCH 1/3] alloc: add some try_* methods Rust-for-Linux needs Based off of https://github.com/Rust-for-Linux/linux/commit/487d7578bd036ba57b0998d7af69e0bdcc595da8 --- library/alloc/src/lib.rs | 1 + library/alloc/src/raw_vec.rs | 55 +++- library/alloc/src/slice.rs | 97 +++++++- library/alloc/src/str.rs | 17 ++ library/alloc/src/sync.rs | 54 +++- library/alloc/src/vec/mod.rs | 318 +++++++++++++++++++++++- library/alloc/src/vec/spec_extend.rs | 83 +++++++ library/std/src/collections/hash/map.rs | 4 +- 8 files changed, 621 insertions(+), 8 deletions(-) diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs index 635708fd4cf6e..f79202b9cca24 100644 --- a/library/alloc/src/lib.rs +++ b/library/alloc/src/lib.rs @@ -126,6 +126,7 @@ #![feature(iter_zip)] #![feature(lang_items)] #![feature(layout_for_ptr)] +#![feature(more_fallible_allocation_methods)] #![feature(negative_impls)] #![feature(never_type)] #![feature(nll)] diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs index 75dbd4678bb47..99249df6d55f6 100644 --- a/library/alloc/src/raw_vec.rs +++ b/library/alloc/src/raw_vec.rs @@ -19,7 +19,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, @@ -94,6 +93,16 @@ impl RawVec { Self::with_capacity_in(capacity, Global) } + /// Tries to create a `RawVec` (on the system heap) with exactly the + /// capacity and alignment requirements for a `[T; capacity]`. This is + /// equivalent to calling `RawVec::new` when `capacity` is `0` or `T` is + /// zero-sized. Note that if `T` is zero-sized this means you will + /// *not* get a `RawVec` with the requested capacity. + #[inline] + pub fn try_with_capacity(capacity: usize) -> Result { + Self::try_with_capacity_in(capacity, Global) + } + /// Like `with_capacity`, but guarantees the buffer is zeroed. #[cfg(not(no_global_oom_handling))] #[must_use] @@ -146,6 +155,13 @@ impl RawVec { Self::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))] @@ -220,6 +236,33 @@ impl RawVec { } } + fn try_allocate_in( + capacity: usize, + init: AllocInit, + alloc: A, + ) -> Result { + if mem::size_of::() == 0 { + return Ok(Self::new_in(alloc)); + } + + let layout = Layout::array::(capacity)?; + alloc_guard(layout.size())?; + let result = match init { + AllocInit::Uninitialized => alloc.allocate(layout), + AllocInit::Zeroed => alloc.allocate_zeroed(layout), + }; + let ptr = match result { + Ok(ptr) => ptr, + Err(_) => return Err(TryReserveErrorKind::AllocError { layout, non_exhaustive: () }), + }; + + Ok(Self { + ptr: unsafe { Unique::new_unchecked(ptr.cast().as_ptr()) }, + cap: Self::capacity_from_bytes(ptr.len()), + alloc, + }) + } + /// Reconstitutes a `RawVec` from a pointer, capacity, and allocator. /// /// # Safety @@ -396,6 +439,16 @@ impl RawVec { pub fn shrink_to_fit(&mut self, amount: usize) { handle_reserve(self.shrink(amount)); } + + /// Tries to shrink the allocation down to the specified amount. If the given amount + /// is 0, actually completely deallocates. + /// + /// # Panics + /// + /// Panics if the given amount is *larger* than the current capacity. + pub fn try_shrink_to_fit(&mut self, amount: usize) -> Result<(), TryReserveError> { + self.shrink(amount) + } } impl RawVec { diff --git a/library/alloc/src/slice.rs b/library/alloc/src/slice.rs index 860f21085f325..957c7a5d33484 100644 --- a/library/alloc/src/slice.rs +++ b/library/alloc/src/slice.rs @@ -93,11 +93,11 @@ use core::mem::size_of; use core::ptr; use crate::alloc::Allocator; -#[cfg(not(no_global_oom_handling))] use crate::alloc::Global; #[cfg(not(no_global_oom_handling))] use crate::borrow::ToOwned; use crate::boxed::Box; +use crate::collections::TryReserveError; use crate::vec::Vec; #[unstable(feature = "slice_range", issue = "76393")] @@ -153,6 +153,7 @@ mod hack { use core::alloc::Allocator; use crate::boxed::Box; + use crate::collections::TryReserveError; use crate::vec::Vec; // We shouldn't add inline attribute to this since this is used in @@ -172,6 +173,14 @@ mod hack { T::to_vec(s, alloc) } + #[inline] + pub fn try_to_vec( + s: &[T], + alloc: A, + ) -> Result, TryReserveError> { + T::try_to_vec(s, alloc) + } + #[cfg(not(no_global_oom_handling))] pub trait ConvertVec { fn to_vec(s: &[Self], alloc: A) -> Vec @@ -179,6 +188,12 @@ mod hack { Self: Sized; } + pub trait TryConvertVec { + fn try_to_vec(s: &[Self], alloc: A) -> Result, TryReserveError> + where + Self: Sized; + } + #[cfg(not(no_global_oom_handling))] impl ConvertVec for T { #[inline] @@ -231,6 +246,45 @@ mod hack { v } } + + impl TryConvertVec for T { + #[inline] + default fn try_to_vec( + s: &[Self], + alloc: A, + ) -> Result, TryReserveError> { + struct DropGuard<'a, T, A: Allocator> { + vec: &'a mut Vec, + num_init: usize, + } + impl<'a, T, A: Allocator> Drop for DropGuard<'a, T, A> { + #[inline] + fn drop(&mut self) { + // SAFETY: + // items were marked initialized in the loop below + unsafe { + self.vec.set_len(self.num_init); + } + } + } + let mut vec = Vec::try_with_capacity_in(s.len(), alloc)?; + let mut guard = DropGuard { vec: &mut vec, num_init: 0 }; + let slots = guard.vec.spare_capacity_mut(); + // .take(slots.len()) is necessary for LLVM to remove bounds checks + // and has better codegen than zip. + for (i, b) in s.iter().enumerate().take(slots.len()) { + guard.num_init = i; + slots[i].write(b.clone()); + } + core::mem::forget(guard); + // SAFETY: + // the vec was allocated and initialized above to at least this length. + unsafe { + vec.set_len(s.len()); + } + Ok(vec) + } + } } #[lang = "slice_alloc"] @@ -470,6 +524,24 @@ impl [T] { self.to_vec_in(Global) } + /// Tries to copy `self` into a new `Vec`. + /// + /// # Examples + /// + /// ``` + /// let s = [10, 40, 30]; + /// let x = s.try_to_vec().unwrap(); + /// // Here, `s` and `x` can be modified independently. + /// ``` + #[inline] + #[unstable(feature = "more_fallible_allocation_methods", issue = "86942")] + pub fn try_to_vec(&self) -> Result, TryReserveError> + where + T: Clone, + { + self.try_to_vec_in(Global) + } + /// Copies `self` into a new `Vec` with an allocator. /// /// # Examples @@ -494,6 +566,29 @@ impl [T] { hack::to_vec(self, alloc) } + /// Tries to copy `self` into a new `Vec` with an allocator. + /// + /// # Examples + /// + /// ``` + /// #![feature(allocator_api)] + /// + /// use std::alloc::System; + /// + /// let s = [10, 40, 30]; + /// let x = s.try_to_vec_in(System).unwrap(); + /// // Here, `s` and `x` can be modified independently. + /// ``` + #[inline] + #[unstable(feature = "more_fallible_allocation_methods", issue = "86942")] + pub fn try_to_vec_in(&self, alloc: A) -> Result, TryReserveError> + where + T: Clone, + { + // N.B., see the `hack` module in this file for more details. + hack::try_to_vec(self, alloc) + } + /// Converts `self` into a vector without clones or allocation. /// /// The resulting vector can be converted back into a box via diff --git a/library/alloc/src/str.rs b/library/alloc/src/str.rs index e1d0ee42f4e90..0d7f986f70391 100644 --- a/library/alloc/src/str.rs +++ b/library/alloc/src/str.rs @@ -36,6 +36,7 @@ use core::unicode::conversions; use crate::borrow::ToOwned; use crate::boxed::Box; +use crate::collections::TryReserveError; use crate::slice::{Concat, Join, SliceIndex}; use crate::string::String; use crate::vec::Vec; @@ -579,6 +580,22 @@ impl str { // make_ascii_lowercase() preserves the UTF-8 invariant. unsafe { String::from_utf8_unchecked(bytes) } } + + /// Tries to create a `String`. + /// + /// # Examples + /// + /// Basic usage: + /// + /// ``` + /// let s: &str = "a"; + /// let ss: String = s.try_to_owned().unwrap(); + /// ``` + #[inline] + #[unstable(feature = "more_fallible_allocation_methods", issue = "86942")] + pub fn try_to_owned(&self) -> Result { + unsafe { Ok(String::from_utf8_unchecked(self.as_bytes().try_to_vec()?)) } + } } /// Converts a boxed slice of bytes to a boxed string slice without checking diff --git a/library/alloc/src/sync.rs b/library/alloc/src/sync.rs index b75e9a2f3c71e..2e8aa48d1cf7e 100644 --- a/library/alloc/src/sync.rs +++ b/library/alloc/src/sync.rs @@ -34,10 +34,10 @@ use crate::alloc::{box_free, WriteCloneIntoRaw}; use crate::alloc::{AllocError, Allocator, Global, Layout}; use crate::borrow::{Cow, ToOwned}; use crate::boxed::Box; +use crate::collections::TryReserveError; use crate::rc::is_dangling; #[cfg(not(no_global_oom_handling))] use crate::string::String; -#[cfg(not(no_global_oom_handling))] use crate::vec::Vec; #[cfg(test)] @@ -1207,6 +1207,19 @@ impl Arc<[T]> { } } + /// Tries to allocate an `ArcInner<[T]>` with the given length. + unsafe fn try_allocate_for_slice(len: usize) -> Result<*mut ArcInner<[T]>, TryReserveError> { + unsafe { + let layout = Layout::array::(len)?; + Self::try_allocate_for_layout( + layout, + |l| Global.allocate(l), + |mem| ptr::slice_from_raw_parts_mut(mem as *mut T, len) as *mut ArcInner<[T]>, + ) + .map_err(|_| TryReserveErrorKind::AllocError { layout, non_exhaustive: () }) + } + } + /// Copy elements from slice into newly allocated Arc<\[T\]> /// /// Unsafe because the caller must either take ownership or bind `T: Copy`. @@ -1221,6 +1234,19 @@ impl Arc<[T]> { } } + /// Tries to copy elements from slice into newly allocated Arc<\[T\]> + /// + /// Unsafe because the caller must either take ownership or bind `T: Copy`. + unsafe fn try_copy_from_slice(v: &[T]) -> Result, TryReserveError> { + unsafe { + let ptr = Self::try_allocate_for_slice(v.len())?; + + ptr::copy_nonoverlapping(v.as_ptr(), &mut (*ptr).data as *mut [T] as *mut T, v.len()); + + Ok(Self::from_ptr(ptr)) + } + } + /// Constructs an `Arc<[T]>` from an iterator known to be of a certain size. /// /// Behavior is undefined should the size be wrong. @@ -2470,6 +2496,32 @@ impl From> for Arc<[T]> { } } +// Avoid `error: specializing impl repeats parameter` implementing `TryFrom`. +impl Arc<[T]> { + /// Tries to allocate a reference-counted slice and move `v`'s items into it. + /// + /// # Example + /// + /// ``` + /// # use std::sync::Arc; + /// let unique: Vec = vec![1, 2, 3]; + /// let shared: Arc<[i32]> = Arc::try_from(unique).unwrap(); + /// assert_eq!(&[1, 2, 3], &shared[..]); + /// ``` + #[unstable(feature = "more_fallible_allocation_methods", issue = "86942")] + #[inline] + pub fn try_from_vec(mut v: Vec) -> Result { + unsafe { + let arc = Arc::try_copy_from_slice(&v)?; + + // Allow the Vec to free its memory, but not destroy its contents + v.set_len(0); + + Ok(arc) + } + } +} + #[stable(feature = "shared_from_cow", since = "1.45.0")] impl<'a, B> From> for Arc where diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index 20a16869cb3f8..bb7b879890137 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -1,3 +1,5 @@ +// ignore-tidy-filelength + //! A contiguous growable array type with heap-allocated contents, written //! `Vec`. //! @@ -118,10 +120,8 @@ use self::spec_from_elem::SpecFromElem; #[cfg(not(no_global_oom_handling))] mod spec_from_elem; -#[cfg(not(no_global_oom_handling))] use self::set_len_on_drop::SetLenOnDrop; -#[cfg(not(no_global_oom_handling))] mod set_len_on_drop; #[cfg(not(no_global_oom_handling))] @@ -145,7 +145,8 @@ mod spec_from_iter; #[cfg(not(no_global_oom_handling))] use self::spec_extend::SpecExtend; -#[cfg(not(no_global_oom_handling))] +use self::spec_extend::TrySpecExtend; + mod spec_extend; /// A contiguous growable array type, written as `Vec` and pronounced 'vector'. @@ -470,6 +471,49 @@ impl Vec { Self::with_capacity_in(capacity, Global) } + /// Tries to construct a new, empty `Vec` with the specified capacity. + /// + /// The vector will be able to hold exactly `capacity` elements without + /// reallocating. If `capacity` is 0, the vector will not allocate. + /// + /// It is important to note that although the returned vector has the + /// *capacity* specified, the vector will have a zero *length*. For an + /// explanation of the difference between length and capacity, see + /// *[Capacity and reallocation]*. + /// + /// [Capacity and reallocation]: #capacity-and-reallocation + /// + /// # Examples + /// + /// ``` + /// let mut vec = Vec::try_with_capacity(10).unwrap(); + /// + /// // The vector contains no items, even though it has capacity for more + /// assert_eq!(vec.len(), 0); + /// assert_eq!(vec.capacity(), 10); + /// + /// // These are all done without reallocating... + /// for i in 0..10 { + /// vec.push(i); + /// } + /// assert_eq!(vec.len(), 10); + /// assert_eq!(vec.capacity(), 10); + /// + /// // ...but this may make the vector reallocate + /// vec.push(11); + /// assert_eq!(vec.len(), 11); + /// assert!(vec.capacity() >= 11); + /// + /// let mut result = Vec::try_with_capacity(usize::MAX); + /// assert!(result.is_err()); + /// ``` + #[inline] + #[doc(alias = "malloc")] + #[unstable(feature = "more_fallible_allocation_methods", issue = "86942")] + pub fn try_with_capacity(capacity: usize) -> Result { + Self::try_with_capacity_in(capacity, Global) + } + /// Creates a `Vec` directly from the raw components of another vector. /// /// # Safety @@ -609,6 +653,53 @@ impl Vec { Vec { buf: RawVec::with_capacity_in(capacity, alloc), len: 0 } } + /// Tries to construct a new, empty `Vec` with the specified capacity + /// with the provided allocator. + /// + /// The vector will be able to hold exactly `capacity` elements without + /// reallocating. If `capacity` is 0, the vector will not allocate. + /// + /// It is important to note that although the returned vector has the + /// *capacity* specified, the vector will have a zero *length*. For an + /// explanation of the difference between length and capacity, see + /// *[Capacity and reallocation]*. + /// + /// [Capacity and reallocation]: #capacity-and-reallocation + /// + /// # Examples + /// + /// ``` + /// #![feature(allocator_api)] + /// + /// use std::alloc::System; + /// + /// let mut vec = Vec::try_with_capacity_in(10, System).unwrap(); + /// + /// // The vector contains no items, even though it has capacity for more + /// assert_eq!(vec.len(), 0); + /// assert_eq!(vec.capacity(), 10); + /// + /// // These are all done without reallocating... + /// for i in 0..10 { + /// vec.push(i); + /// } + /// assert_eq!(vec.len(), 10); + /// assert_eq!(vec.capacity(), 10); + /// + /// // ...but this may make the vector reallocate + /// vec.push(11); + /// assert_eq!(vec.len(), 11); + /// assert!(vec.capacity() >= 11); + /// + /// let mut result = Vec::try_with_capacity_in(usize::MAX, System); + /// assert!(result.is_err()); + /// ``` + #[inline] + #[unstable(feature = "more_fallible_allocation_methods", issue = "86942")] + 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 the raw components of another vector. /// /// # Safety @@ -940,6 +1031,33 @@ impl Vec { } } + /// Tries to shrink the capacity of the vector as much as possible. + /// + /// It will drop down as close as possible to the length but the allocator + /// may still inform the vector that there is space for a few more elements. + /// + /// # Examples + /// + /// ``` + /// let mut vec = Vec::with_capacity(10); + /// vec.extend([1, 2, 3]); + /// assert_eq!(vec.capacity(), 10); + /// vec.try_shrink_to_fit().unwrap(); + /// assert!(vec.capacity() >= 3); + /// ``` + #[doc(alias = "realloc")] + #[unstable(feature = "more_fallible_allocation_methods", issue = "86942")] + pub fn try_shrink_to_fit(&mut self) -> Result<(), TryReserveError> { + // The capacity is never less than the length, and there's nothing to do when + // they are equal, so we can avoid the panic case in `RawVec::try_shrink_to_fit` + // by only calling it with a greater capacity. + if self.capacity() <= self.len { + return Ok(()); + } + + self.buf.try_shrink_to_fit(self.len) + } + /// Shrinks the capacity of the vector with a lower bound. /// /// The capacity will remain at least as large as both the length @@ -1002,6 +1120,41 @@ impl Vec { } } + /// Tries to convert the vector into [`Box<[T]>`][owned slice]. + /// + /// Note that this will drop any excess capacity. + /// + /// [owned slice]: Box + /// + /// # Examples + /// + /// ``` + /// let v = vec![1, 2, 3]; + /// + /// let slice = v.try_into_boxed_slice().unwrap(); + /// ``` + /// + /// Any excess capacity is removed: + /// + /// ``` + /// let mut vec = Vec::with_capacity(10); + /// vec.extend([1, 2, 3]); + /// + /// assert_eq!(vec.capacity(), 10); + /// let slice = vec.try_into_boxed_slice().unwrap(); + /// assert_eq!(slice.into_vec().capacity(), 3); + /// ``` + #[unstable(feature = "more_fallible_allocation_methods", issue = "86942")] + pub fn try_into_boxed_slice(mut self) -> Result, TryReserveError> { + unsafe { + self.try_shrink_to_fit()?; + let me = ManuallyDrop::new(self); + let buf = ptr::read(&me.buf); + let len = me.len(); + Ok(buf.into_box(len).assume_init()) + } + } + /// Shortens the vector, keeping the first `len` elements and dropping /// the rest. /// @@ -1703,6 +1856,29 @@ impl Vec { } } + /// Tries to append an element to the back of a collection. + /// + /// # Examples + /// + /// ``` + /// let mut vec = vec![1, 2]; + /// vec.try_push(3).unwrap(); + /// assert_eq!(vec, [1, 2, 3]); + /// ``` + #[inline] + #[unstable(feature = "more_fallible_allocation_methods", issue = "86942")] + pub fn try_push(&mut self, value: T) -> Result<(), TryReserveError> { + if self.len == self.buf.capacity() { + self.try_reserve(1)?; + } + unsafe { + let end = self.as_mut_ptr().add(self.len); + ptr::write(end, value); + self.len += 1; + } + Ok(()) + } + /// Removes the last element from a vector and returns it, or [`None`] if it /// is empty. /// @@ -1762,6 +1938,17 @@ impl Vec { self.len += count; } + /// Tries to append elements to `Self` from other buffer. + #[inline] + unsafe fn try_append_elements(&mut self, other: *const [T]) -> Result<(), TryReserveError> { + let count = unsafe { (*other).len() }; + self.try_reserve(count)?; + let len = self.len(); + unsafe { ptr::copy_nonoverlapping(other as *const T, self.as_mut_ptr().add(len), count) }; + self.len += count; + Ok(()) + } + /// Creates a draining iterator that removes the specified range in the vector /// and yields the removed items. /// @@ -2164,6 +2351,44 @@ impl Vec { } } + /// Tries to resize the `Vec` in-place so that `len` is equal to `new_len`. + /// + /// If `new_len` is greater than `len`, the `Vec` is extended by the + /// difference, with each additional slot filled with `value`. + /// If `new_len` is less than `len`, the `Vec` is simply truncated. + /// + /// This method requires `T` to implement [`Clone`], + /// in order to be able to clone the passed value. + /// If you need more flexibility (or want to rely on [`Default`] instead of + /// [`Clone`]), use [`Vec::resize_with`]. + /// + /// # Examples + /// + /// ``` + /// let mut vec = vec!["hello"]; + /// vec.try_resize(3, "world").unwrap(); + /// assert_eq!(vec, ["hello", "world", "world"]); + /// + /// let mut vec = vec![1, 2, 3, 4]; + /// vec.try_resize(2, 0).unwrap(); + /// assert_eq!(vec, [1, 2]); + /// + /// let mut vec = vec![42]; + /// let result = vec.try_resize(usize::MAX, 0); + /// assert!(result.is_err()); + /// ``` + #[unstable(feature = "more_fallible_allocation_methods", issue = "86942")] + pub fn try_resize(&mut self, new_len: usize, value: T) -> Result<(), TryReserveError> { + let len = self.len(); + + if new_len > len { + self.try_extend_with(new_len - len, ExtendElement(value)) + } else { + self.truncate(new_len); + Ok(()) + } + } + /// Clones and appends all elements in a slice to the `Vec`. /// /// Iterates over the slice `other`, clones each element, and then appends @@ -2189,6 +2414,30 @@ impl Vec { self.spec_extend(other.iter()) } + /// Tries to clone and append all elements in a slice to the `Vec`. + /// + /// Iterates over the slice `other`, clones each element, and then appends + /// it to this `Vec`. The `other` vector is traversed in-order. + /// + /// Note that this function is same as [`extend`] except that it is + /// specialized to work with slices instead. If and when Rust gets + /// specialization this function will likely be deprecated (but still + /// available). + /// + /// # Examples + /// + /// ``` + /// let mut vec = vec![1]; + /// vec.try_extend_from_slice(&[2, 3, 4]).unwrap(); + /// assert_eq!(vec, [1, 2, 3, 4]); + /// ``` + /// + /// [`extend`]: Vec::extend + #[unstable(feature = "more_fallible_allocation_methods", issue = "86942")] + pub fn try_extend_from_slice(&mut self, other: &[T]) -> Result<(), TryReserveError> { + self.try_spec_extend(other.iter()) + } + /// Copies elements from `src` range to the end of the vector. /// /// # Panics @@ -2293,6 +2542,40 @@ impl Vec { // len set by scope guard } } + + /// Try to extend the vector by `n` values, using the given generator. + fn try_extend_with>( + &mut self, + n: usize, + mut value: E, + ) -> Result<(), TryReserveError> { + self.try_reserve(n)?; + + unsafe { + let mut ptr = self.as_mut_ptr().add(self.len()); + // Use SetLenOnDrop to work around bug where compiler + // might not realize the store through `ptr` through self.set_len() + // don't alias. + let mut local_len = SetLenOnDrop::new(&mut self.len); + + // Write all elements except the last one + for _ in 1..n { + ptr::write(ptr, value.next()); + ptr = ptr.offset(1); + // Increment the length in every step in case next() panics + local_len.increment_len(1); + } + + if n > 0 { + // We can write the last element directly without cloning needlessly + ptr::write(ptr, value.last()); + local_len.increment_len(1); + } + + // len set by scope guard + Ok(()) + } + } } impl Vec { @@ -2627,6 +2910,35 @@ impl Vec { } } + // leaf method to which various SpecFrom/SpecExtend implementations delegate when + // they have no further optimizations to apply + fn try_extend_desugared>( + &mut self, + mut iterator: I, + ) -> Result<(), TryReserveError> { + // This is the case for a general iterator. + // + // This function should be the moral equivalent of: + // + // for item in iterator { + // self.push(item); + // } + while let Some(element) = iterator.next() { + let len = self.len(); + if len == self.capacity() { + let (lower, _) = iterator.size_hint(); + self.try_reserve(lower.saturating_add(1))?; + } + unsafe { + ptr::write(self.as_mut_ptr().add(len), element); + // NB can't overflow since we would have had to alloc the address space + self.set_len(len + 1); + } + } + + Ok(()) + } + /// Creates a splicing iterator that replaces the specified range in the vector /// with the given `replace_with` iterator and yields the removed items. /// `replace_with` does not need to be the same length as `range`. diff --git a/library/alloc/src/vec/spec_extend.rs b/library/alloc/src/vec/spec_extend.rs index c3b4534096de5..6773e2ddb1cce 100644 --- a/library/alloc/src/vec/spec_extend.rs +++ b/library/alloc/src/vec/spec_extend.rs @@ -1,4 +1,5 @@ use crate::alloc::Allocator; +use crate::vec::TryReserveError; use core::iter::TrustedLen; use core::ptr::{self}; use core::slice::{self}; @@ -6,10 +7,17 @@ use core::slice::{self}; use super::{IntoIter, SetLenOnDrop, Vec}; // Specialization trait used for Vec::extend +#[cfg(not(no_global_oom_handling))] pub(super) trait SpecExtend { fn spec_extend(&mut self, iter: I); } +// Specialization trait used for Vec::try_extend +pub(super) trait TrySpecExtend { + fn try_spec_extend(&mut self, iter: I) -> Result<(), TryReserveError>; +} + +#[cfg(not(no_global_oom_handling))] impl SpecExtend for Vec where I: Iterator, @@ -19,6 +27,16 @@ where } } +impl TrySpecExtend for Vec +where + I: Iterator, +{ + default fn try_spec_extend(&mut self, iter: I) -> Result<(), TryReserveError> { + self.try_extend_desugared(iter) + } +} + +#[cfg(not(no_global_oom_handling))] impl SpecExtend for Vec where I: TrustedLen, @@ -57,6 +75,39 @@ where } } +impl TrySpecExtend for Vec +where + I: TrustedLen, +{ + default fn try_spec_extend(&mut self, iterator: I) -> Result<(), TryReserveError> { + // This is the case for a TrustedLen iterator. + let (low, high) = iterator.size_hint(); + if let Some(additional) = high { + debug_assert_eq!( + low, + additional, + "TrustedLen iterator's size hint is not exact: {:?}", + (low, high) + ); + self.try_reserve(additional)?; + unsafe { + let mut ptr = self.as_mut_ptr().add(self.len()); + let mut local_len = SetLenOnDrop::new(&mut self.len); + iterator.for_each(move |element| { + ptr::write(ptr, element); + ptr = ptr.offset(1); + // NB can't overflow since we would have had to alloc the address space + local_len.increment_len(1); + }); + } + Ok(()) + } else { + Err(TryReserveErrorKind::CapacityOverflow) + } + } +} + +#[cfg(not(no_global_oom_handling))] impl SpecExtend> for Vec { fn spec_extend(&mut self, mut iterator: IntoIter) { unsafe { @@ -66,6 +117,17 @@ impl SpecExtend> for Vec { } } +impl TrySpecExtend> for Vec { + fn try_spec_extend(&mut self, mut iterator: IntoIter) -> Result<(), TryReserveError> { + unsafe { + self.try_append_elements(iterator.as_slice() as _)?; + } + iterator.ptr = iterator.end; + Ok(()) + } +} + +#[cfg(not(no_global_oom_handling))] impl<'a, T: 'a, I, A: Allocator + 'a> SpecExtend<&'a T, I> for Vec where I: Iterator, @@ -76,6 +138,17 @@ where } } +impl<'a, T: 'a, I, A: Allocator + 'a> TrySpecExtend<&'a T, I> for Vec +where + I: Iterator, + T: Clone, +{ + default fn try_spec_extend(&mut self, iterator: I) -> Result<(), TryReserveError> { + self.try_spec_extend(iterator.cloned()) + } +} + +#[cfg(not(no_global_oom_handling))] impl<'a, T: 'a, A: Allocator + 'a> SpecExtend<&'a T, slice::Iter<'a, T>> for Vec where T: Copy, @@ -85,3 +158,13 @@ where unsafe { self.append_elements(slice) }; } } + +impl<'a, T: 'a, A: Allocator + 'a> TrySpecExtend<&'a T, slice::Iter<'a, T>> for Vec +where + T: Copy, +{ + fn try_spec_extend(&mut self, iterator: slice::Iter<'a, T>) -> Result<(), TryReserveError> { + let slice = iterator.as_slice(); + unsafe { self.try_append_elements(slice) } + } +} diff --git a/library/std/src/collections/hash/map.rs b/library/std/src/collections/hash/map.rs index 01019344f4f67..89571a7c0af0c 100644 --- a/library/std/src/collections/hash/map.rs +++ b/library/std/src/collections/hash/map.rs @@ -3006,10 +3006,10 @@ fn map_entry<'a, K: 'a, V: 'a>(raw: base::RustcEntry<'a, K, V>) -> Entry<'a, K, #[inline] pub(super) fn map_try_reserve_error(err: hashbrown::TryReserveError) -> TryReserveError { match err { - hashbrown::TryReserveError::CapacityOverflow => { + hashbrown::TryReserveErrorKind::CapacityOverflow => { TryReserveErrorKind::CapacityOverflow.into() } - hashbrown::TryReserveError::AllocError { layout } => { + hashbrown::TryReserveErrorKind::AllocError { layout } => { TryReserveErrorKind::AllocError { layout, non_exhaustive: () }.into() } } From 313490aa2de229b62ce72dcf11822d7206023953 Mon Sep 17 00:00:00 2001 From: John Ericson Date: Wed, 7 Jul 2021 00:57:29 -0400 Subject: [PATCH 2/3] Dedup regular methods from new `try_` ones I did all the non-specialization trait ones. I hope this doesn't cause performance issues --- library/alloc/src/collections/mod.rs | 17 +++++ library/alloc/src/raw_vec.rs | 64 ++++++------------- library/alloc/src/string.rs | 4 +- library/alloc/src/sync.rs | 30 ++++----- library/alloc/src/vec/mod.rs | 96 +++++++--------------------- 5 files changed, 73 insertions(+), 138 deletions(-) diff --git a/library/alloc/src/collections/mod.rs b/library/alloc/src/collections/mod.rs index 77d28bdfe6475..50b495fcc514a 100644 --- a/library/alloc/src/collections/mod.rs +++ b/library/alloc/src/collections/mod.rs @@ -52,7 +52,11 @@ pub use linked_list::LinkedList; #[doc(no_inline)] pub use vec_deque::VecDeque; +#[cfg(not(no_global_oom_handling))] +use crate::alloc::handle_alloc_error; use crate::alloc::{Layout, LayoutError}; +#[cfg(not(no_global_oom_handling))] +use crate::raw_vec::capacity_overflow; use core::fmt::Display; /// The error type for `try_reserve` methods. @@ -105,6 +109,19 @@ pub enum TryReserveErrorKind { }, } +#[cfg(not(no_global_oom_handling))] +impl TryReserveError { + /// Panic, or worse, according to the global handlers for each case. + pub(crate) fn handle(self) -> ! { + match self { + TryReserveErrorKind::CapacityOverflow => capacity_overflow(), + TryReserveErrorKind::AllocError { layout, non_exhaustive: () } => { + handle_alloc_error(layout) + } + } + } +} + #[unstable( feature = "try_reserve_kind", reason = "Uncertain how much info should be exposed", diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs index 99249df6d55f6..eff11fbdeccbc 100644 --- a/library/alloc/src/raw_vec.rs +++ b/library/alloc/src/raw_vec.rs @@ -206,33 +206,9 @@ impl RawVec { #[cfg(not(no_global_oom_handling))] fn allocate_in(capacity: usize, init: AllocInit, alloc: A) -> Self { - if mem::size_of::() == 0 { - 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(), - }; - match alloc_guard(layout.size()) { - Ok(_) => {} - Err(_) => capacity_overflow(), - } - 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), - }; - - Self { - ptr: unsafe { Unique::new_unchecked(ptr.cast().as_ptr()) }, - cap: Self::capacity_from_bytes(ptr.len()), - alloc, - } + match Self::try_allocate_in(capacity, init, alloc) { + Ok(r) => r, + Err(e) => e.handle(), } } @@ -437,7 +413,7 @@ impl RawVec { /// Aborts on OOM. #[cfg(not(no_global_oom_handling))] pub fn shrink_to_fit(&mut self, amount: usize) { - handle_reserve(self.shrink(amount)); + handle_reserve(self.try_shrink_to_fit(amount)); } /// Tries to shrink the allocation down to the specified amount. If the given amount @@ -447,7 +423,19 @@ impl RawVec { /// /// Panics if the given amount is *larger* than the current capacity. pub fn try_shrink_to_fit(&mut self, amount: usize) -> Result<(), TryReserveError> { - self.shrink(amount) + assert!(amount <= self.capacity(), "Tried to shrink to a larger capacity"); + + let (ptr, layout) = if let Some(mem) = self.current_memory() { mem } else { return Ok(()) }; + let new_size = amount * mem::size_of::(); + + let ptr = unsafe { + let new_layout = Layout::from_size_align_unchecked(new_size, layout.align()); + self.alloc + .shrink(ptr, layout, new_layout) + .map_err(|_| AllocError { layout: new_layout, non_exhaustive: () })? + }; + self.set_ptr(ptr); + Ok(()) } } @@ -519,22 +507,6 @@ impl RawVec { self.set_ptr(ptr); Ok(()) } - - fn shrink(&mut self, amount: usize) -> Result<(), TryReserveError> { - assert!(amount <= self.capacity(), "Tried to shrink to a larger capacity"); - - let (ptr, layout) = if let Some(mem) = self.current_memory() { mem } else { return Ok(()) }; - let new_size = amount * mem::size_of::(); - - let ptr = unsafe { - let new_layout = Layout::from_size_align_unchecked(new_size, layout.align()); - self.alloc - .shrink(ptr, layout, new_layout) - .map_err(|_| AllocError { layout: new_layout, non_exhaustive: () })? - }; - self.set_ptr(ptr); - Ok(()) - } } // This function is outside `RawVec` to minimize compile times. See the comment @@ -611,6 +583,6 @@ fn alloc_guard(alloc_size: usize) -> Result<(), TryReserveError> { // 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))] -fn capacity_overflow() -> ! { +pub(crate) fn capacity_overflow() -> ! { panic!("capacity overflow"); } diff --git a/library/alloc/src/string.rs b/library/alloc/src/string.rs index f479bf231b376..11b7cc0b02841 100644 --- a/library/alloc/src/string.rs +++ b/library/alloc/src/string.rs @@ -46,9 +46,9 @@ use core::char::{decode_utf16, REPLACEMENT_CHARACTER}; use core::fmt; use core::hash; +use core::iter::FusedIterator; #[cfg(not(no_global_oom_handling))] -use core::iter::FromIterator; -use core::iter::{from_fn, FusedIterator}; +use core::iter::{from_fn, FromIterator}; #[cfg(not(no_global_oom_handling))] use core::ops::Add; #[cfg(not(no_global_oom_handling))] diff --git a/library/alloc/src/sync.rs b/library/alloc/src/sync.rs index 2e8aa48d1cf7e..7763794fa025f 100644 --- a/library/alloc/src/sync.rs +++ b/library/alloc/src/sync.rs @@ -1199,11 +1199,10 @@ impl Arc<[T]> { #[cfg(not(no_global_oom_handling))] unsafe fn allocate_for_slice(len: usize) -> *mut ArcInner<[T]> { unsafe { - Self::allocate_for_layout( - Layout::array::(len).unwrap(), - |layout| Global.allocate(layout), - |mem| ptr::slice_from_raw_parts_mut(mem as *mut T, len) as *mut ArcInner<[T]>, - ) + match Self::try_allocate_for_slice(len) { + Ok(r) => r, + Err(e) => e.handle(), + } } } @@ -1226,11 +1225,10 @@ impl Arc<[T]> { #[cfg(not(no_global_oom_handling))] unsafe fn copy_from_slice(v: &[T]) -> Arc<[T]> { unsafe { - let ptr = Self::allocate_for_slice(v.len()); - - ptr::copy_nonoverlapping(v.as_ptr(), &mut (*ptr).data as *mut [T] as *mut T, v.len()); - - Self::from_ptr(ptr) + match Self::try_copy_from_slice(v) { + Ok(r) => r, + Err(e) => e.handle(), + } } } @@ -2484,14 +2482,10 @@ impl From> for Arc<[T]> { /// assert_eq!(&[1, 2, 3], &shared[..]); /// ``` #[inline] - fn from(mut v: Vec) -> Arc<[T]> { - unsafe { - let arc = Arc::copy_from_slice(&v); - - // Allow the Vec to free its memory, but not destroy its contents - v.set_len(0); - - arc + fn from(v: Vec) -> Arc<[T]> { + match Arc::try_from_vec(v) { + Ok(r) => r, + Err(e) => e.handle(), } } } diff --git a/library/alloc/src/vec/mod.rs b/library/alloc/src/vec/mod.rs index bb7b879890137..642a040045df2 100644 --- a/library/alloc/src/vec/mod.rs +++ b/library/alloc/src/vec/mod.rs @@ -1023,11 +1023,9 @@ impl Vec { #[cfg(not(no_global_oom_handling))] #[stable(feature = "rust1", since = "1.0.0")] pub fn shrink_to_fit(&mut self) { - // The capacity is never less than the length, and there's nothing to do when - // they are equal, so we can avoid the panic case in `RawVec::shrink_to_fit` - // by only calling it with a greater capacity. - if self.capacity() > self.len { - self.buf.shrink_to_fit(self.len); + match self.try_shrink_to_fit() { + Ok(r) => r, + Err(e) => e.handle(), } } @@ -1110,13 +1108,10 @@ impl Vec { /// ``` #[cfg(not(no_global_oom_handling))] #[stable(feature = "rust1", since = "1.0.0")] - pub fn into_boxed_slice(mut self) -> Box<[T], A> { - unsafe { - self.shrink_to_fit(); - let me = ManuallyDrop::new(self); - let buf = ptr::read(&me.buf); - let len = me.len(); - buf.into_box(len).assume_init() + pub fn into_boxed_slice(self) -> Box<[T], A> { + match self.try_into_boxed_slice() { + Ok(r) => r, + Err(e) => e.handle(), } } @@ -1844,15 +1839,9 @@ impl Vec { #[inline] #[stable(feature = "rust1", since = "1.0.0")] pub fn push(&mut self, value: T) { - // This will panic or abort if we would allocate > isize::MAX bytes - // or if the length increment would overflow for zero-sized types. - if self.len == self.buf.capacity() { - self.reserve(1); - } - unsafe { - let end = self.as_mut_ptr().add(self.len); - ptr::write(end, value); - self.len += 1; + match self.try_push(value) { + Ok(r) => r, + Err(e) => e.handle(), } } @@ -1931,11 +1920,12 @@ impl Vec { #[cfg(not(no_global_oom_handling))] #[inline] unsafe fn append_elements(&mut self, other: *const [T]) { - let count = unsafe { (*other).len() }; - self.reserve(count); - let len = self.len(); - unsafe { ptr::copy_nonoverlapping(other as *const T, self.as_mut_ptr().add(len), count) }; - self.len += count; + unsafe { + match self.try_append_elements(other) { + Ok(r) => r, + Err(e) => e.handle(), + } + } } /// Tries to append elements to `Self` from other buffer. @@ -2515,31 +2505,10 @@ impl T> ExtendWith for ExtendFunc { impl Vec { #[cfg(not(no_global_oom_handling))] /// Extend the vector by `n` values, using the given generator. - fn extend_with>(&mut self, n: usize, mut value: E) { - self.reserve(n); - - unsafe { - let mut ptr = self.as_mut_ptr().add(self.len()); - // Use SetLenOnDrop to work around bug where compiler - // might not realize the store through `ptr` through self.set_len() - // don't alias. - let mut local_len = SetLenOnDrop::new(&mut self.len); - - // Write all elements except the last one - for _ in 1..n { - ptr::write(ptr, value.next()); - ptr = ptr.offset(1); - // Increment the length in every step in case next() panics - local_len.increment_len(1); - } - - if n > 0 { - // We can write the last element directly without cloning needlessly - ptr::write(ptr, value.last()); - local_len.increment_len(1); - } - - // len set by scope guard + fn extend_with>(&mut self, n: usize, value: E) { + match self.try_extend_with(n, value) { + Ok(r) => r, + Err(e) => e.handle(), } } @@ -2886,27 +2855,10 @@ impl Vec { // leaf method to which various SpecFrom/SpecExtend implementations delegate when // they have no further optimizations to apply #[cfg(not(no_global_oom_handling))] - fn extend_desugared>(&mut self, mut iterator: I) { - // This is the case for a general iterator. - // - // This function should be the moral equivalent of: - // - // for item in iterator { - // self.push(item); - // } - while let Some(element) = iterator.next() { - let len = self.len(); - if len == self.capacity() { - let (lower, _) = iterator.size_hint(); - self.reserve(lower.saturating_add(1)); - } - unsafe { - ptr::write(self.as_mut_ptr().add(len), element); - // Since next() executes user code which can panic we have to bump the length - // after each step. - // NB can't overflow since we would have had to alloc the address space - self.set_len(len + 1); - } + fn extend_desugared>(&mut self, iterator: I) { + match self.try_extend_desugared(iterator) { + Ok(r) => r, + Err(e) => e.handle(), } } From c7a67445300ef3b2560c5bf6c468465cf7f4509b Mon Sep 17 00:00:00 2001 From: John Ericson Date: Wed, 7 Jul 2021 10:58:51 -0400 Subject: [PATCH 3/3] Add some more raw vec methods so `AllocInit::Zeroed` is always used --- library/alloc/src/raw_vec.rs | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/library/alloc/src/raw_vec.rs b/library/alloc/src/raw_vec.rs index eff11fbdeccbc..50c80a8998eee 100644 --- a/library/alloc/src/raw_vec.rs +++ b/library/alloc/src/raw_vec.rs @@ -111,6 +111,12 @@ impl RawVec { Self::with_capacity_zeroed_in(capacity, Global) } + /// Like `try_with_capacity`, but guarantees a successfully allocated buffer is zeroed. + #[inline] + pub fn try_with_capacity_zeroed(capacity: usize) -> Result { + Self::try_with_capacity_zeroed_in(capacity, Global) + } + /// Reconstitutes a `RawVec` from a pointer and capacity. /// /// # Safety @@ -170,6 +176,13 @@ impl RawVec { Self::allocate_in(capacity, AllocInit::Zeroed, alloc) } + /// Like `with_capacity_zeroed`, but parameterized over the choice + /// of allocator for the returned `RawVec`. + #[inline] + pub fn try_with_capacity_zeroed_in(capacity: usize, alloc: A) -> Result { + Self::try_allocate_in(capacity, AllocInit::Zeroed, alloc) + } + /// Converts a `Box<[T]>` into a `RawVec`. pub fn from_box(slice: Box<[T], A>) -> Self { unsafe {