Dedup regular methods from new try_ ones

Ericson2314 · Ericson2314 · commit 313490aa2de2 · 2021-10-19T11:08:22.000-04:00
I did all the non-specialization trait ones.

I hope this doesn't cause performance issues
diff --git 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
@@ -206,33 +206,9 @@ impl<T, A: Allocator> RawVec<T, A> {
 
     #[cfg(not(no_global_oom_handling))]
     fn allocate_in(capacity: usize, init: AllocInit, alloc: A) -> Self {
-        if mem::size_of::<T>() == 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::<T>(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<T, A: Allocator> RawVec<T, A> {
     /// 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<T, A: Allocator> RawVec<T, A> {
     ///
     /// 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::<T>();
+
+        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<T, A: Allocator> RawVec<T, A> {
         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::<T>();
-
-        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
@@ -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
@@ -1199,11 +1199,10 @@ impl<T> 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::<T>(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<T> 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<T> From<Vec<T>> for Arc<[T]> {
     /// assert_eq!(&[1, 2, 3], &shared[..]);
     /// ```
     #[inline]
-    fn from(mut v: Vec<T>) -> 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<T>) -> 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
@@ -1023,11 +1023,9 @@ impl<T, A: Allocator> Vec<T, A> {
     #[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<T, A: Allocator> Vec<T, A> {
     /// ```
     #[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<T, A: Allocator> Vec<T, A> {
     #[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<T, A: Allocator> Vec<T, A> {
     #[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, F: FnMut() -> T> ExtendWith<T> for ExtendFunc<F> {
 impl<T, A: Allocator> Vec<T, A> {
     #[cfg(not(no_global_oom_handling))]
     /// Extend the vector by `n` values, using the given generator.
-    fn extend_with<E: ExtendWith<T>>(&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<E: ExtendWith<T>>(&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<T, A: Allocator> Vec<T, A> {
     // 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<I: Iterator<Item = T>>(&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<I: Iterator<Item = T>>(&mut self, iterator: I) {
+        match self.try_extend_desugared(iterator) {
+            Ok(r) => r,
+            Err(e) => e.handle(),
         }
     }
 
