From 3a4a874185d9ad0078e6f71f698f824e1b0b9084 Mon Sep 17 00:00:00 2001
From: Christopher Serr <christopher.serr@gmail.com>
Date: Sat, 27 Mar 2021 18:48:48 +0100
Subject: [PATCH] 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]);
```
---
 library/core/src/array/mod.rs | 82 +++++++++++++++++++++++++++--------
 library/core/tests/array.rs   |  9 ++++
 library/core/tests/lib.rs     |  1 +
 3 files changed, 73 insertions(+), 19 deletions(-)

diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs
index 8f52985d1df71..fd8827d0bcc89 100644
--- a/library/core/src/array/mod.rs
+++ b/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 };
diff --git a/library/core/tests/array.rs b/library/core/tests/array.rs
index 89c2a969c28bb..0452876bcee85 100644
--- a/library/core/tests/array.rs
+++ b/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];
diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs
index 997e618cc511f..e862e3d220c1d 100644
--- a/library/core/tests/lib.rs
+++ b/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)]