Add core::array::from_fn

This adds `core::array::from_fn` for initializing an array from a
closure. The API is meant to closely resemble `core::iter::from_fn` with
the difference that the array dictates the amount of elements produced
as opposed to the function.

This should also cover a lot of cases where one might want to reach for
collecting into an array, but without all the associated problems that
this would have, such as providing too few or too many elements.

```rust
let array = std::array::from_fn(|index| 2 * index);
assert_eq!(array, [0, 2, 4, 6]);
```

Author: CryZe

library/core/src/array/mod.rs

@@ -532,6 +532,68 @@ impl<T, const N: usize> [T; N] {
     }
 }
 
+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);
+        }
+    }
+}
+
+/// Creates a new array where each element is initialized by calling the
+/// provided closure `F: FnMut(usize) -> T` with the index of each element.
+///
+/// This allows initializing an array with elements that are not `Copy` or
+/// `const` and where writing out the entire array literal as opposed to
+/// `[expr; N]` is too verbose or not possible.
+///
+/// # Example
+///
+/// ```
+/// #![feature(array_from_fn)]
+///
+/// let array = std::array::from_fn(|index| 2 * index);
+/// assert_eq!(array, [0, 2, 4, 6]);
+/// ```
+#[unstable(feature = "array_from_fn", issue = "none")]
+pub fn from_fn<T, F, const N: usize>(mut f: F) -> [T; N]
+where
+    F: FnMut(usize) -> T,
+{
+    let mut array = MaybeUninit::uninit_array::<N>();
+    let mut guard: Guard<_, N> =
+        Guard { ptr: MaybeUninit::slice_as_mut_ptr(&mut array), initialized: 0 };
+
+    while guard.initialized < N {
+        let item = f(guard.initialized);
+
+        // 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;
+    }
+
+    mem::forget(guard);
+
+    // SAFETY: the condition above asserts that all elements are
+    // initialized.
+    unsafe { MaybeUninit::array_assume_init(array) }
+}
+
 /// Pulls `N` items from `iter` and returns them as an array. If the iterator
 /// yields fewer than `N` items, this function exhibits undefined behavior.
 ///
@@ -568,7 +630,7 @@ where
 /// `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
+/// If `iter.next()` panics, 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
@@ -579,24 +641,6 @@ where
         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 };

library/core/tests/array.rs

@@ -306,6 +306,15 @@ fn empty_array_is_always_default() {
     let _arr = <[DoesNotImplDefault; 0]>::default();
 }
 
+#[test]
+fn array_from_fn() {
+    let a = array::from_fn(|i| 2 * i + 1);
+    assert_eq!(a, [1, 3, 5, 7]);
+
+    let b = array::from_fn(|i| i.to_string());
+    assert_eq!(b, [String::from("0"), String::from("1")]);
+}
+
 #[test]
 fn array_map() {
     let a = [1, 2, 3];

library/core/tests/lib.rs

@@ -1,5 +1,6 @@
 #![feature(alloc_layout_extra)]
 #![feature(array_chunks)]
+#![feature(array_from_fn)]
 #![feature(array_from_ref)]
 #![feature(array_methods)]
 #![feature(array_map)]