futures_util/stream/stream/

peek.rs

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;
use futures_core::stream::{FusedStream, Stream};
use futures_core::task::{Context, Poll};
#[cfg(feature = "sink")]
use futures_sink::Sink;
use pin_project_lite::pin_project;

pin_project! {
    /// A `Stream` that implements a `peek` method.
    ///
    /// The `peek` method can be used to retrieve a reference
    /// to the next `Stream::Item` if available. A subsequent
    /// call to `poll` will return the owned item.
    #[derive(Debug)]
    #[must_use = "streams do nothing unless polled"]
    pub struct Peekable<St: Stream> {
        #[pin]
        stream: Fuse<St>,
        peeked: Option<St::Item>,
    }
}

impl<St: Stream> Peekable<St> {
    pub(super) fn new(stream: St) -> Self {
        Self { stream: stream.fuse(), peeked: None }
    }

    delegate_access_inner!(stream, St, (.));

    /// Produces a future which retrieves a reference to the next item
    /// in the stream, or `None` if the underlying stream terminates.
    pub fn peek(self: Pin<&mut Self>) -> Peek<'_, St> {
        Peek { inner: Some(self) }
    }

    /// Peek retrieves a reference to the next item in the stream.
    ///
    /// 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<Option<&St::Item>> {
        let mut this = self.project();

        Poll::Ready(loop {
            if this.peeked.is_some() {
                break this.peeked.as_ref();
            } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) {
                *this.peeked = Some(item);
            } else {
                break None;
            }
        })
    }

    /// Produces a future which retrieves a mutable reference to the next item
    /// in the stream, or `None` if the underlying stream terminates.
    ///
    /// # Examples
    ///
    /// ```
    /// # futures::executor::block_on(async {
    /// use futures::stream::{self, StreamExt};
    /// use futures::pin_mut;
    ///
    /// let stream = stream::iter(vec![1, 2, 3]).peekable();
    /// pin_mut!(stream);
    ///
    /// assert_eq!(stream.as_mut().peek_mut().await, Some(&mut 1));
    /// assert_eq!(stream.as_mut().next().await, Some(1));
    ///
    /// // Peek into the stream and modify the value which will be returned next
    /// if let Some(p) = stream.as_mut().peek_mut().await {
    ///     if *p == 2 {
    ///         *p = 5;
    ///     }
    /// }
    ///
    /// assert_eq!(stream.collect::<Vec<_>>().await, vec![5, 3]);
    /// # });
    /// ```
    pub fn peek_mut(self: Pin<&mut Self>) -> PeekMut<'_, St> {
        PeekMut { inner: Some(self) }
    }

    /// Peek retrieves a mutable reference to the next item in the stream.
    pub fn poll_peek_mut(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Option<&mut St::Item>> {
        let mut this = self.project();

        Poll::Ready(loop {
            if this.peeked.is_some() {
                break this.peeked.as_mut();
            } else if let Some(item) = ready!(this.stream.as_mut().poll_next(cx)) {
                *this.peeked = Some(item);
            } else {
                break None;
            }
        })
    }

    /// Creates a future which will 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<F>(self: Pin<&mut Self>, func: F) -> NextIf<'_, St, F>
    where
        F: FnOnce(&St::Item) -> bool,
    {
        NextIf { inner: Some((self, func)) }
    }

    /// Creates a future which will 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<T>,
    {
        NextIfEq {
            inner: NextIf { inner: Some((self, NextIfEqFn { expected, _next: PhantomData })) },
        }
    }
}

impl<St: Stream> FusedStream for Peekable<St> {
    fn is_terminated(&self) -> bool {
        self.peeked.is_none() && self.stream.is_terminated()
    }
}

