use pin_project

Author: starsheriff

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<S> {
-    stream: S,
-    count: usize,
+pin_project! {
+    #[doc(hidden)]
+    #[allow(missing_debug_implementations)]
+    pub struct CountFuture<S> {
+        #[pin]
+        stream: S,
+        count: usize,
+    }
 }
 
 impl<S> CountFuture<S> {
-    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<Self::Output> {
-        let next = futures_core::ready!(self.as_mut().stream().poll_next(cx));
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        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),
         }
     }
 }

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