diff --git a/src/stream/double_ended/mod.rs b/src/stream/double_ended/mod.rs index 101ccbd..982c2c7 100644 --- a/src/stream/double_ended/mod.rs +++ b/src/stream/double_ended/mod.rs @@ -1,9 +1,14 @@ +mod nth_back; + +use nth_back::NthBackFuture; + extension_trait! { use crate::stream::Stream; use std::pin::Pin; use std::task::{Context, Poll}; + #[doc = r#" Something fancy "#] @@ -17,6 +22,36 @@ extension_trait! { Something else "#] pub trait DoubleEndedStreamExt: crate::stream::DoubleEndedStream { + + #[doc = r#" + Returns the nth element from the back of the stream. + + # Examples + + Basic usage: + + ``` + # fn main() { async_std::task::block_on(async { + # + use async_std::stream::double_ended::DoubleEndedStreamExt; + use async_std::stream; + + let mut s = stream::from_iter(vec![1u8, 2, 3, 4, 5]); + + let second = s.nth_back(1).await; + assert_eq!(second, Some(4)); + # + # }) } + ``` + "#] + fn nth_back( + &mut self, + n: usize, + ) -> impl Future> + '_ [NthBackFuture<'_, Self>] + where + Self: Unpin + Sized, + { + NthBackFuture::new(self, n) + } } } - diff --git a/src/stream/double_ended/nth_back.rs b/src/stream/double_ended/nth_back.rs new file mode 100644 index 0000000..e318e79 --- /dev/null +++ b/src/stream/double_ended/nth_back.rs @@ -0,0 +1,39 @@ +use std::pin::Pin; +use std::task::{Context, Poll}; +use std::future::Future; + +use crate::stream::DoubleEndedStream; + +pub struct NthBackFuture<'a, S> { + stream: &'a mut S, + n: usize, +} + +impl<'a, S> NthBackFuture<'a, S> { + pub(crate) fn new(stream: &'a mut S, n: usize) -> Self { + NthBackFuture { stream, n } + } +} + +impl<'a, S> Future for NthBackFuture<'a, S> +where + S: DoubleEndedStream + Sized + Unpin, +{ + type Output = Option; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let next = futures_core::ready!(Pin::new(&mut *self.stream).poll_next_back(cx)); + match next { + Some(v) => match self.n { + 0 => Poll::Ready(Some(v)), + _ => { + self.n -= 1; + cx.waker().wake_by_ref(); + Poll::Pending + } + }, + None => Poll::Ready(None), + } + } +} + diff --git a/src/stream/mod.rs b/src/stream/mod.rs index e15d081..318733b 100644 --- a/src/stream/mod.rs +++ b/src/stream/mod.rs @@ -318,7 +318,7 @@ mod repeat; mod repeat_with; cfg_unstable! { - mod double_ended; + pub mod double_ended; mod double_ended_stream; mod exact_size_stream; mod extend;