diff --git a/src/stream/successors.rs b/src/stream/successors.rs index 32353d0..e708416 100644 --- a/src/stream/successors.rs +++ b/src/stream/successors.rs @@ -14,11 +14,11 @@ use crate::task::{Context, Poll}; #[derive(Debug)] pub struct Successors where - Fut: Future, + Fut: Future>, { successor: F, future: Option, - next: T, + next: Option, _marker: PhantomData, } @@ -33,9 +33,9 @@ where /// use async_std::prelude::*; /// use async_std::stream; /// -/// let s = stream::successors(22, |val| { +/// let s = stream::successors(Some(22), |val| { /// async move { -/// val + 1 +/// Some(val + 1) /// } /// }); /// @@ -43,14 +43,25 @@ where /// assert_eq!(s.next().await, Some(23)); /// assert_eq!(s.next().await, Some(24)); /// assert_eq!(s.next().await, Some(25)); +/// +/// +///let never = stream::successors(None, |val: usize| { +/// async move { +/// Some(val + 1) +/// } +/// }); +/// +/// pin_utils::pin_mut!(never); +/// assert_eq!(never.next().await, None); +/// assert_eq!(never.next().await, None); /// # /// # }) } /// /// ``` -pub fn successors(start: T, func: F) -> Successors +pub fn successors(start: Option, func: F) -> Successors where F: FnMut(T) -> Fut, - Fut: Future, + Fut: Future>, T: Copy, { Successors { @@ -64,26 +75,30 @@ where impl Successors where F: FnMut(T) -> Fut, - Fut: Future, + Fut: Future>, T: Copy, { pin_utils::unsafe_unpinned!(successor: F); - pin_utils::unsafe_unpinned!(next: T); + pin_utils::unsafe_unpinned!(next: Option); pin_utils::unsafe_pinned!(future: Option); } impl Stream for Successors where - Fut: Future, + Fut: Future>, F: FnMut(T) -> Fut, T: Copy, { type Item = T; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.next.is_none() { + return Poll::Ready(None); + } + match &self.future { None => { - let x = self.next; + let x = self.next.unwrap(); let fut = (self.as_mut().successor())(x); self.as_mut().future().set(Some(fut)); } @@ -93,6 +108,6 @@ where let next = futures_core::ready!(self.as_mut().future().as_pin_mut().unwrap().poll(cx)); *self.as_mut().next() = next; self.as_mut().future().set(None); - Poll::Ready(Some(next)) + Poll::Ready(next) } }