From 50a7db2af47ac00845add30bb7b6b14d1e8e342d Mon Sep 17 00:00:00 2001
From: Wonwoo Choi <chwo9843@gmail.com>
Date: Sat, 14 Sep 2019 23:13:36 +0900
Subject: [PATCH] Add Stream::scan

---
 src/stream/mod.rs         |  2 +-
 src/stream/stream/mod.rs  | 45 +++++++++++++++++++++++++++++++++++++++
 src/stream/stream/scan.rs | 41 +++++++++++++++++++++++++++++++++++
 3 files changed, 87 insertions(+), 1 deletion(-)
 create mode 100644 src/stream/stream/scan.rs

diff --git a/src/stream/mod.rs b/src/stream/mod.rs
index 6081912..3146bd3 100644
--- a/src/stream/mod.rs
+++ b/src/stream/mod.rs
@@ -25,7 +25,7 @@ pub use double_ended_stream::DoubleEndedStream;
 pub use empty::{empty, Empty};
 pub use once::{once, Once};
 pub use repeat::{repeat, Repeat};
-pub use stream::{Stream, Take};
+pub use stream::{Scan, Stream, Take};
 
 mod double_ended_stream;
 mod empty;
diff --git a/src/stream/stream/mod.rs b/src/stream/stream/mod.rs
index fb8b38d..9bcccb6 100644
--- a/src/stream/stream/mod.rs
+++ b/src/stream/stream/mod.rs
@@ -29,8 +29,10 @@ mod find_map;
 mod min_by;
 mod next;
 mod nth;
+mod scan;
 mod take;
 
+pub use scan::Scan;
 pub use take::Take;
 
 use all::AllFuture;
@@ -501,6 +503,49 @@ pub trait Stream {
             f,
         }
     }
+
+    /// A stream adaptor similar to [`fold`] that holds internal state and produces a new stream.
+    ///
+    /// [`fold`]: #method.fold
+    ///
+    /// `scan()` takes two arguments: an initial value which seeds the internal state, and a
+    /// closure with two arguments, the first being a mutable reference to the internal state and
+    /// the second a stream element. The closure can assign to the internal state to share state
+    /// between iterations.
+    ///
+    /// On iteration, the closure will be applied to each element of the stream and the return
+    /// value from the closure, an `Option`, is yielded by the stream.
+    ///
+    /// ## Examples
+    ///
+    /// ```
+    /// # fn main() { async_std::task::block_on(async {
+    /// #
+    /// use std::collections::VecDeque;
+    /// use async_std::stream::Stream;
+    ///
+    /// let s: VecDeque<isize> = vec![1, 2, 3].into_iter().collect();
+    /// let mut s = s.scan(1, |state, x| {
+    ///     *state = *state * x;
+    ///     Some(-*state)
+    /// });
+    ///
+    /// assert_eq!(s.next().await, Some(-1));
+    /// assert_eq!(s.next().await, Some(-2));
+    /// assert_eq!(s.next().await, Some(-6));
+    /// assert_eq!(s.next().await, None);
+    /// #
+    /// # }) }
+    /// ```
+    #[inline]
+    fn scan<St, B, F>(self, initial_state: St, f: F) -> Scan<Self, St, F>
+    where
+        Self: Sized,
+        St: Unpin,
+        F: Unpin + FnMut(&mut St, Self::Item) -> Option<B>,
+    {
+        Scan::new(self, initial_state, f)
+    }
 }
 
 impl<T: futures_core::stream::Stream + Unpin + ?Sized> Stream for T {
diff --git a/src/stream/stream/scan.rs b/src/stream/stream/scan.rs
new file mode 100644
index 0000000..0776a7b
--- /dev/null
+++ b/src/stream/stream/scan.rs
@@ -0,0 +1,41 @@
+use crate::task::{Context, Poll};
+
+use std::pin::Pin;
+
+/// A stream to maintain state while polling another stream.
+#[derive(Debug)]
+pub struct Scan<S, St, F> {
+    stream: S,
+    state_f: (St, F),
+}
+
+impl<S, St: Unpin, F: Unpin> Scan<S, St, F> {
+    pub(crate) fn new(stream: S, initial_state: St, f: F) -> Self {
+        Self {
+            stream,
+            state_f: (initial_state, f),
+        }
+    }
+
+    pin_utils::unsafe_pinned!(stream: S);
+    pin_utils::unsafe_unpinned!(state_f: (St, F));
+}
+
+impl<S, St, F, B> futures_core::stream::Stream for Scan<S, St, F>
+where
+    S: futures_core::stream::Stream,
+    St: Unpin,
+    F: Unpin + FnMut(&mut St, S::Item) -> Option<B>,
+{
+    type Item = B;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<B>> {
+        let poll_result = self.as_mut().stream().poll_next(cx);
+        poll_result.map(|item| {
+            item.and_then(|item| {
+                let (state, f) = self.as_mut().state_f();
+                f(state, item)
+            })
+        })
+    }
+}