Rollup merge of rust-lang#82098 - LukasKalbertodt:add-collect-into-array, r=dtolnay

Add internal `collect_into_array[_unchecked]` to remove duplicate code

Unlike the similar PRs  rust-lang#69985, rust-lang#75644 and rust-lang#79659, this PR only adds private functions and does not propose any new public API. The change 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, ...}`, `[T; N]::{cloned, copied, ...}`, ...) which all basically need this functionality. Writing custom `unsafe` code for each of those doesn't seem like a good idea. I added two functions in this PR (and not just the `unsafe` version) because I already know that I need the `Option`-returning version for `Iterator::map_windows`.

This is closely related to rust-lang#81615. I think that all options listed in that issue can be implemented using the function added in this PR. The only instance where `collect_array_into` might not be general enough is when the caller want to handle incomplete arrays manually. Currently, if `iter` yields fewer than `N` items, `None` is returned and the already yielded items are dropped. But as this is just a private function, it can be made more general in future PRs.

And while this was not the goal, this seems to lead to better assembly for `array::map`: https://rust.godbolt.org/z/75qKTa (CC `@JulianKnodt)`

Let me know what you think :)

CC `@matklad` `@bstrie`

Author: Dylan-DPC

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
 }