|
|
|
@ -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};
|
|
|
|
|
|
|
|
|
|
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),
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|
}
|
|
|
|
|