add array::FillError similar to array::IntoIter

library/core/src/array/mod.rs

@@ -12,6 +12,7 @@ use crate::convert::{Infallible, TryFrom};
 use crate::fmt;
 use crate::hash::{self, Hash};
 use crate::marker::Unsize;
+use crate::mem::{self, MaybeUninit};
 use crate::slice::{Iter, IterMut};
 
 mod iter;
@@ -174,6 +175,101 @@ impl<T: fmt::Debug, const N: usize> fmt::Debug for [T; N] {
     }
 }
 
+/// The error returned by the `FromIterator` implementation of
+/// arrays once these are implemented.
+///
+/// Until then `FillError::new().fill(iter)` can be used instead.
+#[unstable(feature = "array_from_iter_impl", issue = "none")]
+pub struct FillError<T, const N: usize> {
+    array: [MaybeUninit<T>; N],
+    len: usize,
+}
+
+#[unstable(feature = "array_from_iter_impl", issue = "none")]
+impl<T, const N: usize> fmt::Display for FillError<T, N> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt::Display::fmt(
+            &format_args!("The iterator only returned {} items, but {} were needed", self.len(), N),
+            f,
+        )
+    }
+}
+
+#[unstable(feature = "array_from_iter_impl", issue = "none")]
+impl<T, const N: usize> fmt::Debug for FillError<T, N> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt::Display::fmt(self, f)
+    }
+}
+
+#[unstable(feature = "array_from_iter_impl", issue = "none")]
+impl<T, const N: usize> Drop for FillError<T, N> {
+    fn drop(&mut self) {
+        // SAFETY: This is safe: `as_mut_slice` returns exactly the sub-slice
+        // of elements that have been initialized and need to be dropped.
+        unsafe { crate::ptr::drop_in_place(self.as_mut_slice()) }
+    }
+}
+
+#[unstable(feature = "array_from_iter_impl", issue = "none")]
+impl<T, const N: usize> FillError<T, N> {
+    /// Creates a new empty `FillError` which can be used
+    /// to build `[T; N]` from an iterator.
+    pub fn new() -> Self {
+        Self { array: MaybeUninit::uninit_array(), len: 0 }
+    }
+
+    /// Returns how many elements were read from the given iterator.
+    pub fn len(&self) -> usize {
+        self.len
+    }
+
+    /// Returns an immutable slice of all initialized elements.
+    pub fn as_slice(&self) -> &[T] {
+        // SAFETY: We know that all elements from 0 to len are properly initialized.
+        unsafe { MaybeUninit::slice_assume_init_ref(&self.array[0..self.len]) }
+    }
+
+    /// Returns a mutable slice of all initialized elements.
+    pub fn as_mut_slice(&mut self) -> &mut [T] {
+        // SAFETY: We know that all elements from 0 to len are properly initialized.
+        unsafe { MaybeUninit::slice_assume_init_mut(&mut self.array[0..self.len]) }
+    }
+
+    /// Tries to initialize the left-over elements using `iter`.
+    pub fn fill<I: IntoIterator<Item = T>>(mut self, iter: I) -> Result<[T; N], FillError<T, N>> {
+        let mut iter = iter.into_iter();
+
+        for i in self.len..N {
+            if let Some(value) = iter.next() {
+                self.array[i].write(value);
+                self.len = i + 1;
+            } else {
+                return Err(self);
+            }
+        }
+
+        // SAFETY: The transmute here is actually safe. The docs of `MaybeUninit`
+        // promise:
+        //
+        // > `MaybeUninit<T>` is guaranteed to have the same size and alignment
+        // > as `T`.
+        //
+        // The docs even show a transmute from an array of `MaybeUninit<T>` to
+        // an array of `T`.
+        //
+        // With that, this initialization satisfies the invariants.
+        //
+        // FIXME: actually use `mem::transmute` here, once it
+        // works with const generics:
+        //     `mem::transmute::<[MaybeUninit<T>; N], [T; N]>(array)`
+        let arr = unsafe { crate::mem::transmute_copy::<_, [T; N]>(&self.array) };
+        // We forget `self` to not drop anything here.
+        mem::forget(self);
+        Ok(arr)
+    }
+}
+
 #[stable(feature = "rust1", since = "1.0.0")]
 impl<'a, T, const N: usize> IntoIterator for &'a [T; N] {
     type Item = &'a T;
@@ -385,42 +481,11 @@ 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,
     {
-        use crate::mem::MaybeUninit;
-        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) }
+        FillError::new().fill(IntoIter::new(self).map(f)).unwrap()
     }
 
     /// Returns a slice containing the entire array. Equivalent to `&s[..]`.

