Make array::{try_from_fn, try_map} and Iterator::try_find generic over Try

Fixes 85115

This only updates unstable functions.

`array::try_map` didn't actually exist before, despite the tracking issue 79711 still being open from the old PR 79713.

8 files changed, 195 insertions, 45 deletions

diff --git a/library/core/src/array/mod.rs b/library/core/src/array/mod.rs
index 39ccbaaaf7b..d635829151e 100644
--- a/library/core/src/array/mod.rs
+++ b/library/core/src/array/mod.rs
@@ -11,7 +11,9 @@ use crate::fmt;
 use crate::hash::{self, Hash};
 use crate::iter::TrustedLen;
 use crate::mem::{self, MaybeUninit};
-use crate::ops::{Index, IndexMut};
+use crate::ops::{
+    ChangeOutputType, ControlFlow, FromResidual, Index, IndexMut, NeverShortCircuit, Residual, Try,
+};
 use crate::slice::{Iter, IterMut};
 
 mod equality;
@@ -49,9 +51,13 @@ where
 }
 
 /// 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
+/// Unlike [`from_fn`], where the element creation can't fail, this version will return an error
 /// if any element creation was unsuccessful.
 ///
+/// The return type of this function depends on the return type of the closure.
+/// If you return `Result<T, E>` from the closure, you'll get a `Result<[T; N]; E>`.
+/// If you return `Option<T>` from the closure, you'll get an `Option<[T; N]>`.
+///
 /// # Arguments
 ///
 /// * `cb`: Callback where the passed argument is the current array index.
@@ -60,27 +66,32 @@ where
 ///
 /// ```rust
 /// #![feature(array_from_fn)]
+/// # // Apparently these doc tests are still on edition2018
+/// # use std::convert::TryInto;
 ///
-/// #[derive(Debug, PartialEq)]
-/// enum SomeError {
-///     Foo,
-/// }
-///
-/// let array = core::array::try_from_fn(|i| Ok::<_, SomeError>(i));
+/// let array: Result<[u8; 5], _> = std::array::try_from_fn(|i| i.try_into());
 /// 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));
+/// let array: Result<[i8; 200], _> = std::array::try_from_fn(|i| i.try_into());
+/// assert!(array.is_err());
+///
+/// let array: Option<[_; 4]> = std::array::try_from_fn(|i| i.checked_add(100));
+/// assert_eq!(array, Some([100, 101, 102, 103]));
+///
+/// let array: Option<[_; 4]> = std::array::try_from_fn(|i| i.checked_sub(100));
+/// assert_eq!(array, None);
 /// ```
 #[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>
+pub fn try_from_fn<F, R, const N: usize>(cb: F) -> ChangeOutputType<R, [R::Output; N]>
 where
-    F: FnMut(usize) -> Result<T, E>,
+    F: FnMut(usize) -> R,
+    R: Try,
+    R::Residual: Residual<[R::Output; N]>,
 {
     // SAFETY: we know for certain that this iterator will yield exactly `N`
     // items.
-    unsafe { collect_into_array_rslt_unchecked(&mut (0..N).map(cb)) }
+    unsafe { try_collect_into_array_unchecked(&mut (0..N).map(cb)) }
 }
 
 /// Converts a reference to `T` into a reference to an array of length 1 (without copying).
@@ -444,6 +455,45 @@ impl<T, const N: usize> [T; N] {
         unsafe { collect_into_array_unchecked(&mut IntoIterator::into_iter(self).map(f)) }
     }
 
