Cleanup

Author: saethlin

library/alloc/src/raw_vec.rs

@@ -83,59 +83,11 @@ struct MonoRawVec<A: Allocator = Global> {
 }
 
 #[inline]
-const fn max_size_for_align(align: usize) -> usize {
-    // (power-of-two implies align != 0.)
-
-    // Rounded up size is:
-    //   size_rounded_up = (size + align - 1) & !(align - 1);
-    //
-    // We know from above that align != 0. If adding (align - 1)
-    // does not overflow, then rounding up will be fine.
-    //
-    // Conversely, &-masking with !(align - 1) will subtract off
-    // only low-order-bits. Thus if overflow occurs with the sum,
-    // the &-mask cannot subtract enough to undo that overflow.
-    //
-    // Above implies that checking for summation overflow is both
-    // necessary and sufficient.
-    isize::MAX as usize - (align - 1)
-}
-
-#[inline]
-fn layout_array(cap: usize, size: usize, align: usize) -> Result<Layout, TryReserveError> {
-    // We need to check two things about the size:
-    //  - That the total size won't overflow a `usize`, and
-    //  - That the total size still fits in an `isize`.
-    // By using division we can check them both with a single threshold.
-    // That'd usually be a bad idea, but thankfully here the element size
-    // and alignment are constants, so the compiler will fold all of it.
-    if size != 0 && cap > max_size_for_align(align) / size {
-        return Err(CapacityOverflow.into());
-    }
-
-    // SAFETY: We just checked that we won't overflow `usize` when we multiply.
-    // This is a useless hint inside this function, but after inlining this helps
-    // deduplicate checks for whether the overall capacity is zero (e.g., in RawVec's
-    // allocation path) before/after this multiplication.
-    let array_size = unsafe { size.unchecked_mul(cap) };
-
-    // SAFETY: We just checked above that the `array_size` will not
-    // exceed `isize::MAX` even when rounded up to the alignment.
-    // And `Alignment` guarantees it's a power of two.
-    unsafe { Ok(Layout::from_size_align_unchecked(array_size, align)) }
+fn layout_array(cap: usize, elem_layout: Layout) -> Result<Layout, TryReserveError> {
+    elem_layout.repeat(cap).map(|(layout, _pad)| layout).map_err(|_| CapacityOverflow.into())
 }
 
 impl<A: Allocator> MonoRawVec<A> {
-    #[inline]
-    fn ptr<T>(&self) -> *mut T {
-        unsafe { core::mem::transmute(self.ptr) }
-    }
-
-    #[inline]
-    fn non_null<T>(&self) -> NonNull<T> {
-        unsafe { core::mem::transmute(self.ptr) }
-    }
-
     /// # Safety:
     ///
     /// `cap` must not exceed `isize::MAX`.
@@ -149,8 +101,8 @@ impl<A: Allocator> MonoRawVec<A> {
     }
 
     #[inline]
-    fn current_memory(&self, size: usize, align: usize) -> Option<(NonNull<u8>, Layout)> {
-        if size == 0 || self.cap.0 == 0 {
+    fn current_memory(&self, elem_layout: Layout) -> Option<(NonNull<u8>, Layout)> {
+        if elem_layout.size() == 0 || self.cap.0 == 0 {
             return None;
         }
 
@@ -159,21 +111,19 @@ impl<A: Allocator> MonoRawVec<A> {
         // has already been allocated so we know it can't overflow and currently Rust does not
         // support such types. So we can do better by skipping some checks and avoid an unwrap.
         unsafe {
-            let size = size.unchecked_mul(self.cap.0);
-            let layout = Layout::from_size_align_unchecked(size, align);
+            let alloc_size = elem_layout.size().unchecked_mul(self.cap.0);
+            let layout = Layout::from_size_align_unchecked(alloc_size, elem_layout.align());
             Some((self.ptr.into(), layout))
         }
     }
 
-    #[inline]
     fn grow_amortized(
         &mut self,
         len: usize,
         additional: usize,
-        size: usize,
-        align: usize,
+        elem_layout: Layout,
     ) -> Result<(), TryReserveError> {
-        if size == 0 {
+        if elem_layout.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.into());
@@ -185,36 +135,34 @@ impl<A: Allocator> MonoRawVec<A> {
         // 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.0 * 2, required_cap);
-        let cap = cmp::max(min_non_zero_cap(size), cap);
+        let cap = cmp::max(min_non_zero_cap(elem_layout.size()), cap);
 
-        let new_layout = layout_array(cap, size, align)?;
+        let new_layout = layout_array(cap, elem_layout)?;
 
         // `finish_grow` is non-generic over `T`.
-        let ptr = finish_grow(new_layout, self.current_memory(size, align), &mut self.alloc)?;
+        let ptr = finish_grow(new_layout, self.current_memory(elem_layout), &mut self.alloc)?;
         // SAFETY: finish_grow would have resulted in a capacity overflow if we tried to allocate more than isize::MAX items
         unsafe { self.set_ptr_and_cap(ptr, cap) };
         Ok(())
     }
 
-    #[inline]
     fn grow_exact(
         &mut self,
         len: usize,
         additional: usize,
-        size: usize,
-        align: usize,
+        elem_layout: Layout,
     ) -> Result<(), TryReserveError> {
-        if size == 0 {
+        if elem_layout.size() == 0 {
             // Since we return a capacity of `usize::MAX` when the type size is
             // 0, getting to here necessarily means the `RawVec` is overfull.
             return Err(CapacityOverflow.into());
         }
 
         let cap = len.checked_add(additional).ok_or(CapacityOverflow)?;
-        let new_layout = layout_array(cap, size, align)?;
+        let new_layout = layout_array(cap, elem_layout)?;
 
         // `finish_grow` is non-generic over `T`.
-        let ptr = finish_grow(new_layout, self.current_memory(size, align), &mut self.alloc)?;
+        let ptr = finish_grow(new_layout, self.current_memory(elem_layout), &mut self.alloc)?;
         // SAFETY: finish_grow would have resulted in a capacity overflow if we tried to allocate more than isize::MAX items
         unsafe {
             self.set_ptr_and_cap(ptr, cap);
@@ -223,27 +171,27 @@ impl<A: Allocator> MonoRawVec<A> {
     }
 
     #[cfg(not(no_global_oom_handling))]
-    #[inline]
-    fn shrink(&mut self, cap: usize, size: usize, align: usize) -> Result<(), TryReserveError> {
-        if size > 0 {
+    fn shrink(&mut self, cap: usize, elem_layout: Layout) -> Result<(), TryReserveError> {
+        if elem_layout.size() > 0 {
             assert!(cap <= self.cap.0, "Tried to shrink to a larger capacity");
         }
 
         let (ptr, layout) =
-            if let Some(mem) = self.current_memory(size, align) { mem } else { return Ok(()) };
+            if let Some(mem) = self.current_memory(elem_layout) { mem } else { return Ok(()) };
 
         // If shrinking to 0, deallocate the buffer. We don't reach this point
         // for the T::IS_ZST case since current_memory() will have returned
         // None.
         if cap == 0 {
             unsafe { self.alloc.deallocate(ptr, layout) };
-            self.ptr = unsafe { Unique::new_unchecked(ptr::without_provenance_mut(align)) };
+            self.ptr =
+                unsafe { Unique::new_unchecked(ptr::without_provenance_mut(elem_layout.align())) };
             self.cap = Cap::ZERO;
         } else {
             let ptr = unsafe {
                 // `Layout::array` cannot overflow here because it would have
                 // overflowed earlier when capacity was larger.
-                let new_size = size.unchecked_mul(cap);
+                let new_size = elem_layout.size().unchecked_mul(cap);
                 let new_layout = Layout::from_size_align_unchecked(new_size, layout.align());
                 self.alloc
                     .shrink(ptr, layout, new_layout)
@@ -264,21 +212,19 @@ impl<A: Allocator> MonoRawVec<A> {
         Self { ptr, cap: Cap::ZERO, alloc }
     }
 
-    #[inline]
     fn try_allocate_in(
         capacity: usize,
         init: AllocInit,
         alloc: A,
-        size: usize,
-        align: usize,
+        elem_layout: Layout,
     ) -> Result<Self, TryReserveError> {
         // Don't allocate here because `Drop` will not deallocate when `capacity` is 0.
-        if size == 0 || capacity == 0 {
-            Ok(Self::new_in(alloc, align))
+        if elem_layout.size() == 0 || capacity == 0 {
+            Ok(Self::new_in(alloc, elem_layout.align()))
         } else {
             // We avoid `unwrap_or_else` here because it bloats the amount of
             // LLVM IR generated.
-            let layout = match layout_array(capacity, size, align) {
+            let layout = match layout_array(capacity, elem_layout) {
                 Ok(layout) => layout,
                 Err(_) => return Err(CapacityOverflow.into()),
             };
@@ -304,10 +250,8 @@ impl<A: Allocator> MonoRawVec<A> {
         }
     }
 
-    #[inline]
-    #[rustc_no_mir_inline]
-    fn drop_if_needed(&mut self, size: usize, align: usize) {
-        if let Some((ptr, layout)) = self.current_memory(size, align) {
+    fn drop_if_needed(&mut self, elem_layout: Layout) {
+        if let Some((ptr, layout)) = self.current_memory(elem_layout) {
             unsafe { self.alloc.deallocate(ptr, layout) }
         }
     }
@@ -449,13 +393,7 @@ impl<T, A: Allocator> RawVec<T, A> {
         init: AllocInit,
         alloc: A,
     ) -> Result<Self, TryReserveError> {
-        match MonoRawVec::try_allocate_in(
-            capacity,
-            init,
-            alloc,
-            mem::size_of::<T>(),
-            mem::align_of::<T>(),
-        ) {
+        match MonoRawVec::try_allocate_in(capacity, init, alloc, T::LAYOUT) {
             Ok(inner) => Ok(Self { inner, _phantom: PhantomData }),
             Err(e) => Err(e),
         }
@@ -499,23 +437,24 @@ impl<T, A: Allocator> RawVec<T, A> {
     /// be careful.
     #[inline]
     pub fn ptr(&self) -> *mut T {
-        self.inner.ptr()
+        unsafe { core::mem::transmute(self.inner.ptr) }
     }
 
     #[inline]
     pub fn non_null(&self) -> NonNull<T> {
-        self.inner.non_null()
+        unsafe { core::mem::transmute(self.inner.ptr) }
     }
 
     /// Gets the capacity of the allocation.
     ///
     /// This will always be `usize::MAX` if `T` is zero-sized.
-    #[inline(always)]
+    #[inline]
     pub fn capacity(&self) -> usize {
         if T::IS_ZST { usize::MAX } else { self.inner.cap.0 }
     }
 
     /// Returns a shared reference to the allocator backing this `RawVec`.
+    #[inline]
     pub fn allocator(&self) -> &A {
         &self.inner.alloc
     }
@@ -645,6 +584,7 @@ impl<T, A: Allocator> RawVec<T, A> {
 impl<T, A: Allocator> RawVec<T, A> {
     /// Returns if the buffer needs to grow to fulfill the needed extra capacity.
     /// Mainly used to make inlining reserve-calls possible without inlining `grow`.
+    #[inline]
     fn needs_to_grow(&self, len: usize, additional: usize) -> bool {
         additional > self.capacity().wrapping_sub(len)
     }
@@ -656,27 +596,25 @@ impl<T, A: Allocator> RawVec<T, A> {
     // so that all of the code that depends on `T` is within it, while as much
     // of the code that doesn't depend on `T` as possible is in functions that
     // are non-generic over `T`.
+    #[inline]
     fn grow_amortized(&mut self, len: usize, additional: usize) -> Result<(), TryReserveError> {
-        const { assert!(mem::size_of::<T>() % mem::align_of::<T>() == 0) };
-
         // This is ensured by the calling contexts.
         debug_assert!(additional > 0);
 
-        self.inner.grow_amortized(len, additional, mem::size_of::<T>(), mem::align_of::<T>())
+        self.inner.grow_amortized(len, additional, T::LAYOUT)
     }
 
     // The constraints on this method are much the same as those on
     // `grow_amortized`, but this method is usually instantiated less often so
     // it's less critical.
+    #[inline]
     fn grow_exact(&mut self, len: usize, additional: usize) -> Result<(), TryReserveError> {
-        const { assert!(mem::size_of::<T>() % mem::align_of::<T>() == 0) };
-        self.inner.grow_exact(len, additional, mem::size_of::<T>(), mem::align_of::<T>())
+        self.inner.grow_exact(len, additional, T::LAYOUT)
     }
 
     #[cfg(not(no_global_oom_handling))]
     fn shrink(&mut self, cap: usize) -> Result<(), TryReserveError> {
-        const { assert!(mem::size_of::<T>() % mem::align_of::<T>() == 0) };
-        self.inner.shrink(cap, mem::size_of::<T>(), mem::align_of::<T>())
+        self.inner.shrink(cap, T::LAYOUT)
     }
 }
 
@@ -712,8 +650,7 @@ where
 unsafe impl<#[may_dangle] T, A: Allocator> Drop for RawVec<T, A> {
     /// Frees the memory owned by the `RawVec` *without* trying to drop its contents.
     fn drop(&mut self) {
-        const { assert!(mem::size_of::<T>() % mem::align_of::<T>() == 0) };
-        self.inner.drop_if_needed(mem::size_of::<T>(), mem::align_of::<T>());
+        self.inner.drop_if_needed(T::LAYOUT);
     }
 }
 

library/core/src/mem/mod.rs

@@ -5,6 +5,7 @@
 
 #![stable(feature = "rust1", since = "1.0.0")]
 
+use crate::alloc::Layout;
 use crate::clone;
 use crate::cmp;
 use crate::fmt;
@@ -1235,6 +1236,12 @@ pub trait SizedTypeProperties: Sized {
     #[doc(hidden)]
     #[unstable(feature = "sized_type_properties", issue = "none")]
     const IS_ZST: bool = size_of::<Self>() == 0;
+
+    #[doc(hidden)]
+    #[unstable(feature = "sized_type_properties", issue = "none")]
+    const LAYOUT: Layout =
+        // SAFETY: By construction
+        unsafe { Layout::from_size_align_unchecked(size_of::<Self>(), align_of::<Self>()) };
 }
 #[doc(hidden)]
 #[unstable(feature = "sized_type_properties", issue = "none")]