impl<S: Stream> Stream for Peekable<S> {
    type Item = S::Item;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.project();
        if let Some(item) = this.peeked.take() {
            return Poll::Ready(Some(item));
        }
        this.stream.poll_next(cx)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let peek_len = usize::from(self.peeked.is_some());
        let (lower, upper) = self.stream.size_hint();
        let lower = lower.saturating_add(peek_len);
        let upper = match upper {
            Some(x) => x.checked_add(peek_len),
            None => None,
        };
        (lower, upper)
    }
}

// Forwarding impl of Sink from the underlying stream
#[cfg(feature = "sink")]
impl<S, Item> Sink<Item> for Peekable<S>
where
    S: Sink<Item> + Stream,
{
    type Error = S::Error;

    delegate_sink!(stream, Item);
}

pin_project! {
    /// Future for the [`Peekable::peek`](self::Peekable::peek) method.
    #[must_use = "futures do nothing unless polled"]
    pub struct Peek<'a, St: Stream> {
        inner: Option<Pin<&'a mut Peekable<St>>>,
    }
}

impl<St> fmt::Debug for Peek<'_, St>
where
    St: Stream + fmt::Debug,
    St::Item: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Peek").field("inner", &self.inner).finish()
    }
}

impl<St: Stream> FusedFuture for Peek<'_, St> {
    fn is_terminated(&self) -> bool {
        self.inner.is_none()
    }
}

impl<'a, St> Future for Peek<'a, St>
where
    St: Stream,
{
    type Output = Option<&'a St::Item>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let inner = self.project().inner;
        if let Some(peekable) = inner {
            ready!(peekable.as_mut().poll_peek(cx));

            inner.take().unwrap().poll_peek(cx)
        } else {
            panic!("Peek polled after completion")
        }
    }
}

pin_project! {
    /// Future for the [`Peekable::peek_mut`](self::Peekable::peek_mut) method.
    #[must_use = "futures do nothing unless polled"]
    pub struct PeekMut<'a, St: Stream> {
        inner: Option<Pin<&'a mut Peekable<St>>>,
    }
}

impl<St> fmt::Debug for PeekMut<'_, St>
where
    St: Stream + fmt::Debug,
    St::Item: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("PeekMut").field("inner", &self.inner).finish()
    }
}

impl<St: Stream> FusedFuture for PeekMut<'_, St> {
    fn is_terminated(&self) -> bool {
        self.inner.is_none()
    }
}

impl<'a, St> Future for PeekMut<'a, St>
where
    St: Stream,
{
    type Output = Option<&'a mut St::Item>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let inner = self.project().inner;
        if let Some(peekable) = inner {
            ready!(peekable.as_mut().poll_peek_mut(cx));

            inner.take().unwrap().poll_peek_mut(cx)
        } else {
            panic!("PeekMut polled after completion")
        }
    }
}

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<St>>, F)>,
    }
}

impl<St, F> 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<St, F> 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<St, F> Future for NextIf<'_, St, F>
where
    St: Stream,
    F: for<'a> FnOnce1<&'a St::Item, Output = bool>,
{
    type Output = Option<St::Item>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        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<St, T> 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<St, T> FusedFuture for NextIfEq<'_, St, T>
where
    St: Stream,
    T: ?Sized,
    St::Item: PartialEq<T>,
{
    fn is_terminated(&self) -> bool {
        self.inner.is_terminated()
    }
}

impl<St, T> Future for NextIfEq<'_, St, T>
where
    St: Stream,
    T: ?Sized,
    St::Item: PartialEq<T>,
{
    type Output = Option<St::Item>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        self.project().inner.poll(cx)
    }
}

struct NextIfEqFn<'a, T: ?Sized, Item> {
    expected: &'a T,
    _next: PhantomData<Item>,
}

impl<T, Item> FnOnce1<&Item> for NextIfEqFn<'_, T, Item>
where
    T: ?Sized,
    Item: PartialEq<T>,
{
    type Output = bool;

    fn call_once(self, next: &Item) -> Self::Output {
        next == self.expected
    }
}