Add stream::try_unfold

futures-util/src/stream/mod.rs

@@ -38,9 +38,9 @@ pub use self::stream::{ReuniteError, SplitSink, SplitStream};
 
 mod try_stream;
 pub use self::try_stream::{
-    AndThen, ErrInto, InspectErr, InspectOk, IntoStream, MapErr, MapOk, OrElse, TryCollect,
-    TryConcat, TryFilter, TryFilterMap, TryFlatten, TryFold, TryForEach, TryNext, TrySkipWhile,
-    TryStreamExt,
+    try_unfold, AndThen, ErrInto, InspectErr, InspectOk, IntoStream, MapErr, MapOk, OrElse,
+    TryCollect, TryConcat, TryFilter, TryFilterMap, TryFlatten, TryFold, TryForEach, TryNext,
+    TrySkipWhile, TryStreamExt, TryUnfold,
 };
 
 #[cfg(feature = "io")]

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

@@ -76,6 +76,10 @@ mod try_fold;
 #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
 pub use self::try_fold::TryFold;
 
+mod try_unfold;
+#[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
+pub use self::try_unfold::{try_unfold, TryUnfold};
+
 mod try_skip_while;
 #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411
 pub use self::try_skip_while::TrySkipWhile;

futures-util/src/stream/try_stream/try_unfold.rs

@@ -0,0 +1,134 @@
+use core::fmt;
+use core::pin::Pin;
+use futures_core::future::TryFuture;
+use futures_core::stream::Stream;
+use futures_core::task::{Context, Poll};
+use pin_utils::{unsafe_pinned, unsafe_unpinned};
+
+/// Creates a `TryStream` from a seed and a closure returning a `TryFuture`.
+///
+/// This function is the dual for the `TryStream::try_fold()` adapter: while
+/// `TryStream::try_fold()` reduces a `TryStream` to one single value,
+/// `try_unfold()` creates a `TryStream` from a seed value.
+///
+/// `try_unfold()` will call the provided closure with the provided seed, then
+/// wait for the returned `TryFuture` to complete with `(a, b)`. It will then
+/// yield the value `a`, and use `b` as the next internal state.
+///
+/// If the closure returns `None` instead of `Some(TryFuture)`, then the
+/// `try_unfold()` will stop producing items and return `Poll::Ready(None)` in
+/// future calls to `poll()`.
+///
+/// In case of error generated by the returned `TryFuture`, the error will be
+/// returned by the `TryStream`. The `TryStream` will then yield
+/// `Poll::Ready(None)` in future calls to `poll()`.
+///
+/// This function can typically be used when wanting to go from the "world of
+/// futures" to the "world of streams": the provided closure can build a
+/// `TryFuture` using other library functions working on futures, and
+/// `try_unfold()` will turn it into a `TryStream` by repeating the operation.
+///
+/// # Example
+///
+/// ```
+/// # #[derive(Debug, PartialEq)]
+/// # struct SomeError;
+/// # futures::executor::block_on(async {
+/// use futures::stream::{self, TryStreamExt};
+///
+/// let stream = stream::try_unfold(0, |state| async move {
+///     if state < 0 {
+///         return Err(SomeError);
+///     }
+///
+///     if state <= 2 {
+///         let next_state = state + 1;
+///         let yielded = state * 2;
+///         Ok(Some((yielded, next_state)))
+///     } else {
+///         Ok(None)
+///     }
+/// });
+///
+/// let result: Result<Vec<i32>, _> = stream.try_collect().await;
+/// assert_eq!(result, Ok(vec![0, 2, 4]));
+/// # });
+/// ```
+pub fn try_unfold<T, F, Fut, Item>(init: T, f: F) -> TryUnfold<T, F, Fut>
+where
+    F: FnMut(T) -> Fut,
+    Fut: TryFuture<Ok = Option<(Item, T)>>,
+{
+    TryUnfold {
+        f,
+        state: Some(init),
+        fut: None,
+    }
+}
+
+/// Stream for the [`try_unfold`] function.
+#[must_use = "streams do nothing unless polled"]
+pub struct TryUnfold<T, F, Fut> {
+    f: F,
+    state: Option<T>,
+    fut: Option<Fut>,
+}
+
+impl<T, F, Fut: Unpin> Unpin for TryUnfold<T, F, Fut> {}
+
+impl<T, F, Fut> fmt::Debug for TryUnfold<T, F, Fut>
+where
+    T: fmt::Debug,
+    Fut: fmt::Debug,
+{
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("TryUnfold")
+            .field("state", &self.state)
+            .field("fut", &self.fut)
+            .finish()
+    }
+}
+
+impl<T, F, Fut> TryUnfold<T, F, Fut> {
+    unsafe_unpinned!(f: F);
+    unsafe_unpinned!(state: Option<T>);
+    unsafe_pinned!(fut: Option<Fut>);
+}
+
+impl<T, F, Fut, Item> Stream for TryUnfold<T, F, Fut>
+where
+    F: FnMut(T) -> Fut,
+    Fut: TryFuture<Ok = Option<(Item, T)>>,
+{
+    type Item = Result<Item, Fut::Error>;
+
+    fn poll_next(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<Result<Item, Fut::Error>>> {
+        if let Some(state) = self.as_mut().state().take() {
+            let fut = (self.as_mut().f())(state);
+            self.as_mut().fut().set(Some(fut));
+        }
+
+        match self.as_mut().fut().as_pin_mut() {
+            None => {
+                // The future previously errored
+                Poll::Ready(None)
+            }
+            Some(fut) => {
+                let step = ready!(fut.try_poll(cx));
+                self.as_mut().fut().set(None);
+
+                match step {
+                    Ok(Some((item, next_state))) => {
+                        *self.as_mut().state() = Some(next_state);
+                        Poll::Ready(Some(Ok(item)))
+                    }
+                    Ok(None) => Poll::Ready(None),
+                    Err(e) => Poll::Ready(Some(Err(e))),
+                }
+            }
+        }
+    }
+}

futures/src/lib.rs

@@ -439,6 +439,7 @@ pub mod stream {
         poll_fn, PollFn,
         select, Select,
         unfold, Unfold,
+        try_unfold, TryUnfold,
 
         StreamExt,
         Chain, Collect, Concat, Enumerate, Filter, FilterMap, Flatten, Fold,