diff --git a/src/stream/stream/count.rs b/src/stream/stream/count.rs index b6d53ca..221b0f0 100644 --- a/src/stream/stream/count.rs +++ b/src/stream/stream/count.rs @@ -1,20 +1,22 @@ +use std::future::Future; use std::pin::Pin; -use crate::future::Future; +use pin_project_lite::pin_project; + use crate::stream::Stream; use crate::task::{Context, Poll}; -#[doc(hidden)] -#[allow(missing_debug_implementations)] -pub struct CountFuture { - stream: S, - count: usize, +pin_project! { + #[doc(hidden)] + #[allow(missing_debug_implementations)] + pub struct CountFuture { + #[pin] + stream: S, + count: usize, + } } impl CountFuture { - pin_utils::unsafe_pinned!(stream: S); - pin_utils::unsafe_unpinned!(count: usize); - pub(crate) fn new(stream: S) -> Self { CountFuture { stream, count: 0 } } @@ -26,16 +28,17 @@ where { type Output = usize; - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let next = futures_core::ready!(self.as_mut().stream().poll_next(cx)); + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = self.project(); + let next = futures_core::ready!(this.stream.poll_next(cx)); match next { Some(_) => { cx.waker().wake_by_ref(); - *self.as_mut().count() += 1; + *this.count += 1; Poll::Pending } - None => Poll::Ready(self.count), + None => Poll::Ready(*this.count), } } } diff --git a/src/stream/stream/mod.rs b/src/stream/stream/mod.rs index 2f21c38..d659227 100644 --- a/src/stream/stream/mod.rs +++ b/src/stream/stream/mod.rs @@ -1813,10 +1813,10 @@ extension_trait! { # fn main() { async_std::task::block_on(async { # use async_std::prelude::*; - use std::collections::VecDeque; + use async_std::stream; - let s1 = VecDeque::from(vec![0]); - let s2 = VecDeque::from(vec![1, 2, 3]); + let s1 = stream::from_iter(vec![0]); + let s2 = stream::from_iter(vec![1, 2, 3]); assert_eq!(s1.count().await, 1); assert_eq!(s2.count().await, 3);