library/core/src/mem/maybe_uninit.rs

@@ -819,7 +819,7 @@ impl<T> MaybeUninit<T> {
     #[unstable(feature = "maybe_uninit_slice", issue = "63569")]
     #[inline(always)]
     pub unsafe fn slice_assume_init_mut(slice: &mut [Self]) -> &mut [T] {
-        // SAFETY: similar to safety notes for `slice_get_ref`, but we have a
+        // SAFETY: similar to safety notes for `slice_assume_init_ref`, but we have a
         // mutable reference which is also guaranteed to be valid for writes.
         unsafe { &mut *(slice as *mut [Self] as *mut [T]) }
     }

library/core/tests/array.rs

@@ -1,4 +1,4 @@
-use core::array::{FixedSizeArray, IntoIter};
+use core::array::{FillError, FixedSizeArray, IntoIter};
 use core::convert::TryFrom;
 
 #[test]
@@ -330,3 +330,60 @@ fn array_map_drop_safety() {
     assert_eq!(DROPPED.load(Ordering::SeqCst), num_to_create);
     panic!("test succeeded")
 }
+
+#[test]
+fn array_collects() {
+    let v = vec![1, 2, 3, 4, 5];
+    let a: [i32; 5] = FillError::new().fill(v.clone()).unwrap();
+
+    assert_eq!(v[..], a[..]);
+}
+
+#[test]
+#[should_panic(expected = "test succeeded")]
+fn array_collect_panic_safety() {
+    use core::sync::atomic::AtomicUsize;
+    use core::sync::atomic::Ordering;
+    static DROPPED: AtomicUsize = AtomicUsize::new(0);
+    struct DropCounter;
+    impl Drop for DropCounter {
+        fn drop(&mut self) {
+            DROPPED.fetch_add(1, Ordering::SeqCst);
+        }
+    }
+
+    let num_to_create = 5;
+    let success = std::panic::catch_unwind(|| {
+        let items = [0; 10];
+        let mut nth = 0;
+        FillError::<DropCounter, 10>::new()
+            .fill(items.iter().map(|_| {
+                assert!(nth < num_to_create);
+                nth += 1;
+                DropCounter
+            }))
+            .unwrap()
+    });
+    assert!(success.is_err());
+    assert_eq!(DROPPED.load(Ordering::SeqCst), num_to_create);
+    panic!("test succeeded")
+}
+
+#[test]
+fn array_collect_no_double_drop() {
+    use core::sync::atomic::AtomicUsize;
+    use core::sync::atomic::Ordering;
+    static DROPPED: AtomicUsize = AtomicUsize::new(0);
+    struct DropCounter;
+    impl Drop for DropCounter {
+        fn drop(&mut self) {
+            DROPPED.fetch_add(1, Ordering::SeqCst);
+        }
+    }
+
+    const LEN: usize = 10;
+    let items = [0; LEN];
+    let _ = FillError::<DropCounter, 10>::new().fill(items.iter().map(|_| DropCounter)).unwrap();
+
+    assert_eq!(DROPPED.load(Ordering::SeqCst), LEN);
+}

library/core/tests/lib.rs

@@ -56,6 +56,7 @@
 #![feature(unsafe_block_in_unsafe_fn)]
 #![feature(int_bits_const)]
 #![deny(unsafe_op_in_unsafe_fn)]
+#![feature(array_from_iter_impl)]
 
 extern crate test;