Add StreamExt::try_scan

futures-util/src/stream/stream/mod.rs

@@ -156,6 +156,10 @@ pub use self::ready_chunks::ReadyChunks;
 mod scan;
 pub use self::scan::Scan;
 
+mod try_scan;
+#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
+pub use self::try_scan::TryScan;
+
 #[cfg_attr(target_os = "none", cfg(target_has_atomic = "ptr"))]
 #[cfg(feature = "alloc")]
 mod buffer_unordered;
@@ -1038,6 +1042,52 @@ pub trait StreamExt: Stream {
         assert_stream::<Self::Item, _>(TakeUntil::new(self, fut))
     }
 
+    /// Combinator similar to [`fold`](StreamExt::fold) that holds internal state
+    /// and produces a new stream.
+    ///
+    /// Accepts initial state and closure which will be applied to each element
+    /// of the stream until provided closure returns `None`. Once `None` is
+    /// returned, stream will be terminated.
+    ///
+    /// This method is similar to [`scan`](StreamExt::scan), but will
+    /// exit early if an error is encountered in either the stream or the
+    /// provided closure.
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// # futures::executor::block_on(async {
+    /// use futures::future;
+    /// use futures::stream::{self, StreamExt, TryStreamExt};
+    ///
+    /// let stream = stream::iter(1..=10);
+    ///
+    /// let stream = stream.try_scan(0, |mut state, x| {
+    ///     state += x;
+    ///     future::ready(if state < 10 { Ok::<_, ()>(Some((state, x))) } else { Ok(None) })
+    /// });
+    ///
+    /// assert_eq!(Ok(vec![1, 2, 3]), stream.try_collect::<Vec<_>>().await);
+    ///
+    /// let stream = stream::iter(1..=10);
+    ///
+    /// let stream = stream.try_scan(0, |mut state, x| {
+    ///     state += x;
+    ///     future::ready(if state < 10 { Ok(Some((state, x))) } else { Err(()) })
+    /// });
+    ///
+    /// assert_eq!(Err(()), stream.try_collect::<Vec<_>>().await);
+    /// # });
+    /// ```
+    fn try_scan<S, B, Fut, F>(self, initial_state: S, f: F) -> TryScan<Self, S, Fut, F>
+    where
+        F: FnMut(S, Self::Item) -> Fut,
+        Fut: TryFuture<Ok = Option<(S, B)>>,
+        Self: Sized,
+    {
+        assert_stream::<Result<B, Fut::Error>, _>(TryScan::new(self, initial_state, f))
+    }
+
     /// Runs this stream to completion, executing the provided asynchronous
     /// closure for each element on the stream.
     ///

futures-util/src/stream/stream/try_scan.rs

@@ -0,0 +1,127 @@
+use crate::unfold_state::UnfoldState;
+use core::fmt;
+use core::pin::Pin;
+use futures_core::stream::{FusedStream, Stream};
+use futures_core::task::{Context, Poll};
+use futures_core::{ready, TryFuture};
+#[cfg(feature = "sink")]
+use futures_sink::Sink;
+use pin_project_lite::pin_project;
+
+pin_project! {
+    /// TryStream for the [`try_scan`](super::StreamExt::try_scan) method.
+    #[must_use = "streams do nothing unless polled"]
+    pub struct TryScan<St, S, Fut, F> {
+        #[pin]
+        stream: St,
+        f: F,
+        #[pin]
+        state: UnfoldState<S, Fut>,
+    }
+}
+
+impl<St, S, Fut, F> fmt::Debug for TryScan<St, S, Fut, F>
+where
+    St: Stream + fmt::Debug,
+    St::Item: fmt::Debug,
+    S: fmt::Debug,
+    Fut: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("TryScan")
+            .field("stream", &self.stream)
+            .field("state", &self.state)
+            .field("done_taking", &self.is_done_taking())
+            .finish()
+    }
+}
+
+impl<St: Stream, S, Fut, F> TryScan<St, S, Fut, F> {
+    /// Checks if internal state is `None`.
+    fn is_done_taking(&self) -> bool {
+        matches!(self.state, UnfoldState::Empty)
+    }
+}
+
+impl<B, St, S, Fut, F> TryScan<St, S, Fut, F>
+where
+    St: Stream,
+    F: FnMut(S, St::Item) -> Fut,
+    Fut: TryFuture<Ok = Option<(S, B)>>,
+{
+    pub(super) fn new(stream: St, initial_state: S, f: F) -> Self {
+        Self { stream, f, state: UnfoldState::Value { value: initial_state } }
+    }
+}
+
+impl<B, St, S, Fut, F> Stream for TryScan<St, S, Fut, F>
+where
+    St: Stream,
+    F: FnMut(S, St::Item) -> Fut,
+    Fut: TryFuture<Ok = Option<(S, B)>>,
+{
+    type Item = Result<B, Fut::Error>;
+
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        if self.is_done_taking() {
+            return Poll::Ready(None);
+        }
+
+        let mut this = self.project();
+
+        Poll::Ready(loop {
+            if let Some(fut) = this.state.as_mut().project_future() {
+                match ready!(fut.try_poll(cx)) {
+                    Ok(None) => {
+                        this.state.set(UnfoldState::Empty);
+                        break None;
+                    }
+                    Ok(Some((state, item))) => {
+                        this.state.set(UnfoldState::Value { value: state });
+                        break Some(Ok(item));
+                    }
+                    Err(e) => {
+                        this.state.set(UnfoldState::Empty);
+                        break Some(Err(e));
+                    }
+                }
+            } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) {
+                let state = this.state.as_mut().take_value().unwrap();
+                this.state.set(UnfoldState::Future { future: (this.f)(state, item) })
+            } else {
+                break None;
+            }
+        })
+    }
+
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        if self.is_done_taking() {
+            (0, Some(0))
+        } else {
+            self.stream.size_hint() // can't know a lower bound, due to the predicate
+        }
+    }
+}
+
+impl<B, St, S, Fut, F> FusedStream for TryScan<St, S, Fut, F>
+where
+    St: FusedStream,
+    F: FnMut(S, St::Item) -> Fut,
+    Fut: TryFuture<Ok = Option<(S, B)>>,
+{
+    fn is_terminated(&self) -> bool {
+        self.is_done_taking()
+            || !matches!(self.state, UnfoldState::Future { .. }) && self.stream.is_terminated()
+    }
+}
+
+// Forwarding impl of Sink from the underlying stream
+#[cfg(feature = "sink")]
+impl<St, S, Fut, F, Item> Sink<Item> for TryScan<St, S, Fut, F>
+where
+    St: Stream + Sink<Item>,
+{
+    type Error = St::Error;
+
+    delegate_sink!(stream, Item);
+}