+    /// A fallible function `f` applied to each element on array `self` in order to
+    /// return an array the same size as `self` or the first error encountered.
+    ///
+    /// The return type of this function depends on the return type of the closure.
+    /// If you return `Result<T, E>` from the closure, you'll get a `Result<[T; N]; E>`.
+    /// If you return `Option<T>` from the closure, you'll get an `Option<[T; N]>`.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// #![feature(array_try_map)]
+    /// let a = ["1", "2", "3"];
+    /// let b = a.try_map(|v| v.parse::<u32>()).unwrap().map(|v| v + 1);
+    /// assert_eq!(b, [2, 3, 4]);
+    ///
+    /// let a = ["1", "2a", "3"];
+    /// let b = a.try_map(|v| v.parse::<u32>());
+    /// assert!(b.is_err());
+    ///
+    /// use std::num::NonZeroU32;
+    /// let z = [1, 2, 0, 3, 4];
+    /// assert_eq!(z.try_map(NonZeroU32::new), None);
+    /// let a = [1, 2, 3];
+    /// let b = a.try_map(NonZeroU32::new);
+    /// let c = b.map(|x| x.map(NonZeroU32::get));
+    /// assert_eq!(c, Some(a));
+    /// ```
+    #[unstable(feature = "array_try_map", issue = "79711")]
+    pub fn try_map<F, R>(self, f: F) -> ChangeOutputType<R, [R::Output; N]>
+    where
+        F: FnMut(T) -> R,
+        R: Try,
+        R::Residual: Residual<[R::Output; N]>,
+    {
+        // SAFETY: we know for certain that this iterator will yield exactly `N`
+        // items.
+        unsafe { try_collect_into_array_unchecked(&mut IntoIterator::into_iter(self).map(f)) }
+    }
+
     /// 'Zips up' two arrays into a single array of pairs.
     ///
     /// `zip()` returns a new array where every element is a tuple where the
@@ -621,42 +671,42 @@ impl<T, const N: usize> [T; N] {
 /// Pulls `N` items from `iter` and returns them as an array. If the iterator
 /// yields fewer than `N` items, this function exhibits undefined behavior.
 ///
-/// See [`collect_into_array`] for more information.
+/// See [`try_collect_into_array`] for more information.
 ///
 ///
 /// # Safety
 ///
 /// 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_rslt_unchecked<E, I, T, const N: usize>(
-    iter: &mut I,
-) -> Result<[T; N], E>
+unsafe fn try_collect_into_array_unchecked<I, T, R, const N: usize>(iter: &mut I) -> R::TryType
 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<Item = Result<T, E>> + TrustedLen,
+    I: Iterator + TrustedLen,
+    I::Item: Try<Output = T, Residual = R>,
+    R: Residual<[T; N]>,
 {
     debug_assert!(N <= iter.size_hint().1.unwrap_or(usize::MAX));
     debug_assert!(N <= iter.size_hint().0);
 
     // SAFETY: covered by the function contract.
-    unsafe { collect_into_array(iter).unwrap_unchecked() }
+    unsafe { try_collect_into_array(iter).unwrap_unchecked() }
 }
 
-// Infallible version of `collect_into_array_rslt_unchecked`.
+// Infallible version of `try_collect_into_array_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(Ok::<_, Infallible>);
+    let mut map = iter.map(NeverShortCircuit);
 
     // SAFETY: The same safety considerations w.r.t. the iterator length
-    // apply for `collect_into_array_rslt_unchecked` as for
+    // apply for `try_collect_into_array_unchecked` as for
     // `collect_into_array_unchecked`
-    match unsafe { collect_into_array_rslt_unchecked(&mut map) } {
-        Ok(array) => array,
+    match unsafe { try_collect_into_array_unchecked(&mut map) } {
+        NeverShortCircuit(array) => array,
     }
 }
 
@@ -670,13 +720,15 @@ where
 ///
 /// If `iter.next()` panicks, all items already yielded by the iterator are
 /// dropped.
-fn collect_into_array<E, I, T, const N: usize>(iter: &mut I) -> Option<Result<[T; N], E>>
+fn try_collect_into_array<I, T, R, const N: usize>(iter: &mut I) -> Option<R::TryType>
 where
