diff --git a/tokio-stream/src/stream_ext/peekable.rs b/tokio-stream/src/stream_ext/peekable.rs index 4963a3922..a0452f8b5 100644 --- a/tokio-stream/src/stream_ext/peekable.rs +++ b/tokio-stream/src/stream_ext/peekable.rs @@ -34,6 +34,33 @@ impl Peekable { self.peek.as_ref() } } + + /// Peek at the next item in the stream as a mutable reference. + pub async fn peek_mut(&mut self) -> Option<&mut T::Item> + where + T: Unpin, + { + if let Some(ref mut it) = self.peek { + Some(it) + } else { + self.peek = self.next().await; + self.peek.as_mut() + } + } + + /// Poll to peek at the next item in the stream as a mutable reference. + pub fn poll_peek(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let mut this = self.project(); + + if this.peek.is_none() { + match this.stream.as_mut().poll_next(cx) { + Poll::Ready(item) => *this.peek = item, + Poll::Pending => return Poll::Pending, + } + } + + Poll::Ready(this.peek.as_mut()) + } } impl Stream for Peekable { diff --git a/tokio-stream/tests/stream_peekable.rs b/tokio-stream/tests/stream_peekable.rs index 44cd8b733..3b9ef2f19 100644 --- a/tokio-stream/tests/stream_peekable.rs +++ b/tokio-stream/tests/stream_peekable.rs @@ -1,4 +1,8 @@ use tokio_stream::{self as stream, Stream, StreamExt}; +use tokio_test::{assert_pending, assert_ready, task}; + +use std::pin::Pin; +use std::task::{Context, Poll}; #[tokio::test] async fn size_hint_without_peek() { @@ -121,3 +125,93 @@ async fn size_hint_unbounded_upper() { let _ = s.peek().await; assert_eq!(s.size_hint(), (2, None)); // still unbounded after peek } + +#[tokio::test] +async fn peek_does_not_consume() { + let mut stream = stream::iter(vec![1, 2, 3]).peekable(); + + assert_eq!(stream.peek().await, Some(&1)); + assert_eq!(stream.peek().await, Some(&1)); + assert_eq!(stream.next().await, Some(1)); + assert_eq!(stream.next().await, Some(2)); +} + +#[tokio::test] +async fn peek_mut_mutates_yielded_item() { + let mut stream = stream::iter(vec![1, 2, 3]).peekable(); + + if let Some(item) = stream.peek_mut().await { + *item += 10; + } + + assert_eq!(stream.next().await, Some(11)); + assert_eq!(stream.next().await, Some(2)); +} + +#[tokio::test] +async fn peek_on_empty_stream() { + let mut stream = stream::iter(Vec::::new()).peekable(); + + assert_eq!(stream.peek().await, None); + assert_eq!(stream.peek_mut().await, None); + assert_eq!(stream.next().await, None); +} + +#[test] +fn poll_peek_does_not_advance_stream() { + let mut stream = task::spawn(stream::iter(vec![1, 2, 3]).peekable()); + + let first = stream.enter(|cx, s| s.poll_peek(cx).map(|opt| opt.copied())); + assert_eq!(assert_ready!(first), Some(1)); + + let second = stream.enter(|cx, s| s.poll_peek(cx).map(|opt| opt.copied())); + assert_eq!(assert_ready!(second), Some(1)); + + let next = stream.enter(|cx, s| s.poll_next(cx)); + assert_eq!(assert_ready!(next), Some(1)); + + let next = stream.enter(|cx, s| s.poll_next(cx)); + assert_eq!(assert_ready!(next), Some(2)); +} + +#[test] +fn poll_peek_mutates_buffered_item() { + let mut stream = task::spawn(stream::iter(vec![1, 2, 3]).peekable()); + + stream.enter(|cx, s| { + if let Poll::Ready(Some(item)) = s.poll_peek(cx) { + *item += 100; + } + }); + + let next = stream.enter(|cx, s| s.poll_next(cx)); + assert_eq!(assert_ready!(next), Some(101)); +} + +struct PendingOnce { + polled: bool, +} + +impl Stream for PendingOnce { + type Item = i32; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + if self.polled { + Poll::Ready(Some(7)) + } else { + self.polled = true; + cx.waker().wake_by_ref(); + Poll::Pending + } + } +} + +#[test] +fn poll_peek_propagates_pending() { + let mut stream = task::spawn(PendingOnce { polled: false }.peekable()); + + assert_pending!(stream.enter(|cx, s| s.poll_peek(cx).map(|opt| opt.copied()))); + + let ready = stream.enter(|cx, s| s.poll_peek(cx).map(|opt| opt.copied())); + assert_eq!(assert_ready!(ready), Some(7)); +}