From b0bff71211fa163bfb727ef4218c1656f454d08f Mon Sep 17 00:00:00 2001 From: Taiki Endo Date: Sat, 27 Mar 2021 00:46:19 +0900 Subject: [PATCH] Add stream::Peekable::{next_if, next_if_eq} --- futures-util/src/stream/mod.rs | 2 +- futures-util/src/stream/stream/mod.rs | 2 +- futures-util/src/stream/stream/peek.rs | 226 ++++++++++++++++++++++++- futures/tests/stream_peekable.rs | 26 +++ 4 files changed, 246 insertions(+), 10 deletions(-) diff --git a/futures-util/src/stream/mod.rs b/futures-util/src/stream/mod.rs index f3b2baa408..ecffdc7e41 100644 --- a/futures-util/src/stream/mod.rs +++ b/futures-util/src/stream/mod.rs @@ -19,7 +19,7 @@ pub use futures_core::stream::{FusedStream, Stream, TryStream}; mod stream; pub use self::stream::{ Chain, Collect, Concat, Cycle, Enumerate, Filter, FilterMap, FlatMap, Flatten, Fold, ForEach, - Fuse, Inspect, Map, Next, Peek, Peekable, Scan, SelectNextSome, Skip, SkipWhile, StreamExt, + Fuse, Inspect, Map, Next, NextIf, NextIfEq, Peek, Peekable, Scan, SelectNextSome, Skip, SkipWhile, StreamExt, StreamFuture, Take, TakeUntil, TakeWhile, Then, Unzip, Zip, }; diff --git a/futures-util/src/stream/stream/mod.rs b/futures-util/src/stream/stream/mod.rs index bade8119a2..8571afe4e6 100644 --- a/futures-util/src/stream/stream/mod.rs +++ b/futures-util/src/stream/stream/mod.rs @@ -123,7 +123,7 @@ pub use self::select_next_some::SelectNextSome; mod peek; #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 -pub use self::peek::{Peek, Peekable}; +pub use self::peek::{Peek, Peekable, NextIf, NextIfEq}; mod skip; #[allow(unreachable_pub)] // https://github.com/rust-lang/rust/issues/57411 diff --git a/futures-util/src/stream/stream/peek.rs b/futures-util/src/stream/stream/peek.rs index a4031102ab..266a96711b 100644 --- a/futures-util/src/stream/stream/peek.rs +++ b/futures-util/src/stream/stream/peek.rs @@ -1,5 +1,7 @@ +use crate::fns::FnOnce1; use crate::stream::{Fuse, StreamExt}; use core::fmt; +use core::marker::PhantomData; use core::pin::Pin; use futures_core::future::{FusedFuture, Future}; use futures_core::ready; @@ -44,10 +46,7 @@ impl Peekable { /// /// This method polls the underlying stream and return either a reference /// to the next item if the stream is ready or passes through any errors. - pub fn poll_peek( - self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { + pub fn poll_peek(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { let mut this = self.project(); Poll::Ready(loop { @@ -60,6 +59,94 @@ impl Peekable { } }) } + + /// Consume and return the next value of this stream if a condition is true. + /// + /// If `func` returns `true` for the next value of this stream, consume and return it. + /// Otherwise, return `None`. + /// + /// # Examples + /// + /// Consume a number if it's equal to 0. + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// use futures::pin_mut; + /// + /// let stream = stream::iter(0..5).peekable(); + /// pin_mut!(stream); + /// // The first item of the stream is 0; consume it. + /// assert_eq!(stream.as_mut().next_if(|&x| x == 0).await, Some(0)); + /// // The next item returned is now 1, so `consume` will return `false`. + /// assert_eq!(stream.as_mut().next_if(|&x| x == 0).await, None); + /// // `next_if` saves the value of the next item if it was not equal to `expected`. + /// assert_eq!(stream.next().await, Some(1)); + /// # }); + /// ``` + /// + /// Consume any number less than 10. + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// use futures::pin_mut; + /// + /// let stream = stream::iter(1..20).peekable(); + /// pin_mut!(stream); + /// // Consume all numbers less than 10 + /// while stream.as_mut().next_if(|&x| x < 10).await.is_some() {} + /// // The next value returned will be 10 + /// assert_eq!(stream.next().await, Some(10)); + /// # }); + /// ``` + pub fn next_if(self: Pin<&mut Self>, func: F) -> NextIf<'_, St, F> + where + F: FnOnce(&St::Item) -> bool, + { + NextIf { + inner: Some((self, func)), + } + } + + /// Consume and return the next item if it is equal to `expected`. + /// + /// # Example + /// + /// Consume a number if it's equal to 0. + /// + /// ``` + /// # futures::executor::block_on(async { + /// use futures::stream::{self, StreamExt}; + /// use futures::pin_mut; + /// + /// let stream = stream::iter(0..5).peekable(); + /// pin_mut!(stream); + /// // The first item of the stream is 0; consume it. + /// assert_eq!(stream.as_mut().next_if_eq(&0).await, Some(0)); + /// // The next item returned is now 1, so `consume` will return `false`. + /// assert_eq!(stream.as_mut().next_if_eq(&0).await, None); + /// // `next_if_eq` saves the value of the next item if it was not equal to `expected`. + /// assert_eq!(stream.next().await, Some(1)); + /// # }); + /// ``` + pub fn next_if_eq<'a, T>(self: Pin<&'a mut Self>, expected: &'a T) -> NextIfEq<'a, St, T> + where + T: ?Sized, + St::Item: PartialEq, + { + NextIfEq { + inner: NextIf { + inner: Some(( + self, + NextIfEqFn { + expected, + _next: PhantomData, + }, + )), + }, + } + } } impl FusedStream for Peekable { @@ -103,7 +190,7 @@ where } pin_project! { - /// Future for the [`Peekable::peek()`](self::Peekable::peek) function from [`Peekable`] + /// Future for the [`Peekable::peek`](self::Peekable::peek) method. #[must_use = "futures do nothing unless polled"] pub struct Peek<'a, St: Stream> { inner: Option>>, @@ -116,9 +203,7 @@ where St::Item: fmt::Debug, { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("Peek") - .field("inner", &self.inner) - .finish() + f.debug_struct("Peek").field("inner", &self.inner).finish() } } @@ -133,6 +218,7 @@ where St: Stream, { type Output = Option<&'a St::Item>; + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let inner = self.project().inner; if let Some(peekable) = inner { @@ -144,3 +230,127 @@ where } } } + +pin_project! { + /// Future for the [`Peekable::next_if`](self::Peekable::next_if) method. + #[must_use = "futures do nothing unless polled"] + pub struct NextIf<'a, St: Stream, F> { + inner: Option<(Pin<&'a mut Peekable>, F)>, + } +} + +impl fmt::Debug for NextIf<'_, St, F> +where + St: Stream + fmt::Debug, + St::Item: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("NextIf") + .field("inner", &self.inner.as_ref().map(|(s, _f)| s)) + .finish() + } +} + +#[allow(single_use_lifetimes)] // https://github.com/rust-lang/rust/issues/55058 +impl FusedFuture for NextIf<'_, St, F> +where + St: Stream, + F: for<'a> FnOnce1<&'a St::Item, Output = bool>, +{ + fn is_terminated(&self) -> bool { + self.inner.is_none() + } +} + +#[allow(single_use_lifetimes)] // https://github.com/rust-lang/rust/issues/55058 +impl Future for NextIf<'_, St, F> +where + St: Stream, + F: for<'a> FnOnce1<&'a St::Item, Output = bool>, +{ + type Output = Option; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let inner = self.project().inner; + if let Some((peekable, _)) = inner { + let res = ready!(peekable.as_mut().poll_next(cx)); + + let (peekable, func) = inner.take().unwrap(); + match res { + Some(ref matched) if func.call_once(matched) => Poll::Ready(res), + other => { + let peekable = peekable.project(); + // Since we called `self.next()`, we consumed `self.peeked`. + assert!(peekable.peeked.is_none()); + *peekable.peeked = other; + Poll::Ready(None) + } + } + } else { + panic!("NextIf polled after completion") + } + } +} + +pin_project! { + /// Future for the [`Peekable::next_if_eq`](self::Peekable::next_if_eq) method. + #[must_use = "futures do nothing unless polled"] + pub struct NextIfEq<'a, St: Stream, T: ?Sized> { + #[pin] + inner: NextIf<'a, St, NextIfEqFn<'a, T, St::Item>>, + } +} + +impl fmt::Debug for NextIfEq<'_, St, T> +where + St: Stream + fmt::Debug, + St::Item: fmt::Debug, + T: ?Sized, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("NextIfEq") + .field("inner", &self.inner.inner.as_ref().map(|(s, _f)| s)) + .finish() + } +} + +impl FusedFuture for NextIfEq<'_, St, T> +where + St: Stream, + T: ?Sized, + St::Item: PartialEq, +{ + fn is_terminated(&self) -> bool { + self.inner.is_terminated() + } +} + +impl Future for NextIfEq<'_, St, T> +where + St: Stream, + T: ?Sized, + St::Item: PartialEq, +{ + type Output = Option; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.project().inner.poll(cx) + } +} + +struct NextIfEqFn<'a, T: ?Sized, Item> { + expected: &'a T, + _next: PhantomData, +} + +impl FnOnce1<&Item> for NextIfEqFn<'_, T, Item> +where + T: ?Sized, + Item: PartialEq, +{ + type Output = bool; + + fn call_once(self, next: &Item) -> Self::Output { + next == self.expected + } +} diff --git a/futures/tests/stream_peekable.rs b/futures/tests/stream_peekable.rs index b65a0572cb..2fa7f3a4f5 100644 --- a/futures/tests/stream_peekable.rs +++ b/futures/tests/stream_peekable.rs @@ -11,3 +11,29 @@ fn peekable() { assert_eq!(peekable.collect::>().await, vec![1, 2, 3]); }); } + +#[test] +fn peekable_next_if_eq() { + block_on(async { + // first, try on references + let s = stream::iter(vec!["Heart", "of", "Gold"]).peekable(); + pin_mut!(s); + // try before `peek()` + assert_eq!(s.as_mut().next_if_eq(&"trillian").await, None); + assert_eq!(s.as_mut().next_if_eq(&"Heart").await, Some("Heart")); + // try after peek() + assert_eq!(s.as_mut().peek().await, Some(&"of")); + assert_eq!(s.as_mut().next_if_eq(&"of").await, Some("of")); + assert_eq!(s.as_mut().next_if_eq(&"zaphod").await, None); + // make sure `next()` still behaves + assert_eq!(s.next().await, Some("Gold")); + + // make sure comparison works for owned values + let s = stream::iter(vec![String::from("Ludicrous"), "speed".into()]).peekable(); + pin_mut!(s); + // make sure basic functionality works + assert_eq!(s.as_mut().next_if_eq("Ludicrous").await, Some("Ludicrous".into())); + assert_eq!(s.as_mut().next_if_eq("speed").await, Some("speed".into())); + assert_eq!(s.as_mut().next_if_eq("").await, None); + }); +}