Faster Layout::array

library/alloc/src/raw_vec/tests.rs

@@ -77,3 +77,87 @@ fn reserve_does_not_overallocate() {
         assert!(v.capacity() >= 12 + 12 / 2);
     }
 }
+
+struct ZST;
+
+// A `RawVec` holding zero-sized elements should always look like this.
+fn zst_sanity<T>(v: &RawVec<T>) {
+    assert_eq!(v.capacity(), usize::MAX);
+    assert_eq!(v.ptr(), core::ptr::Unique::<T>::dangling().as_ptr());
+    assert_eq!(v.current_memory(), None);
+}
+
+#[test]
+fn zst() {
+    let cap_err = Err(crate::collections::TryReserveErrorKind::CapacityOverflow.into());
+
+    assert_eq!(std::mem::size_of::<ZST>(), 0);
+
+    // All these different ways of creating the RawVec produce the same thing.
+
+    let v: RawVec<ZST> = RawVec::new();
+    zst_sanity(&v);
+
+    let v: RawVec<ZST> = RawVec::with_capacity_in(100, Global);
+    zst_sanity(&v);
+
+    let v: RawVec<ZST> = RawVec::with_capacity_in(100, Global);
+    zst_sanity(&v);
+
+    let v: RawVec<ZST> = RawVec::allocate_in(0, AllocInit::Uninitialized, Global);
+    zst_sanity(&v);
+
+    let v: RawVec<ZST> = RawVec::allocate_in(100, AllocInit::Uninitialized, Global);
+    zst_sanity(&v);
+
+    let mut v: RawVec<ZST> = RawVec::allocate_in(usize::MAX, AllocInit::Uninitialized, Global);
+    zst_sanity(&v);
+
+    // Check all these operations work as expected with zero-sized elements.
+
+    assert!(!v.needs_to_grow(100, usize::MAX - 100));
+    assert!(v.needs_to_grow(101, usize::MAX - 100));
+    zst_sanity(&v);
+
+    v.reserve(100, usize::MAX - 100);
+    //v.reserve(101, usize::MAX - 100); // panics, in `zst_reserve_panic` below
+    zst_sanity(&v);
+
+    v.reserve_exact(100, usize::MAX - 100);
+    //v.reserve_exact(101, usize::MAX - 100); // panics, in `zst_reserve_exact_panic` below
+    zst_sanity(&v);
+
+    assert_eq!(v.try_reserve(100, usize::MAX - 100), Ok(()));
+    assert_eq!(v.try_reserve(101, usize::MAX - 100), cap_err);
+    zst_sanity(&v);
+
+    assert_eq!(v.try_reserve_exact(100, usize::MAX - 100), Ok(()));
+    assert_eq!(v.try_reserve_exact(101, usize::MAX - 100), cap_err);
+    zst_sanity(&v);
+
+    assert_eq!(v.grow_amortized(100, usize::MAX - 100), cap_err);
+    assert_eq!(v.grow_amortized(101, usize::MAX - 100), cap_err);
+    zst_sanity(&v);
+
+    assert_eq!(v.grow_exact(100, usize::MAX - 100), cap_err);
+    assert_eq!(v.grow_exact(101, usize::MAX - 100), cap_err);
+    zst_sanity(&v);
+}
+
+#[test]
+#[should_panic(expected = "capacity overflow")]
+fn zst_reserve_panic() {
+    let mut v: RawVec<ZST> = RawVec::new();
+    zst_sanity(&v);
+
+    v.reserve(101, usize::MAX - 100);
+}
+
+#[test]
+#[should_panic(expected = "capacity overflow")]
+fn zst_reserve_exact_panic() {
+    let mut v: RawVec<ZST> = RawVec::new();
+    zst_sanity(&v);
+
+    v.reserve_exact(101, usize::MAX - 100);
+}

library/core/src/alloc/layout.rs

@@ -281,7 +281,9 @@ impl Layout {
         // > `usize::MAX`)
         let new_size = self.size() + pad;
 
-        Layout::from_size_align(new_size, self.align()).unwrap()
+        // SAFETY: self.align is already known to be valid and new_size has been
+        // padded already.
+        unsafe { Layout::from_size_align_unchecked(new_size, self.align()) }
     }
 
     /// Creates a layout describing the record for `n` instances of
@@ -403,9 +405,17 @@ impl Layout {
     #[stable(feature = "alloc_layout_manipulation", since = "1.44.0")]
     #[inline]
     pub fn array<T>(n: usize) -> Result<Self, LayoutError> {
-        let (layout, offset) = Layout::new::<T>().repeat(n)?;
-        debug_assert_eq!(offset, mem::size_of::<T>());
-        Ok(layout.pad_to_align())
+        let array_size = mem::size_of::<T>().checked_mul(n).ok_or(LayoutError)?;
+
+        // SAFETY:
+        // - Size: `array_size` cannot be too big because `size_of::<T>()` must
+        //   be a multiple of `align_of::<T>()`. Therefore, `array_size`
+        //   rounded up to the nearest multiple of `align_of::<T>()` is just
+        //   `array_size`. And `array_size` cannot be too big because it was
+        //   just checked by the `checked_mul()`.
+        // - Alignment: `align_of::<T>()` will always give an acceptable
+        //   (non-zero, power of two) alignment.
+        Ok(unsafe { Layout::from_size_align_unchecked(array_size, mem::align_of::<T>()) })
     }
 }