Shrink grow_amortized() some more

src/liballoc/raw_vec.rs

@@ -1,7 +1,7 @@
 #![unstable(feature = "raw_vec_internals", reason = "implementation detail", issue = "none")]
 #![doc(hidden)]
 
-use core::alloc::{LayoutErr, MemoryBlock};
+use core::alloc::MemoryBlock;
 use core::cmp;
 use core::mem::{self, ManuallyDrop, MaybeUninit};
 use core::ops::Drop;
@@ -172,7 +172,8 @@ impl<T, A: AllocRef> RawVec<T, A> {
         if mem::size_of::<T>() == 0 {
             Self::new_in(alloc)
         } else {
-            let layout = Layout::array::<T>(capacity).unwrap_or_else(|_| capacity_overflow());
+            let layout =
+                array_layout(Layout::new::<T>(), capacity).unwrap_or_else(|_| capacity_overflow());
             alloc_guard(layout.size()).unwrap_or_else(|_| capacity_overflow());
 
             let memory = alloc.alloc(layout, init).unwrap_or_else(|_| handle_alloc_error(layout));
@@ -290,10 +291,14 @@ impl<T, A: AllocRef> RawVec<T, A> {
     /// # }
     /// ```
     pub fn reserve(&mut self, len: usize, additional: usize) {
-        match self.try_reserve(len, additional) {
-            Err(CapacityOverflow) => capacity_overflow(),
-            Err(AllocError { layout, .. }) => handle_alloc_error(layout),
-            Ok(()) => { /* yay */ }
+        // This function is marginally shorter if it calls `try_reserve`, but
+        // that results in more LLVM IR being generated.
+        if self.needs_to_grow(len, additional) {
+            match self.grow_amortized(len, additional) {
+                Ok(()) => { /* yay */ }
+                Err(CapacityOverflow) => capacity_overflow(),
+                Err(AllocError { layout, .. }) => handle_alloc_error(layout),
+            }
         }
     }
 
@@ -326,10 +331,14 @@ impl<T, A: AllocRef> RawVec<T, A> {
     ///
     /// Aborts on OOM.
     pub fn reserve_exact(&mut self, len: usize, additional: usize) {
-        match self.try_reserve_exact(len, additional) {
-            Err(CapacityOverflow) => capacity_overflow(),
-            Err(AllocError { layout, .. }) => handle_alloc_error(layout),
-            Ok(()) => { /* yay */ }
+        // This function is marginally shorter if it calls `try_reserve_exact`,
+        // but that results in more LLVM IR being generated.
+        if self.needs_to_grow(len, additional) {
+            match self.grow_exact(len, additional) {
+                Ok(()) => { /* yay */ }
+                Err(CapacityOverflow) => capacity_overflow(),
+                Err(AllocError { layout, .. }) => handle_alloc_error(layout),
+            }
         }
     }
 
@@ -386,42 +395,13 @@ impl<T, A: AllocRef> RawVec<T, A> {
     // of the code that doesn't depend on `T` as possible is in functions that
     // are non-generic over `T`.
     fn grow_amortized(&mut self, len: usize, additional: usize) -> Result<(), TryReserveError> {
-        // This is ensured by the calling contexts.
-        debug_assert!(additional > 0);
-
-        if mem::size_of::<T>() == 0 {
-            // Since we return a capacity of `usize::MAX` when `elem_size` is
-            // 0, getting to here necessarily means the `RawVec` is overfull.
-            return Err(CapacityOverflow);
-        }
-
-        // Nothing we can really do about these checks, sadly.
-        let required_cap = len.checked_add(additional).ok_or(CapacityOverflow)?;
-
-        // This guarantees exponential growth. The doubling cannot overflow
-        // because `cap <= isize::MAX` and the type of `cap` is `usize`.
-        let cap = cmp::max(self.cap * 2, required_cap);
-
-        // Tiny Vecs are dumb. Skip to:
-        // - 8 if the element size is 1, because any heap allocators is likely
-        //   to round up a request of less than 8 bytes to at least 8 bytes.
-        // - 4 if elements are moderate-sized (<= 1 KiB).
-        // - 1 otherwise, to avoid wasting too much space for very short Vecs.
-        // Note that `min_non_zero_cap` is computed statically.
         let elem_size = mem::size_of::<T>();
-        let min_non_zero_cap = if elem_size == 1 {
-            8
-        } else if elem_size <= 1024 {
-            4
-        } else {
-            1
-        };
-        let cap = cmp::max(min_non_zero_cap, cap);
+        let cap = cap_amortized(elem_size, len, additional, self.cap)?;
 
-        let new_layout = Layout::array::<T>(cap);
+        let elem_layout = Layout::new::<T>();
 
         // `finish_grow` is non-generic over `T`.
-        let memory = finish_grow(new_layout, self.current_memory(), &mut self.alloc)?;
+        let memory = finish_grow(cap, elem_layout, self.current_memory(), &mut self.alloc)?;
         self.set_memory(memory);
         Ok(())
     }
@@ -437,10 +417,11 @@ impl<T, A: AllocRef> RawVec<T, A> {
         }
 
         let cap = len.checked_add(additional).ok_or(CapacityOverflow)?;
-        let new_layout = Layout::array::<T>(cap);
+
+        let elem_layout = Layout::new::<T>();
 
         // `finish_grow` is non-generic over `T`.
-        let memory = finish_grow(new_layout, self.current_memory(), &mut self.alloc)?;
+        let memory = finish_grow(cap, elem_layout, self.current_memory(), &mut self.alloc)?;
         self.set_memory(memory);
         Ok(())
     }
@@ -468,20 +449,61 @@ impl<T, A: AllocRef> RawVec<T, A> {
     }
 }
 
+// This function is outside `RawVec` to minimize compile times. See the comment
+// above `RawVec::grow_amortized` for details.
+#[inline]
+fn cap_amortized(
+    elem_size: usize,
+    len: usize,
+    additional: usize,
+    curr_cap: usize,
+) -> Result<usize, TryReserveError> {
+    // This is ensured by the calling contexts.
+    debug_assert!(additional > 0);
+
+    // Tiny Vecs are dumb. Skip to:
+    // - 8 if the element size is 1, because any heap allocators is likely
+    //   to round up a request of less than 8 bytes to at least 8 bytes.
+    // - 4 if elements are moderate-sized (<= 1 KiB).
+    // - 1 otherwise, to avoid wasting too much space for very short Vecs.
+    // Note that `min_non_zero_cap` is computed statically.
+    let min_non_zero_cap = if elem_size == 0 {
+        // Since we return a capacity of `usize::MAX` when `elem_size` is
+        // 0, getting to here necessarily means the `RawVec` is overfull.
+        return Err(CapacityOverflow);
+    } else if elem_size == 1 {
+        8
+    } else if elem_size <= 1024 {
+        4
+    } else {
+        1
+    };
+
+    // Nothing we can really do about these checks, sadly.
+    let required_cap = len.checked_add(additional).ok_or(CapacityOverflow)?;
+
+    // This guarantees exponential growth. The doubling cannot overflow
+    // because `cap <= isize::MAX` and the type of `cap` is `usize`.
+    let cap = cmp::max(curr_cap * 2, required_cap);
+
+    Ok(cmp::max(min_non_zero_cap, cap))
+}
+
 // This function is outside `RawVec` to minimize compile times. See the comment
 // above `RawVec::grow_amortized` for details. (The `A` parameter isn't
 // significant, because the number of different `A` types seen in practice is
 // much smaller than the number of `T` types.)
+#[inline]
 fn finish_grow<A>(
-    new_layout: Result<Layout, LayoutErr>,
+    cap: usize,
+    elem_layout: Layout,
     current_memory: Option<(NonNull<u8>, Layout)>,
     alloc: &mut A,
 ) -> Result<MemoryBlock, TryReserveError>
 where
     A: AllocRef,
 {
-    // Check for the error here to minimize the size of `RawVec::grow_*`.
-    let new_layout = new_layout.map_err(|_| CapacityOverflow)?;
+    let new_layout = array_layout(elem_layout, cap)?;
 
     alloc_guard(new_layout.size())?;
 
@@ -496,6 +518,15 @@ where
     Ok(memory)
 }
 
+// This is equivalent to Layout::array, but is non-generic and has a different
+// error type in its result. It helps reduce the amount of LLVM IR generated.
+#[inline]
+fn array_layout(elem_layout: Layout, n: usize) -> Result<Layout, TryReserveError> {
+    let (new_layout, offset) = elem_layout.repeat(n).map_err(|_| CapacityOverflow)?;
+    debug_assert_eq!(offset, elem_layout.size());
+    Ok(new_layout.pad_to_align())
+}
+
 unsafe impl<#[may_dangle] T, A: AllocRef> Drop for RawVec<T, A> {
     /// Frees the memory owned by the `RawVec` *without* trying to drop its contents.
     fn drop(&mut self) {