Add 'core::array::from_fn' and 'core::array::try_from_fn'

1 files changed, 91 insertions, 7 deletions

diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs
index 8c33a43ab33..b823c33ee23 100644
--- a/library/core/src/array/mod.rs
+++ b/library/core/src/array/mod.rs
@@ -20,6 +20,69 @@ mod iter;
 #[stable(feature = "array_value_iter", since = "1.51.0")]
 pub use iter::IntoIter;
 
+/// Creates an array `[T; N]` where each array element `T` is returned by the `cb` call.
+///
+/// # Arguments
+///
+/// * `cb`: Callback where the passed argument is the current array index.
+///
+/// # Example
+///
+/// ```rust
+/// #![feature(array_from_fn)]
+///
+/// let array = core::array::from_fn(|i| i);
+/// assert_eq!(array, [0, 1, 2, 3, 4]);
+/// ```
+#[inline]
+#[unstable(feature = "array_from_fn", issue = "89379")]
+pub fn from_fn<F, T, const N: usize>(mut cb: F) -> [T; N]
+where
+    F: FnMut(usize) -> T,
+{
+    let mut idx = 0;
+    [(); N].map(|_| {
+        let res = cb(idx);
+        idx += 1;
+        res
+    })
+}
+
+/// Creates an array `[T; N]` where each fallible array element `T` is returned by the `cb` call.
+/// Unlike `core::array::from_fn`, where the element creation can't fail, this version will return an error
+/// if any element creation was unsuccessful.
+///
+/// # Arguments
+///
+/// * `cb`: Callback where the passed argument is the current array index.
+///
+/// # Example
+///
+/// ```rust
+/// #![feature(array_from_fn)]
+///
+/// #[derive(Debug, PartialEq)]
+/// enum SomeError {
+///     Foo,
+/// }
+///
+/// let array = core::array::try_from_fn(|i| Ok::<_, SomeError>(i));
+/// assert_eq!(array, Ok([0, 1, 2, 3, 4]));
+///
+/// let another_array = core::array::try_from_fn::<SomeError, _, (), 2>(|_| Err(SomeError::Foo));
+/// assert_eq!(another_array, Err(SomeError::Foo));
+/// ```
+#[inline]
+#[unstable(feature = "array_from_fn", issue = "89379")]
+pub fn try_from_fn<E, F, T, const N: usize>(cb: F) -> Result<[T; N], E>
+where
+    F: FnMut(usize) -> Result<T, E>,
+{
+    // SAFETY: we know for certain that this iterator will yield exactly `N`
+    // items.
+    unsafe { collect_into_array_rslt_unchecked(&mut (0..N).map(cb)) }
+}
+
 /// Converts a reference to `T` into a reference to an array of length 1 (without copying).
 #[stable(feature = "array_from_ref", since = "1.53.0")]
 pub fn from_ref<T>(s: &T) -> &[T; 1] {
@@ -448,13 +511,15 @@ impl<T, const N: usize> [T; N] {
 ///
 /// 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]
+unsafe fn collect_into_array_rslt_unchecked<E, I, T, const N: usize>(
+    iter: &mut I,
+) -> Result<[T; N], E>
 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,
+    I: Iterator<Item = Result<T, E>> + TrustedLen,
 {
     debug_assert!(N <= iter.size_hint().1.unwrap_or(usize::MAX));
     debug_assert!(N <= iter.size_hint().0);
@@ -463,6 +528,18 @@ where
     unsafe { collect_into_array(iter).unwrap_unchecked() }
 }
 
+// Infallible version of `collect_into_array_rslt_unchecked`.
+unsafe fn collect_into_array_unchecked<I, const N: usize>(iter: &mut I) -> [I::Item; N]
+where
+    I: Iterator + TrustedLen,
+{
+    let mut map = iter.map(|el| Ok::<_, Infallible>(el));
+
+    // SAFETY: Valid array elements are covered by the fact that all passed values
+    // to `collect_into_array` are `Ok`.
+    unsafe { collect_into_array_rslt_unchecked(&mut map).unwrap_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.
@@ -473,13 +550,13 @@ where
 ///
 /// 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]>
+fn collect_into_array<E, I, T, const N: usize>(iter: &mut I) -> Option<Result<[T; N], E>>
 where
-    I: Iterator,
+    I: Iterator<Item = Result<T, E>>,
 {
     if N == 0 {
         // SAFETY: An empty array is always inhabited and has no validity invariants.
-        return unsafe { Some(mem::zeroed()) };
+        return unsafe { Some(Ok(mem::zeroed())) };
     }
 
     struct Guard<T, const N: usize> {
@@ -504,7 +581,14 @@ where
     let mut guard: Guard<_, N> =
         Guard { ptr: MaybeUninit::slice_as_mut_ptr(&mut array), initialized: 0 };
 
-    while let Some(item) = iter.next() {
+    while let Some(item_rslt) = iter.next() {
+        let item = match item_rslt {
+            Err(err) => {
+                return Some(Err(err));
+            }
+            Ok(elem) => elem,
+        };
+
         // 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()`).
@@ -520,7 +604,7 @@ where
             // SAFETY: the condition above asserts that all elements are
             // initialized.
             let out = unsafe { MaybeUninit::array_assume_init(array) };
-            return Some(out);
+            return Some(Ok(out));
         }
     }