Add internal collect_into_array[_unchecked] to remove duplicate code

This does not suggest adding such a function to the public API. This is
just for the purpose of avoiding duplicate code. Many array methods
already contained the same kind of code and there are still many array
related methods to come (e.g. `Iterator::{chunks, map_windows, next_n,
...}`) which all basically need this functionality. Writing custom
`unsafe` code for each of those seems not like a good idea.

Author: LukasKalbertodt

library/core/src/array/mod.rs

@@ -11,8 +11,9 @@ use crate::cmp::Ordering;
 use crate::convert::{Infallible, TryFrom};
 use crate::fmt;
 use crate::hash::{self, Hash};
+use crate::iter::TrustedLen;
 use crate::marker::Unsize;
-use crate::mem::MaybeUninit;
+use crate::mem::{self, MaybeUninit};
 use crate::ops::{Index, IndexMut};
 use crate::slice::{Iter, IterMut};
 
@@ -426,41 +427,13 @@ impl<T, const N: usize> [T; N] {
     /// assert_eq!(y, [6, 9, 3, 3]);
     /// ```
     #[unstable(feature = "array_map", issue = "75243")]
-    pub fn map<F, U>(self, mut f: F) -> [U; N]
+    pub fn map<F, U>(self, f: F) -> [U; N]
     where
         F: FnMut(T) -> U,
     {
-        struct Guard<T, const N: usize> {
-            dst: *mut T,
-            initialized: usize,
-        }
-
-        impl<T, const N: usize> Drop for Guard<T, N> {
-            fn drop(&mut self) {
-                debug_assert!(self.initialized <= N);
-
-                let initialized_part =
-                    crate::ptr::slice_from_raw_parts_mut(self.dst, self.initialized);
-                // SAFETY: this raw slice will contain only initialized objects
-                // that's why, it is allowed to drop it.
-                unsafe {
-                    crate::ptr::drop_in_place(initialized_part);
-                }
-            }
-        }
-        let mut dst = MaybeUninit::uninit_array::<N>();
-        let mut guard: Guard<U, N> =
-            Guard { dst: MaybeUninit::slice_as_mut_ptr(&mut dst), initialized: 0 };
-        for (src, dst) in IntoIter::new(self).zip(&mut dst) {
-            dst.write(f(src));
-            guard.initialized += 1;
-        }
-        // FIXME: Convert to crate::mem::transmute once it works with generics.
-        // unsafe { crate::mem::transmute::<[MaybeUninit<U>; N], [U; N]>(dst) }
-        crate::mem::forget(guard);
-        // SAFETY: At this point we've properly initialized the whole array
-        // and we just need to cast it to the correct type.
-        unsafe { crate::mem::transmute_copy::<_, [U; N]>(&dst) }
+        // SAFETY: we know for certain that this iterator will yield exactly `N`
+        // items.
+        unsafe { collect_into_array_unchecked(&mut IntoIter::new(self).map(f)) }
     }
 
     /// 'Zips up' two arrays into a single array of pairs.
@@ -481,15 +454,11 @@ impl<T, const N: usize> [T; N] {
     /// ```
     #[unstable(feature = "array_zip", issue = "80094")]
     pub fn zip<U>(self, rhs: [U; N]) -> [(T, U); N] {
-        let mut dst = MaybeUninit::uninit_array::<N>();
-        for (i, (lhs, rhs)) in IntoIter::new(self).zip(IntoIter::new(rhs)).enumerate() {
-            dst[i].write((lhs, rhs));
-        }
-        // FIXME: Convert to crate::mem::transmute once it works with generics.
-        // unsafe { crate::mem::transmute::<[MaybeUninit<U>; N], [U; N]>(dst) }
-        // SAFETY: At this point we've properly initialized the whole array
-        // and we just need to cast it to the correct type.
-        unsafe { crate::mem::transmute_copy::<_, [(T, U); N]>(&dst) }
+        let mut iter = IntoIter::new(self).zip(IntoIter::new(rhs));
+
+        // SAFETY: we know for certain that this iterator will yield exactly `N`
+        // items.
+        unsafe { collect_into_array_unchecked(&mut iter) }
     }
 
     /// Returns a slice containing the entire array. Equivalent to `&s[..]`.
@@ -535,16 +504,9 @@ impl<T, const N: usize> [T; N] {
     /// ```
     #[unstable(feature = "array_methods", issue = "76118")]
     pub fn each_ref(&self) -> [&T; N] {
-        // Unlike in `map`, we don't need a guard here, as dropping a reference
-        // is a noop.
-        let mut out = MaybeUninit::uninit_array::<N>();
-        for (src, dst) in self.iter().zip(&mut out) {
-            dst.write(src);
-        }
-
-        // SAFETY: All elements of `dst` are properly initialized and
-        // `MaybeUninit<T>` has the same layout as `T`, so this cast is valid.
-        unsafe { (&mut out as *mut _ as *mut [&T; N]).read() }
+        // SAFETY: we know for certain that this iterator will yield exactly `N`
+        // items.
+        unsafe { collect_into_array_unchecked(&mut self.iter()) }
     }
 
     /// Borrows each element mutably and returns an array of mutable references
@@ -564,15 +526,103 @@ impl<T, const N: usize> [T; N] {
     /// ```
     #[unstable(feature = "array_methods", issue = "76118")]
     pub fn each_mut(&mut self) -> [&mut T; N] {
-        // Unlike in `map`, we don't need a guard here, as dropping a reference
-        // is a noop.
-        let mut out = MaybeUninit::uninit_array::<N>();
-        for (src, dst) in self.iter_mut().zip(&mut out) {
-            dst.write(src);
+        // SAFETY: we know for certain that this iterator will yield exactly `N`
+        // items.
+        unsafe { collect_into_array_unchecked(&mut self.iter_mut()) }
+    }
+}
+
+/// Pulls `N` items from `iter` and returns them as an array. If the iterator
+/// yields fewer than `N` items, this function exhibits undefined behavior.
+///
+/// See [`collect_into_array`] for more information.
+///
+///
+/// # Safety
+///
+/// It is up to the caller to guarantee that `iter` yields at least `N` items.
+/// Violating this condition causes undefined behavior.
+unsafe fn collect_into_array_unchecked<I, const N: usize>(iter: &mut I) -> [I::Item; N]
+where
+    // Note: `TrustedLen` here is somewhat of an experiment. This is just an
+    // internal function, so feel free to remove if this bound turns out to be a
+    // bad idea. In that case, remember to also remove the lower bound
+    // `debug_assert!` below!
+    I: Iterator + TrustedLen,
+{
+    debug_assert!(N <= iter.size_hint().1.unwrap_or(usize::MAX));
+    debug_assert!(N <= iter.size_hint().0);
+
+    match collect_into_array(iter) {
+        Some(array) => array,
+        // SAFETY: covered by the function contract.
+        None => unsafe { crate::hint::unreachable_unchecked() },
+    }
+}
+
+/// Pulls `N` items from `iter` and returns them as an array. If the iterator
+/// yields fewer than `N` items, `None` is returned and all already yielded
+/// items are dropped.
+///
+/// Since the iterator is passed as mutable reference and this function calls
+/// `next` at most `N` times, the iterator can still be used afterwards to
+/// retrieve the remaining items.
+///
+/// If `iter.next()` panicks, all items already yielded by the iterator are
+/// dropped.
+fn collect_into_array<I, const N: usize>(iter: &mut I) -> Option<[I::Item; N]>
+where
+    I: Iterator,
+{
+    if N == 0 {
+        // SAFETY: An empty array is always inhabited and has no validity invariants.
+        return unsafe { Some(mem::zeroed()) };
+    }
+
+    struct Guard<T, const N: usize> {
+        ptr: *mut T,
+        initialized: usize,
+    }
+
+    impl<T, const N: usize> Drop for Guard<T, N> {
+        fn drop(&mut self) {
+            debug_assert!(self.initialized <= N);
+
+            let initialized_part = crate::ptr::slice_from_raw_parts_mut(self.ptr, self.initialized);
+
+            // SAFETY: this raw slice will contain only initialized objects.
+            unsafe {
+                crate::ptr::drop_in_place(initialized_part);
+            }
+        }
+    }
+
+    let mut array = MaybeUninit::uninit_array::<N>();
+    let mut guard: Guard<_, N> =
+        Guard { ptr: MaybeUninit::slice_as_mut_ptr(&mut array), initialized: 0 };
+
+    while let Some(item) = iter.next() {
+        // SAFETY: `guard.initialized` starts at 0, is increased by one in the
+        // loop and the loop is aborted once it reaches N (which is
+        // `array.len()`).
+        unsafe {
+            array.get_unchecked_mut(guard.initialized).write(item);
         }
+        guard.initialized += 1;
+
+        // Check if the whole array was initialized.
+        if guard.initialized == N {
+            mem::forget(guard);
 
-        // SAFETY: All elements of `dst` are properly initialized and
-        // `MaybeUninit<T>` has the same layout as `T`, so this cast is valid.
-        unsafe { (&mut out as *mut _ as *mut [&mut T; N]).read() }
+            // SAFETY: the condition above asserts that all elements are
+            // initialized.
+            let out = unsafe { MaybeUninit::array_assume_init(array) };
+            return Some(out);
+        }
     }
+
+    // This is only reached if the iterator is exhausted before
+    // `guard.initialized` reaches `N`. Also note that `guard` is dropped here,
+    // dropping all already initialized elements.
+    None
 }