use pin_project

new-scheduler
Johannes Weissmann 5 years ago
parent 7d2282dbd2
commit 37922408e5

@ -1,20 +1,22 @@
use std::future::Future;
use std::pin::Pin; use std::pin::Pin;
use crate::future::Future; use pin_project_lite::pin_project;
use crate::stream::Stream; use crate::stream::Stream;
use crate::task::{Context, Poll}; use crate::task::{Context, Poll};
#[doc(hidden)] pin_project! {
#[allow(missing_debug_implementations)] #[doc(hidden)]
pub struct CountFuture<S> { #[allow(missing_debug_implementations)]
stream: S, pub struct CountFuture<S> {
count: usize, #[pin]
stream: S,
count: usize,
}
} }
impl<S> CountFuture<S> { impl<S> CountFuture<S> {
pin_utils::unsafe_pinned!(stream: S);
pin_utils::unsafe_unpinned!(count: usize);
pub(crate) fn new(stream: S) -> Self { pub(crate) fn new(stream: S) -> Self {
CountFuture { stream, count: 0 } CountFuture { stream, count: 0 }
} }
@ -26,16 +28,17 @@ where
{ {
type Output = usize; type Output = usize;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let next = futures_core::ready!(self.as_mut().stream().poll_next(cx)); let this = self.project();
let next = futures_core::ready!(this.stream.poll_next(cx));
match next { match next {
Some(_) => { Some(_) => {
cx.waker().wake_by_ref(); cx.waker().wake_by_ref();
*self.as_mut().count() += 1; *this.count += 1;
Poll::Pending Poll::Pending
} }
None => Poll::Ready(self.count), None => Poll::Ready(*this.count),
} }
} }
} }

@ -1813,10 +1813,10 @@ extension_trait! {
# fn main() { async_std::task::block_on(async { # fn main() { async_std::task::block_on(async {
# #
use async_std::prelude::*; use async_std::prelude::*;
use std::collections::VecDeque; use async_std::stream;
let s1 = VecDeque::from(vec![0]); let s1 = stream::from_iter(vec![0]);
let s2 = VecDeque::from(vec![1, 2, 3]); let s2 = stream::from_iter(vec![1, 2, 3]);
assert_eq!(s1.count().await, 1); assert_eq!(s1.count().await, 1);
assert_eq!(s2.count().await, 3); assert_eq!(s2.count().await, 3);

Loading…
Cancel
Save