From 4ebef6726b6cbd6cd2f6777fc4ea560a2f4d4f9a Mon Sep 17 00:00:00 2001
From: Tony Arcieri <bascule@gmail.com>
Date: Wed, 10 Jan 2024 20:52:43 -0700
Subject: [PATCH] Add `FromFn::try_from_fn`

Adds support for fallibly constructing arrays from a provided callback
---
 src/lib.rs   | 45 +++++++++++++++++++++++++++++++++++++++++++++
 tests/mod.rs | 18 ++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/src/lib.rs b/src/lib.rs
index 4eeb183..defc2a9 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -70,6 +70,15 @@ where
         Self(FromFn::from_fn(cb))
     }
 
+    /// Create array fallibly where each array element `T` is returned by the `cb` call, or return
+    /// an error if any are encountered.
+    pub fn try_from_fn<E, F>(cb: F) -> Result<Self, E>
+    where
+        F: FnMut(usize) -> Result<T, E>,
+    {
+        FromFn::try_from_fn(cb).map(Self)
+    }
+
     /// Returns an iterator over the array.
     #[inline]
     pub fn iter(&self) -> Iter<'_, T> {
@@ -831,6 +840,12 @@ pub trait FromFn<T>: Sized {
     fn from_fn<F>(cb: F) -> Self
     where
         F: FnMut(usize) -> T;
+
+    /// Create an array using the given callback function for each element, returning any errors
+    /// which are encountered in the given callback.
+    fn try_from_fn<E, F>(cb: F) -> Result<Self, E>
+    where
+        F: FnMut(usize) -> Result<T, E>;
 }
 
 impl<T, U> FromFn<T> for Array<T, U>
@@ -844,6 +859,13 @@ where
     {
         Array::from_fn(cb)
     }
+
+    fn try_from_fn<E, F>(cb: F) -> Result<Self, E>
+    where
+        F: FnMut(usize) -> Result<T, E>,
+    {
+        Array::try_from_fn(cb)
+    }
 }
 
 impl<T, const N: usize> FromFn<T> for [T; N] {
@@ -854,6 +876,29 @@ impl<T, const N: usize> FromFn<T> for [T; N] {
     {
         core::array::from_fn(cb)
     }
+
+    // TODO(tarcieri): use `array::try_from_fn` when stable
+    fn try_from_fn<E, F>(mut cb: F) -> Result<Self, E>
+    where
+        F: FnMut(usize) -> Result<T, E>,
+    {
+        // SAFETY: an array of `MaybeUninit`s is always valid.
+        let mut array: [MaybeUninit<T>; N] = unsafe { MaybeUninit::uninit().assume_init() };
+
+        #[allow(clippy::needless_range_loop)]
+        for i in 0..N {
+            array[i].write(cb(i)?);
+        }
+
+        // TODO(tarcieri): use `MaybeUninit::array_assume_init` when stable
+        let mut iter = array.into_iter();
+        Ok(Self::from_fn(|_| {
+            let item = iter.next().expect("should have enough items");
+
+            // SAFETY: if we got here, every element of the array was initialized
+            unsafe { item.assume_init() }
+        }))
+    }
 }
 
 /// Splits the shared slice into a slice of `N`-element arrays, starting at the beginning
diff --git a/tests/mod.rs b/tests/mod.rs
index 90949ee..0073b52 100644
--- a/tests/mod.rs
+++ b/tests/mod.rs
@@ -73,6 +73,24 @@ fn split_ref_mut() {
     assert_eq!(suffix.as_slice(), &EXAMPLE_SLICE[4..]);
 }
 
+#[test]
+fn from_fn() {
+    let array = Array::<u8, U6>::from_fn(|n| (n + 1) as u8);
+    assert_eq!(array.as_slice(), EXAMPLE_SLICE);
+}
+
+#[test]
+fn try_from_fn() {
+    let array = Array::<u8, U6>::try_from_fn::<(), _>(|n| Ok((n + 1) as u8)).unwrap();
+    assert_eq!(array.as_slice(), EXAMPLE_SLICE);
+
+    let err = Array::<u8, U6>::try_from_fn::<&'static str, _>(|_| Err("err"))
+        .err()
+        .unwrap();
+
+    assert_eq!(err, "err");
+}
+
 #[test]
 fn from_iterator_correct_size() {
     let array: Array<u8, U6> = EXAMPLE_SLICE.iter().copied().collect();