From 3b9c7b1715777b6db69d420c1530aa845e6306c3 Mon Sep 17 00:00:00 2001 From: Artem Vorotnikov Date: Sat, 21 Dec 2019 07:17:05 +0300 Subject: [PATCH] stream: filtering utilities (#2001) Adds `StreamExt::filter` and `StreamExt::filter_map`. --- tokio/src/stream/filter.rs | 62 +++++++++++++++++++++++++ tokio/src/stream/filter_map.rs | 62 +++++++++++++++++++++++++ tokio/src/stream/mod.rs | 82 ++++++++++++++++++++++++++++++++++ 3 files changed, 206 insertions(+) create mode 100644 tokio/src/stream/filter.rs create mode 100644 tokio/src/stream/filter_map.rs diff --git a/tokio/src/stream/filter.rs b/tokio/src/stream/filter.rs new file mode 100644 index 000000000..88da15b7f --- /dev/null +++ b/tokio/src/stream/filter.rs @@ -0,0 +1,62 @@ +use crate::stream::Stream; + +use core::fmt; +use core::pin::Pin; +use core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream returned by the [`filter`](super::StreamExt::filter) method. + #[must_use = "streams do nothing unless polled"] + pub struct Filter { + #[pin] + stream: St, + f: F, + } +} + +impl fmt::Debug for Filter +where + St: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Filter") + .field("stream", &self.stream) + .finish() + } +} + +impl Filter +where + St: Stream, + F: FnMut(&St::Item) -> bool, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f } + } +} + +impl Stream for Filter +where + St: Stream, + F: FnMut(&St::Item) -> bool, +{ + type Item = St::Item; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + match ready!(self.as_mut().project().stream.poll_next(cx)) { + Some(e) => { + if (self.as_mut().project().f)(&e) { + return Poll::Ready(Some(e)); + } + } + None => return Poll::Ready(None), + } + } + } + + fn size_hint(&self) -> (usize, Option) { + (0, self.stream.size_hint().1) // can't know a lower bound, due to the predicate + } +} diff --git a/tokio/src/stream/filter_map.rs b/tokio/src/stream/filter_map.rs new file mode 100644 index 000000000..3aeb036f4 --- /dev/null +++ b/tokio/src/stream/filter_map.rs @@ -0,0 +1,62 @@ +use crate::stream::Stream; + +use core::fmt; +use core::pin::Pin; +use core::task::{Context, Poll}; +use pin_project_lite::pin_project; + +pin_project! { + /// Stream returned by the [`filter_map`](super::StreamExt::filter_map) method. + #[must_use = "streams do nothing unless polled"] + pub struct FilterMap { + #[pin] + stream: St, + f: F, + } +} + +impl fmt::Debug for FilterMap +where + St: fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("FilterMap") + .field("stream", &self.stream) + .finish() + } +} + +impl FilterMap +where + St: Stream, + F: FnMut(St::Item) -> Option, +{ + pub(super) fn new(stream: St, f: F) -> Self { + Self { stream, f } + } +} + +impl Stream for FilterMap +where + St: Stream, + F: FnMut(St::Item) -> Option, +{ + type Item = T; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + match ready!(self.as_mut().project().stream.poll_next(cx)) { + Some(e) => { + if let Some(e) = (self.as_mut().project().f)(e) { + return Poll::Ready(Some(e)); + } + } + None => return Poll::Ready(None), + } + } + } + + fn size_hint(&self) -> (usize, Option) { + (0, self.stream.size_hint().1) // can't know a lower bound, due to the predicate + } +} diff --git a/tokio/src/stream/mod.rs b/tokio/src/stream/mod.rs index 2d48bc7ed..329ee8a96 100644 --- a/tokio/src/stream/mod.rs +++ b/tokio/src/stream/mod.rs @@ -4,6 +4,12 @@ //! //! This module provides helpers to work with them. +mod filter; +use filter::Filter; + +mod filter_map; +use filter_map::FilterMap; + mod iter; pub use iter::{iter, Iter}; @@ -88,6 +94,82 @@ pub trait StreamExt: Stream { { Map::new(self, f) } + + /// Filters the values produced by this stream according to the provided + /// predicate. + /// + /// As values of this stream are made available, the provided predicate `f` + /// will be run against them. If the predicate + /// resolves to `true`, then the stream will yield the value, but if the + /// predicate resolves to `false`, then the value + /// will be discarded and the next value will be produced. + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it, similar to [`Iterator::filter`] method in the + /// standard library. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use tokio::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=8); + /// let mut evens = stream.filter(|x| x % 2 == 0); + /// + /// assert_eq!(Some(2), evens.next().await); + /// assert_eq!(Some(4), evens.next().await); + /// assert_eq!(Some(6), evens.next().await); + /// assert_eq!(Some(8), evens.next().await); + /// assert_eq!(None, evens.next().await); + /// # } + /// ``` + fn filter(self, f: F) -> Filter + where + F: FnMut(&Self::Item) -> bool, + Self: Sized, + { + Filter::new(self, f) + } + + /// Filters the values produced by this stream while simultaneously mapping + /// them to a different type according to the provided closure. + /// + /// As values of this stream are made available, the provided function will + /// be run on them. If the predicate `f` resolves to + /// [`Some(item)`](Some) then the stream will yield the value `item`, but if + /// it resolves to [`None`] then the next value will be produced. + /// + /// Note that this function consumes the stream passed into it and returns a + /// wrapped version of it, similar to [`Iterator::filter_map`] method in the + /// standard library. + /// + /// # Examples + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use tokio::stream::{self, StreamExt}; + /// + /// let stream = stream::iter(1..=8); + /// let mut evens = stream.filter_map(|x| { + /// if x % 2 == 0 { Some(x + 1) } else { None } + /// }); + /// + /// assert_eq!(Some(3), evens.next().await); + /// assert_eq!(Some(5), evens.next().await); + /// assert_eq!(Some(7), evens.next().await); + /// assert_eq!(Some(9), evens.next().await); + /// assert_eq!(None, evens.next().await); + /// # } + /// ``` + fn filter_map(self, f: F) -> FilterMap + where + F: FnMut(Self::Item) -> Option, + Self: Sized, + { + FilterMap::new(self, f) + } } impl StreamExt for T where T: Stream {}