Add TryFutureExt::map_ok_or_else method

futures-util/src/future/try_future/map_ok_or_else.rs

@@ -0,0 +1,59 @@
+use core::pin::Pin;
+use futures_core::future::{FusedFuture, Future, TryFuture};
+use futures_core::task::{Context, Poll};
+use pin_utils::{unsafe_pinned, unsafe_unpinned};
+
+/// Future for the [`map_ok_or_else`](super::TryFutureExt::map_ok_or_else) method.
+#[derive(Debug)]
+#[must_use = "futures do nothing unless you `.await` or poll them"]
+pub struct MapOkOrElse<Fut, F, E> {
+    future: Fut,
+    f: Option<F>,
+    e: Option<E>,
+}
+
+impl<Fut, F, E> MapOkOrElse<Fut, F, E> {
+    unsafe_pinned!(future: Fut);
+    unsafe_unpinned!(f: Option<F>);
+    unsafe_unpinned!(e: Option<E>);
+
+    /// Creates a new MapOkOrElse.
+    pub(super) fn new(future: Fut, e: E, f: F) -> Self {
+        Self { future, f: Some(f), e: Some(e) }
+    }
+}
+
+impl<Fut: Unpin, F, E> Unpin for MapOkOrElse<Fut, F, E> {}
+
+impl<Fut, F, E, T> FusedFuture for MapOkOrElse<Fut, F, E>
+    where Fut: TryFuture,
+          F: FnOnce(Fut::Ok) -> T,
+          E: FnOnce(Fut::Error) -> T,
+{
+    fn is_terminated(&self) -> bool {
+        self.f.is_none() || self.e.is_none()
+    }
+}
+
+impl<Fut, F, E, T> Future for MapOkOrElse<Fut, F, E>
+    where Fut: TryFuture,
+          F: FnOnce(Fut::Ok) -> T,
+          E: FnOnce(Fut::Error) -> T,
+{
+    type Output = T;
+
+    fn poll(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Self::Output> {
+        self.as_mut()
+            .future()
+            .try_poll(cx)
+            .map(|result| {
+                match result {
+                    Ok(i) => (self.as_mut().f().take().expect("MapOkOrElse must not be polled after it returned `Poll::Ready`"))(i),
+                    Err(e) => (self.as_mut().e().take().expect("MapOkOrElse must not be polled after it returned `Poll::Ready`"))(e),
+                }
+            })
+    }
+}

futures-util/src/future/try_future/mod.rs

@@ -50,6 +50,10 @@ mod map_ok;
 #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
 pub use self::map_ok::MapOk;
 
+mod map_ok_or_else;
+#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
+pub use self::map_ok_or_else::MapOkOrElse;
+
 mod or_else;
 #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
 pub use self::or_else::OrElse;
@@ -160,6 +164,47 @@ pub trait TryFutureExt: TryFuture {
         MapOk::new(self, f)
     }
 
+    /// Maps this future's success value to a different value, and permits for error handling resulting in the same type.
+    ///
+    /// This method can be used to coalesce your [`Ok`](TryFuture::Ok) type and [`Error`](TryFuture::Error) into another type,
+    /// where that type is the same for both outcomes.
+    ///
+    /// The provided closure `f` will only be called if this future is resolved
+    /// to an [`Ok`]. If it resolves to an [`Err`], panics, or is dropped, then
+    /// the provided closure will never be invoked.
+    /// 
+    /// The provided closure `e` will only be called if this future is resolved
+    /// to an [`Err`]. If it resolves to an [`Ok`], panics, or is dropped, then
+    /// the provided closure will never be invoked.
+    ///
+    /// Note that this method consumes the future it is called on and returns a
+    /// wrapped version of it.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use futures::future::TryFutureExt;
+    ///
+    /// # futures::executor::block_on(async {
+    /// let future = async { Ok::<i32, i32>(5) };
+    /// let future = future.map_ok_or_else(|x| x * 2, |x| x + 3);
+    /// assert_eq!(future.await, 8);
+    /// 
+    /// let future = async { Err::<i32, i32>(5) };
+    /// let future = future.map_ok_or_else(|x| x * 2, |x| x + 3);
+    /// assert_eq!(future.await, 10);
+    /// # });
+    /// ```
+    /// 
+    fn map_ok_or_else<T, E, F>(self, e: E, f: F) -> MapOkOrElse<Self, F, E>
+    where
+        F: FnOnce(Self::Ok) -> T,
+        E: FnOnce(Self::Error) -> T,
+        Self: Sized,
+    {
+        MapOkOrElse::new(self, e, f)
+    }
+
     /// Maps this future's error value to a different value.
     ///
     /// This method can be used to change the [`Error`](TryFuture::Error) type