-    I: Iterator<Item = Result<T, E>>,
+    I: Iterator,
+    I::Item: Try<Output = T, Residual = R>,
+    R: Residual<[T; N]>,
 {
     if N == 0 {
         // SAFETY: An empty array is always inhabited and has no validity invariants.
-        return unsafe { Some(Ok(mem::zeroed())) };
+        return unsafe { Some(Try::from_output(mem::zeroed())) };
     }
 
     struct Guard<'a, T, const N: usize> {
@@ -701,11 +753,11 @@ where
     let mut guard = Guard { array_mut: &mut array, initialized: 0 };
 
     while let Some(item_rslt) = iter.next() {
-        let item = match item_rslt {
-            Err(err) => {
-                return Some(Err(err));
+        let item = match item_rslt.branch() {
+            ControlFlow::Break(r) => {
+                return Some(FromResidual::from_residual(r));
             }
-            Ok(elem) => elem,
+            ControlFlow::Continue(elem) => elem,
         };
 
         // SAFETY: `guard.initialized` starts at 0, is increased by one in the
@@ -723,7 +775,7 @@ where
             // SAFETY: the condition above asserts that all elements are
             // initialized.
             let out = unsafe { MaybeUninit::array_assume_init(array) };
-            return Some(Ok(out));
+            return Some(Try::from_output(out));
         }
     }
 
diff --git a/library/core/src/iter/traits/iterator.rs b/library/core/src/iter/traits/iterator.rs
index 35ce9400f8f..88e7623eba1 100644
--- a/library/core/src/iter/traits/iterator.rs
+++ b/library/core/src/iter/traits/iterator.rs
@@ -1,5 +1,5 @@
 use crate::cmp::{self, Ordering};
-use crate::ops::{ControlFlow, Try};
+use crate::ops::{ChangeOutputType, ControlFlow, FromResidual, Residual, Try};
 
 use super::super::TrustedRandomAccessNoCoerce;
 use super::super::{Chain, Cloned, Copied, Cycle, Enumerate, Filter, FilterMap, Fuse};
@@ -2418,6 +2418,10 @@ pub trait Iterator {
     /// Applies function to the elements of iterator and returns
     /// the first true result or the first error.
     ///
+    /// The return type of this method depends on the return type of the closure.
+    /// If you return `Result<bool, E>` from the closure, you'll get a `Result<Option<Self::Item>; E>`.
+    /// If you return `Option<bool>` from the closure, you'll get an `Option<Option<Self::Item>>`.
+    ///
     /// # Examples
     ///
     /// ```
@@ -2435,32 +2439,48 @@ pub trait Iterator {
     /// let result = a.iter().try_find(|&&s| is_my_num(s, 5));
     /// assert!(result.is_err());
     /// ```
+    ///
+    /// This also supports other types which implement `Try`, not just `Result`.
+    /// ```
+    /// #![feature(try_find)]
+    ///
+    /// use std::num::NonZeroU32;
+    /// let a = [3, 5, 7, 4, 9, 0, 11];
+    /// let result = a.iter().try_find(|&&x| NonZeroU32::new(x).map(|y| y.is_power_of_two()));
+    /// assert_eq!(result, Some(Some(&4)));
+    /// let result = a.iter().take(3).try_find(|&&x| NonZeroU32::new(x).map(|y| y.is_power_of_two()));
+    /// assert_eq!(result, Some(None));
+    /// let result = a.iter().rev().try_find(|&&x| NonZeroU32::new(x).map(|y| y.is_power_of_two()));
+    /// assert_eq!(result, None);
+    /// ```
     #[inline]
     #[unstable(feature = "try_find", reason = "new API", issue = "63178")]
-    fn try_find<F, R, E>(&mut self, f: F) -> Result<Option<Self::Item>, E>
+    fn try_find<F, R>(&mut self, f: F) -> ChangeOutputType<R, Option<Self::Item>>
     where
         Self: Sized,
         F: FnMut(&Self::Item) -> R,
         R: Try<Output = bool>,
-        // FIXME: This bound is rather strange, but means minimal breakage on nightly.
-        // See #85115 for the issue tracking a holistic solution for this and try_map.
-        R: Try<Residual = Result<crate::convert::Infallible, E>>,
+        R::Residual: Residual<Option<Self::Item>>,
     {
         #[inline]
-        fn check<F, T, R, E>(mut f: F) -> impl FnMut((), T) -> ControlFlow<Result<T, E>>
+        fn check<I, V, R>(
+            mut f: impl FnMut(&I) -> V,
+        ) -> impl FnMut((), I) -> ControlFlow<R::TryType>
         where
-            F: FnMut(&T) -> R,
-            R: Try<Output = bool>,
-            R: Try<Residual = Result<crate::convert::Infallible, E>>,
+            V: Try<Output = bool, Residual = R>,
+            R: Residual<Option<I>>,
         {
             move |(), x| match f(&x).branch() {
                 ControlFlow::Continue(false) => ControlFlow::CONTINUE,
-                ControlFlow::Continue(true) => ControlFlow::Break(Ok(x)),
-                ControlFlow::Break(Err(x)) => ControlFlow::Break(Err(x)),
+                ControlFlow::Continue(true) => ControlFlow::Break(Try::from_output(Some(x))),
+                ControlFlow::Break(r) => ControlFlow::Break(FromResidual::from_residual(r)),
             }
         }
 
-        self.try_fold((), check(f)).break_value().transpose()
+        match self.try_fold((), check(f)) {
+            ControlFlow::Break(x) => x,
+            ControlFlow::Continue(()) => Try::from_output(None),
+        }
     }
 
     /// Searches for an element in an iterator, returning its index.
diff --git a/library/core/src/ops/control_flow.rs b/library/core/src/ops/control_flow.rs
index b0c15898a1f..5ee8e377de1 100644
--- a/library/core/src/ops/control_flow.rs
+++ b/library/core/src/ops/control_flow.rs
@@ -123,6 +123,11 @@ impl<B, C> ops::FromResidual for ControlFlow<B, C> {
     }
 }
 
+#[unstable(feature = "try_trait_v2_residual", issue = "91285")]
+impl<B, C> ops::Residual<C> for ControlFlow<B, convert::Infallible> {
+    type TryType = ControlFlow<B, C>;
+}
+
 impl<B, C> ControlFlow<B, C> {
     /// Returns `true` if this is a `Break` variant.
     ///
diff --git a/library/core/src/ops/mod.rs b/library/core/src/ops/mod.rs
index 6c866031370..9d1e7e81b0e 100644
--- a/library/core/src/ops/mod.rs
+++ b/library/core/src/ops/mod.rs
@@ -187,6 +187,11 @@ pub use self::range::OneSidedRange;
 #[unstable(feature = "try_trait_v2", issue = "84277")]
 pub use self::try_trait::{FromResidual, Try};
 
+#[unstable(feature = "try_trait_v2_residual", issue = "91285")]
+pub use self::try_trait::Residual;
+
+pub(crate) use self::try_trait::{ChangeOutputType, NeverShortCircuit};
+
 #[unstable(feature = "generator_trait", issue = "43122")]
 pub use self::generator::{Generator, GeneratorState};
 
diff --git a/library/core/src/ops/try_trait.rs b/library/core/src/ops/try_trait.rs
index 6bdcda775ee..f4f0a589809 100644
--- a/library/core/src/ops/try_trait.rs
+++ b/library/core/src/ops/try_trait.rs
@@ -338,3 +338,61 @@ pub trait FromResidual<R = <Self as Try>::Residual> {
     #[unstable(feature = "try_trait_v2", issue = "84277")]
     fn from_residual(residual: R) -> Self;
 }
+
+/// Allows retrieving the canonical type implementing [`Try`] that has this type
+/// as its residual and allows it to hold an `O` as its output.
+///
+/// If you think of the `Try` trait as splitting a type into its [`Try::Output`]
+/// and [`Try::Residual`] components, this allows putting them back together.
+///
+/// For example,
+/// `Result<T, E>: Try<Output = T, Residual = Result<Infallible, E>>`,
+/// and in the other direction,
+/// `<Result<Infallible, E> as Residual<T>>::TryType = Result<T, E>`.
+#[unstable(feature = "try_trait_v2_residual", issue = "91285")]
+pub trait Residual<O> {
+    /// The "return" type of this meta-function.
+    #[unstable(feature = "try_trait_v2_residual", issue = "91285")]
+    type TryType: Try<Output = O, Residual = Self>;
+}
+
+#[unstable(feature = "pub_crate_should_not_need_unstable_attr", issue = "none")]
+pub(crate) type ChangeOutputType<T, V> = <<T as Try>::Residual as Residual<V>>::TryType;
+
+/// An adapter for implementing non-try methods via the `Try` implementation.
+///
+/// Conceptually the same as `Result<T, !>`, but requiring less work in trait
+/// solving and inhabited-ness checking and such, by being an obvious newtype
+/// and not having `From` bounds lying around.
+///
+/// Not currently planned to be exposed publicly, so just `pub(crate)`.
+#[repr(transparent)]
+pub(crate) struct NeverShortCircuit<T>(pub T);
+
+pub(crate) enum NeverShortCircuitResidual {}
+
+impl<T> Try for NeverShortCircuit<T> {
+    type Output = T;
+    type Residual = NeverShortCircuitResidual;
+
+    #[inline]
+    fn branch(self) -> ControlFlow<NeverShortCircuitResidual, T> {
+        ControlFlow::Continue(self.0)
+    }
+
+    #[inline]
+    fn from_output(x: T) -> Self {
+        NeverShortCircuit(x)
+    }
+}
+
+impl<T> FromResidual for NeverShortCircuit<T> {
+    #[inline]
+    fn from_residual(never: NeverShortCircuitResidual) -> Self {
+        match never {}
+    }
+}
+
+impl<T> Residual<T> for NeverShortCircuitResidual {
+    type TryType = NeverShortCircuit<T>;
+}
diff --git a/library/core/src/option.rs b/library/core/src/option.rs
index 4eeb5e43943..4e11162f8c2 100644
--- a/library/core/src/option.rs
+++ b/library/core/src/option.rs
@@ -2088,6 +2088,11 @@ impl<T> const ops::FromResidual for Option<T> {
     }
 }
 
+#[unstable(feature = "try_trait_v2_residual", issue = "91285")]
+impl<T> ops::Residual<T> for Option<convert::Infallible> {
+    type TryType = Option<T>;
+}
+
 impl<T> Option<Option<T>> {
     /// Converts from `Option<Option<T>>` to `Option<T>`.
     ///
diff --git a/library/core/src/result.rs b/library/core/src/result.rs
index a494c089f68..e6b8c8ec338 100644
--- a/library/core/src/result.rs
+++ b/library/core/src/result.rs
@@ -1959,3 +1959,8 @@ impl<T, E, F: From<E>> ops::FromResidual<Result<convert::Infallible, E>> for Res
         }
     }
 }
+
+#[unstable(feature = "try_trait_v2_residual", issue = "91285")]
+impl<T, E> ops::Residual<T> for Result<convert::Infallible, E> {
+    type TryType = Result<T, E>;
+}
diff --git a/library/core/tests/array.rs b/library/core/tests/array.rs
index 7dc071b7423..43d30d12355 100644
--- a/library/core/tests/array.rs
+++ b/library/core/tests/array.rs
@@ -377,7 +377,7 @@ fn array_try_from_fn() {
     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));
+    let another_array = core::array::try_from_fn::<_, Result<(), _>, 2>(|_| Err(SomeError::Foo));
     assert_eq!(another_array, Err(SomeError::Foo));
 }