Change StreamExt::scan to pass state to closure by value

Author: taiki-e

futures-util/src/sink/unfold.rs

@@ -10,10 +10,10 @@ pin_project! {
     /// Sink for the [`unfold`] function.
     #[derive(Debug)]
     #[must_use = "sinks do nothing unless polled"]
-    pub struct Unfold<T, F, R> {
+    pub struct Unfold<T, F, Fut> {
         function: F,
         #[pin]
-        state: UnfoldState<T, R>,
+        state: UnfoldState<T, Fut>,
     }
 }
 
@@ -36,18 +36,18 @@ pin_project! {
 /// unfold.send(5).await?;
 /// # Ok::<(), std::convert::Infallible>(()) }).unwrap();
 /// ```
-pub fn unfold<T, F, R, Item, E>(init: T, function: F) -> Unfold<T, F, R>
+pub fn unfold<T, F, Fut, Item, E>(init: T, function: F) -> Unfold<T, F, Fut>
 where
-    F: FnMut(T, Item) -> R,
-    R: Future<Output = Result<T, E>>,
+    F: FnMut(T, Item) -> Fut,
+    Fut: Future<Output = Result<T, E>>,
 {
     assert_sink::<Item, E, _>(Unfold { function, state: UnfoldState::Value { value: init } })
 }
 
-impl<T, F, R, Item, E> Sink<Item> for Unfold<T, F, R>
+impl<T, F, Fut, Item, E> Sink<Item> for Unfold<T, F, Fut>
 where
-    F: FnMut(T, Item) -> R,
-    R: Future<Output = Result<T, E>>,
+    F: FnMut(T, Item) -> Fut,
+    Fut: Future<Output = Result<T, E>>,
 {
     type Error = E;
 

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

@@ -715,6 +715,10 @@ pub trait StreamExt: Stream {
     /// of the stream until provided closure returns `None`. Once `None` is
     /// returned, stream will be terminated.
     ///
+    /// Unlike [`Iterator::scan`], the closure takes the state by value instead of
+    /// mutable reference to avoid [the limitation of the async
+    /// block](https://github.com/rust-lang/futures-rs/issues/2171).
+    ///
     /// # Examples
     ///
     /// ```
@@ -724,18 +728,22 @@ pub trait StreamExt: Stream {
     ///
     /// let stream = stream::iter(1..=10);
     ///
-    /// let stream = stream.scan(0, |state, x| {
-    ///     *state += x;
-    ///     future::ready(if *state < 10 { Some(x) } else { None })
+    /// let stream = stream.scan(0, |mut state, x| async move {
+    ///     state += x;
+    ///     if state < 10 {
+    ///         Some((state, x))
+    ///     } else {
+    ///         None
+    ///     }
     /// });
     ///
     /// assert_eq!(vec![1, 2, 3], stream.collect::<Vec<_>>().await);
     /// # });
     /// ```
     fn scan<S, B, Fut, F>(self, initial_state: S, f: F) -> Scan<Self, S, Fut, F>
     where
-        F: FnMut(&mut S, Self::Item) -> Fut,
-        Fut: Future<Output = Option<B>>,
+        F: FnMut(S, Self::Item) -> Fut,
+        Fut: Future<Output = Option<(S, B)>>,
         Self: Sized,
     {
         assert_stream::<B, _>(Scan::new(self, initial_state, f))

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

@@ -1,3 +1,4 @@
+use crate::unfold_state::UnfoldState;
 use core::fmt;
 use core::pin::Pin;
 use futures_core::future::Future;
@@ -8,20 +9,15 @@ use futures_core::task::{Context, Poll};
 use futures_sink::Sink;
 use pin_project_lite::pin_project;
 
-struct StateFn<S, F> {
-    state: S,
-    f: F,
-}
-
 pin_project! {
     /// Stream for the [`scan`](super::StreamExt::scan) method.
     #[must_use = "streams do nothing unless polled"]
     pub struct Scan<St: Stream, S, Fut, F> {
         #[pin]
         stream: St,
-        state_f: Option<StateFn<S, F>>,
+        f: F,
         #[pin]
-        future: Option<Fut>,
+        state: UnfoldState<S, Fut>,
     }
 }
 
@@ -35,8 +31,7 @@ where
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         f.debug_struct("Scan")
             .field("stream", &self.stream)
-            .field("state", &self.state_f.as_ref().map(|s| &s.state))
-            .field("future", &self.future)
+            .field("state", &self.state)
             .field("done_taking", &self.is_done_taking())
             .finish()
     }
@@ -45,18 +40,18 @@ where
 impl<St: Stream, S, Fut, F> Scan<St, S, Fut, F> {
     /// Checks if internal state is `None`.
     fn is_done_taking(&self) -> bool {
-        self.state_f.is_none()
+        self.state.is_empty()
     }
 }
 
 impl<B, St, S, Fut, F> Scan<St, S, Fut, F>
 where
     St: Stream,
-    F: FnMut(&mut S, St::Item) -> Fut,
-    Fut: Future<Output = Option<B>>,
+    F: FnMut(S, St::Item) -> Fut,
+    Fut: Future<Output = Option<(S, B)>>,
 {
     pub(super) fn new(stream: St, initial_state: S, f: F) -> Self {
-        Self { stream, state_f: Some(StateFn { state: initial_state, f }), future: None }
+        Self { stream, f, state: UnfoldState::Value { value: initial_state } }
     }
 
     delegate_access_inner!(stream, St, ());
@@ -65,8 +60,8 @@ where
 impl<B, St, S, Fut, F> Stream for Scan<St, S, Fut, F>
 where
     St: Stream,
-    F: FnMut(&mut S, St::Item) -> Fut,
-    Fut: Future<Output = Option<B>>,
+    F: FnMut(S, St::Item) -> Fut,
+    Fut: Future<Output = Option<(S, B)>>,
 {
     type Item = B;
 
@@ -78,18 +73,22 @@ where
         let mut this = self.project();
 
         Poll::Ready(loop {
-            if let Some(fut) = this.future.as_mut().as_pin_mut() {
+            if let Some(fut) = this.state.as_mut().project_future() {
                 let item = ready!(fut.poll(cx));
-                this.future.set(None);
 
-                if item.is_none() {
-                    *this.state_f = None;
+                match item {
+                    None => {
+                        this.state.set(UnfoldState::Empty);
+                        break None;
+                    }
+                    Some((state, item)) => {
+                        this.state.set(UnfoldState::Value { value: state });
+                        break Some(item);
+                    }
                 }
-
-                break item;
             } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) {
-                let state_f = this.state_f.as_mut().unwrap();
-                this.future.set(Some((state_f.f)(&mut state_f.state, item)))
+                let state = this.state.as_mut().take_value().unwrap();
+                this.state.set(UnfoldState::Future { future: (this.f)(state, item) })
             } else {
                 break None;
             }
@@ -108,11 +107,11 @@ where
 impl<B, St, S, Fut, F> FusedStream for Scan<St, S, Fut, F>
 where
     St: FusedStream,
-    F: FnMut(&mut S, St::Item) -> Fut,
-    Fut: Future<Output = Option<B>>,
+    F: FnMut(S, St::Item) -> Fut,
+    Fut: Future<Output = Option<(S, B)>>,
 {
     fn is_terminated(&self) -> bool {
-        self.is_done_taking() || self.future.is_none() && self.stream.is_terminated()
+        self.is_done_taking() || !self.state.is_future() && self.stream.is_terminated()
     }
 }
 

futures-util/src/unfold_state.rs

@@ -7,20 +7,34 @@ pin_project! {
     #[project = UnfoldStateProj]
     #[project_replace = UnfoldStateProjReplace]
     #[derive(Debug)]
-    pub(crate) enum UnfoldState<T, R> {
+    pub(crate) enum UnfoldState<T, Fut> {
         Value {
             value: T,
         },
         Future {
             #[pin]
-            future: R,
+            future: Fut,
         },
         Empty,
     }
 }
 
-impl<T, R> UnfoldState<T, R> {
-    pub(crate) fn project_future(self: Pin<&mut Self>) -> Option<Pin<&mut R>> {
+impl<T, Fut> UnfoldState<T, Fut> {
+    pub(crate) fn is_empty(&self) -> bool {
+        match self {
+            Self::Empty => true,
+            _ => false,
+        }
+    }
+
+    pub(crate) fn is_future(&self) -> bool {
+        match self {
+            Self::Future { .. } => true,
+            _ => false,
+        }
+    }
+
+    pub(crate) fn project_future(self: Pin<&mut Self>) -> Option<Pin<&mut Fut>> {
         match self.project() {
             UnfoldStateProj::Future { future } => Some(future),
             _ => None,

futures/tests/stream.rs

@@ -38,15 +38,32 @@ fn flat_map() {
 fn scan() {
     block_on(async {
         let values = stream::iter(vec![1u8, 2, 3, 4, 6, 8, 2])
-            .scan(1, |state, e| {
-                *state += 1;
-                futures::future::ready(if e < *state { Some(e) } else { None })
+            .scan(1, |mut state, e| async move {
+                state += 1;
+                if e < state {
+                    Some((state, e))
+                } else {
+                    None
+                }
             })
             .collect::<Vec<_>>()
             .await;
 
         assert_eq!(values, vec![1u8, 2, 3, 4]);
     });
+
+    block_on(async {
+        let mut state = vec![];
+        let values = stream::iter(vec![1u8, 2, 3, 4, 6, 8, 2])
+            .scan(&mut state, |state, e| async move {
+                state.push(e);
+                Some((state, e))
+            })
+            .collect::<Vec<_>>()
+            .await;
+
+        assert_eq!(values, state);
+    });
 }
 
 #[test]