Add Stream::scan

pull/192/head
Wonwoo Choi 5 years ago
parent e060326910
commit 50a7db2af4

@ -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;

@ -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 {

@ -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)
})
})
}
}
Loading…
Cancel
Save