From 35043aedfbb1c2237df470d386ccc17ad03a43e7 Mon Sep 17 00:00:00 2001
From: Joshua Liebow-Feeser <hello@joshlf.com>
Date: Mon, 1 Jul 2024 14:02:36 -0700
Subject: [PATCH] FromZeros boxed slice method supports slice DSTs

Makes progress on #29
---
 src/error.rs |  4 ++++
 src/lib.rs   | 37 +++++++++++++++++++------------------
 2 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index bca4130ffd7..a2f636bfd40 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -570,6 +570,10 @@ impl<Src, Dst: ?Sized + TryFromBytes> TryReadError<Src, Dst> {
     }
 }
 
+/// The error type of a failed allocation.
+#[derive(Copy, Clone, PartialEq, Eq, Debug)]
+pub struct AllocError;
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/src/lib.rs b/src/lib.rs
index 7c3020597bd..642f1ec6b51 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2189,14 +2189,16 @@ pub unsafe trait FromZeros: TryFromBytes {
     #[cfg(feature = "alloc")]
     #[cfg_attr(doc_cfg, doc(cfg(feature = "alloc")))]
     #[inline]
-    fn new_box_slice_zeroed(len: usize) -> Box<[Self]>
+    fn new_box_zeroed_with_elems(count: usize) -> Option<Box<Self>>
     where
-        Self: Sized,
+        Self: KnownLayout<PointerMetadata = usize>,
     {
-        let size = mem::size_of::<Self>()
-            .checked_mul(len)
-            .expect("mem::size_of::<Self>() * len overflows `usize`");
-        let align = mem::align_of::<Self>();
+        let size = match count.size_for_metadata(Self::LAYOUT) {
+            Some(size) => size,
+            None => return None,
+        };
+
+        let align = Self::LAYOUT.align.get();
         // On stable Rust versions <= 1.64.0, `Layout::from_size_align` has a
         // bug in which sufficiently-large allocations (those which, when
         // rounded up to the alignment, overflow `isize`) are not rejected,
@@ -2208,29 +2210,28 @@ pub unsafe trait FromZeros: TryFromBytes {
         assert!(size <= max_alloc);
         // TODO(https://github.com/rust-lang/rust/issues/55724): Use
         // `Layout::repeat` once it's stabilized.
-        let layout =
-            Layout::from_size_align(size, align).expect("total allocation size overflows `isize`");
+        let layout = Layout::from_size_align(size, align).ok()?;
 
         let ptr = if layout.size() != 0 {
             // TODO(#429): Add a "SAFETY" comment and remove this `allow`.
             #[allow(clippy::undocumented_unsafe_blocks)]
-            let ptr = unsafe { alloc::alloc::alloc_zeroed(layout).cast::<Self>() };
-            if ptr.is_null() {
-                alloc::alloc::handle_alloc_error(layout);
+            let ptr = unsafe { alloc::alloc::alloc_zeroed(layout) };
+            match NonNull::new(ptr) {
+                Some(ptr) => ptr,
+                None => alloc::alloc::handle_alloc_error(layout),
             }
-            ptr
         } else {
             // `Box<[T]>` does not allocate when `T` is zero-sized or when `len`
             // is zero, but it does require a non-null dangling pointer for its
             // allocation.
-            NonNull::<Self>::dangling().as_ptr()
+            NonNull::<u8>::dangling()
         };
 
+        let ptr = Self::raw_from_ptr_len(ptr, count);
+
         // TODO(#429): Add a "SAFETY" comment and remove this `allow`.
         #[allow(clippy::undocumented_unsafe_blocks)]
-        unsafe {
-            Box::from_raw(slice::from_raw_parts_mut(ptr, len))
-        }
+        Some(unsafe { Box::from_raw(ptr.as_ptr()) })
     }
 
     /// Creates a `Vec<Self>` from zeroed bytes.
@@ -2257,11 +2258,11 @@ pub unsafe trait FromZeros: TryFromBytes {
     #[cfg(feature = "alloc")]
     #[cfg_attr(doc_cfg, doc(cfg(feature = "alloc")))]
     #[inline(always)]
-    fn new_vec_zeroed(len: usize) -> Vec<Self>
+    fn new_vec_zeroed(len: usize) -> Result<Vec<Self>, AllocError>
     where
         Self: Sized,
     {
-        Self::new_box_slice_zeroed(len).into()
+        <[Self]>::new_box_zeroed_with_elems(len).map(Into::into).ok_or(AllocError)
     }
 }