diff --git a/src/stream/mod.rs b/src/stream/mod.rs index 3b5f4610a..437fa9f1f 100644 --- a/src/stream/mod.rs +++ b/src/stream/mod.rs @@ -31,6 +31,7 @@ pub use repeat_with::{repeat_with, RepeatWith}; pub use stream::{ Chain, Filter, Fuse, Inspect, Scan, Skip, SkipWhile, StepBy, Stream, Take, TakeWhile, Zip, }; +pub use successors::{successors, Successors}; pub(crate) mod stream; @@ -39,6 +40,7 @@ mod from_fn; mod once; mod repeat; mod repeat_with; +mod successors; cfg_if! { if #[cfg(any(feature = "unstable", feature = "docs"))] { diff --git a/src/stream/successors.rs b/src/stream/successors.rs new file mode 100644 index 000000000..559554508 --- /dev/null +++ b/src/stream/successors.rs @@ -0,0 +1,61 @@ +use std::pin::Pin; + +use crate::stream::Stream; +use crate::task::{Context, Poll}; + +/// Creates a new stream where each successive item is computed based on the preceding one. +/// The stream starts with the given first item (if any) and calls the given FnMut(&T) -> Option closure to compute each item’s successor. +/// +/// # Examples +/// +/// ``` +/// # async_std::task::block_on(async { +/// # +/// use async_std::prelude::*; +/// use async_std::stream; +/// +/// let powers_of_10 = stream::successors(Some(1_u16), |n| n.checked_mul(10)); +/// assert_eq!(powers_of_10.collect::>().await, &[1, 10, 100, 1_000, 10_000]); +/// +/// # +/// # }) +/// ``` +pub fn successors(first: Option, succ: F) -> Successors +where + F: FnMut(&T) -> Option, +{ + Successors { next: first, succ } +} + +/// An new stream where each successive item is computed based on the preceding one. +/// +/// This `struct` is created by the [`successors`] function. +/// See its documentation for more. +/// +/// [`successors`]: fn.successors.html +#[derive(Clone, Debug)] +pub struct Successors { + next: Option, + succ: F, +} + +impl Successors { + pin_utils::unsafe_unpinned!(next: Option); + pin_utils::unsafe_unpinned!(succ: F); +} + +impl Stream for Successors +where + F: FnMut(&T) -> Option, +{ + type Item = T; + + fn poll_next(mut self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + let item = self.as_mut().next().take().map(|item| { + *self.as_mut().next() = (self.as_mut().succ())(&item); + item + }); + + Poll::Ready(item) + } +}