stream: add inner stream accessors to adaptors (#8272)

tokio-stream/src/stream_ext/filter.rs

@@ -31,6 +31,32 @@ impl<St, F> Filter<St, F> {
     pub(super) fn new(stream: St, f: F) -> Self {
         Self { stream, f }
     }
+
+    /// Returns a reference to the inner stream.
+    pub fn get_ref(&self) -> &St {
+        &self.stream
+    }
+
+    /// Returns a mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_mut(&mut self) -> &mut St {
+        &mut self.stream
+    }
+
+    /// Returns a pinned mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> {
+        self.project().stream
+    }
+
+    /// Consumes this combinator and returns the inner stream.
+    ///
+    /// This may discard intermediate combinator state.
+    pub fn into_inner(self) -> St {
+        self.stream
+    }
 }
 
 impl<St, F> Stream for Filter<St, F>

tokio-stream/src/stream_ext/filter_map.rs

@@ -31,6 +31,32 @@ impl<St, F> FilterMap<St, F> {
     pub(super) fn new(stream: St, f: F) -> Self {
         Self { stream, f }
     }
+
+    /// Returns a reference to the inner stream.
+    pub fn get_ref(&self) -> &St {
+        &self.stream
+    }
+
+    /// Returns a mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_mut(&mut self) -> &mut St {
+        &mut self.stream
+    }
+
+    /// Returns a pinned mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> {
+        self.project().stream
+    }
+
+    /// Consumes this combinator and returns the inner stream.
+    ///
+    /// This may discard intermediate combinator state.
+    pub fn into_inner(self) -> St {
+        self.stream
+    }
 }
 
 impl<St, F, T> Stream for FilterMap<St, F>

tokio-stream/src/stream_ext/map.rs

@@ -29,6 +29,32 @@ impl<St, F> Map<St, F> {
     pub(super) fn new(stream: St, f: F) -> Self {
         Map { stream, f }
     }
+
+    /// Returns a reference to the inner stream.
+    pub fn get_ref(&self) -> &St {
+        &self.stream
+    }
+
+    /// Returns a mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_mut(&mut self) -> &mut St {
+        &mut self.stream
+    }
+
+    /// Returns a pinned mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> {
+        self.project().stream
+    }
+
+    /// Consumes this combinator and returns the inner stream.
+    ///
+    /// This may discard intermediate combinator state.
+    pub fn into_inner(self) -> St {
+        self.stream
+    }
 }
 
 impl<St, F, T> Stream for Map<St, F>

tokio-stream/src/stream_ext/map_while.rs

@@ -37,6 +37,32 @@ impl<St, F> MapWhile<St, F> {
             done: false,
         }
     }
+
+    /// Returns a reference to the inner stream.
+    pub fn get_ref(&self) -> &St {
+        &self.stream
+    }
+
+    /// Returns a mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_mut(&mut self) -> &mut St {
+        &mut self.stream
+    }
+
+    /// Returns a pinned mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> {
+        self.project().stream
+    }
+
+    /// Consumes this combinator and returns the inner stream.
+    ///
+    /// This may discard intermediate combinator state.
+    pub fn into_inner(self) -> St {
+        self.stream
+    }
 }
 
 impl<St, F, T> Stream for MapWhile<St, F>

tokio-stream/src/stream_ext/skip.rs

@@ -31,6 +31,32 @@ impl<St> Skip<St> {
     pub(super) fn new(stream: St, remaining: usize) -> Self {
         Self { stream, remaining }
     }
+
+    /// Returns a reference to the inner stream.
+    pub fn get_ref(&self) -> &St {
+        &self.stream
+    }
+
+    /// Returns a mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_mut(&mut self) -> &mut St {
+        &mut self.stream
+    }
+
+    /// Returns a pinned mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> {
+        self.project().stream
+    }
+
+    /// Consumes this combinator and returns the inner stream.
+    ///
+    /// This may discard intermediate combinator state.
+    pub fn into_inner(self) -> St {
+        self.stream
+    }
 }
 
 impl<St> Stream for Skip<St>

tokio-stream/src/stream_ext/skip_while.rs

@@ -34,6 +34,32 @@ impl<St, F> SkipWhile<St, F> {
             predicate: Some(predicate),
         }
     }
+
+    /// Returns a reference to the inner stream.
+    pub fn get_ref(&self) -> &St {
+        &self.stream
+    }
+
+    /// Returns a mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_mut(&mut self) -> &mut St {
+        &mut self.stream
+    }
+
+    /// Returns a pinned mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> {
+        self.project().stream
+    }
+
+    /// Consumes this combinator and returns the inner stream.
+    ///
+    /// This may discard intermediate combinator state.
+    pub fn into_inner(self) -> St {
+        self.stream
+    }
 }
 
 impl<St, F> Stream for SkipWhile<St, F>

tokio-stream/src/stream_ext/take.rs

@@ -32,6 +32,32 @@ impl<St> Take<St> {
     pub(super) fn new(stream: St, remaining: usize) -> Self {
         Self { stream, remaining }
     }
+
+    /// Returns a reference to the inner stream.
+    pub fn get_ref(&self) -> &St {
+        &self.stream
+    }
+
+    /// Returns a mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_mut(&mut self) -> &mut St {
+        &mut self.stream
+    }
+
+    /// Returns a pinned mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> {
+        self.project().stream
+    }
+
+    /// Consumes this combinator and returns the inner stream.
+    ///
+    /// This may discard intermediate combinator state.
+    pub fn into_inner(self) -> St {
+        self.stream
+    }
 }
 
 impl<St> Stream for Take<St>

tokio-stream/src/stream_ext/take_while.rs

@@ -37,6 +37,32 @@ impl<St, F> TakeWhile<St, F> {
             done: false,
         }
     }
+
+    /// Returns a reference to the inner stream.
+    pub fn get_ref(&self) -> &St {
+        &self.stream
+    }
+
+    /// Returns a mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_mut(&mut self) -> &mut St {
+        &mut self.stream
+    }
+
+    /// Returns a pinned mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> {
+        self.project().stream
+    }
+
+    /// Consumes this combinator and returns the inner stream.
+    ///
+    /// This may discard intermediate combinator state.
+    pub fn into_inner(self) -> St {
+        self.stream
+    }
 }
 
 impl<St, F> Stream for TakeWhile<St, F>

tokio-stream/src/stream_ext/then.rs

@@ -38,6 +38,32 @@ impl<St, Fut, F> Then<St, Fut, F> {
             f,
         }
     }
+
+    /// Returns a reference to the inner stream.
+    pub fn get_ref(&self) -> &St {
+        &self.stream
+    }
+
+    /// Returns a mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_mut(&mut self) -> &mut St {
+        &mut self.stream
+    }
+
+    /// Returns a pinned mutable reference to the inner stream.
+    ///
+    /// Mutating the inner stream may confuse this combinator.
+    pub fn get_pin_mut(self: Pin<&mut Self>) -> Pin<&mut St> {
+        self.project().stream
+    }
+
+    /// Consumes this combinator and returns the inner stream.
+    ///
+    /// This may discard intermediate combinator state.
+    pub fn into_inner(self) -> St {
+        self.stream
+    }
 }
 
 impl<St, F, Fut> Stream for Then<St, Fut, F>

tokio-stream/tests/stream_accessors.rs

@@ -0,0 +1,152 @@
+use std::pin::Pin;
+
+use tokio_stream::{self as stream, Iter, StreamExt};
+
+fn base() -> Iter<std::vec::IntoIter<i32>> {
+    stream::iter(vec![1, 2, 3])
+}
+
+#[tokio::test]
+async fn map_accessors() {
+    let mut map = base().map(|x| x * 2);
+
+    let _: &Iter<_> = map.get_ref();
+    let _: &mut Iter<_> = map.get_mut();
+    let _: Pin<&mut Iter<_>> = Pin::new(&mut map).get_pin_mut();
+
+    assert_eq!(map.next().await, Some(2));
+
+    // The recovered inner stream continues from where the combinator left it.
+    let mut inner = map.into_inner();
+    assert_eq!(inner.next().await, Some(2));
+    assert_eq!(inner.next().await, Some(3));
+    assert_eq!(inner.next().await, None);
+}
+
+#[tokio::test]
+async fn take_accessors() {
+    let mut take = base().take(2);
+
+    let _: &Iter<_> = take.get_ref();
+    let _: &mut Iter<_> = take.get_mut();
+    let _: Pin<&mut Iter<_>> = Pin::new(&mut take).get_pin_mut();
+
+    assert_eq!(take.next().await, Some(1));
+
+    // `take(2)` would stop after item 2, but the inner stream still has
+    // everything that was not yet pulled.
+    let mut inner = take.into_inner();
+    assert_eq!(inner.next().await, Some(2));
+    assert_eq!(inner.next().await, Some(3));
+}
+
+#[tokio::test]
+async fn skip_accessors() {
+    let mut skip = base().skip(1);
+
+    let _: &Iter<_> = skip.get_ref();
+    let _: &mut Iter<_> = skip.get_mut();
+    let _: Pin<&mut Iter<_>> = Pin::new(&mut skip).get_pin_mut();
+
+    assert_eq!(skip.next().await, Some(2));
+
+    let mut inner = skip.into_inner();
+    assert_eq!(inner.next().await, Some(3));
+}
+
+#[tokio::test]
+async fn filter_accessors() {
+    let mut filter = base().filter(|&x| x % 2 == 1);
+
+    let _: &Iter<_> = filter.get_ref();
+    let _: &mut Iter<_> = filter.get_mut();
+    let _: Pin<&mut Iter<_>> = Pin::new(&mut filter).get_pin_mut();
+
+    assert_eq!(filter.next().await, Some(1));
+
+    let mut inner = filter.into_inner();
+    assert_eq!(inner.next().await, Some(2));
+}
+
+#[tokio::test]
+async fn filter_map_accessors() {
+    let mut filter_map = base().filter_map(|x| (x % 2 == 1).then_some(x * 10));
+
+    let _: &Iter<_> = filter_map.get_ref();
+    let _: &mut Iter<_> = filter_map.get_mut();
+    let _: Pin<&mut Iter<_>> = Pin::new(&mut filter_map).get_pin_mut();
+
+    assert_eq!(filter_map.next().await, Some(10));
+
+    let mut inner = filter_map.into_inner();
+    assert_eq!(inner.next().await, Some(2));
+}
+
+#[tokio::test]
+async fn map_while_accessors() {
+    let mut map_while = base().map_while(|x| (x < 3).then_some(x + 100));
+
+    let _: &Iter<_> = map_while.get_ref();
+    let _: &mut Iter<_> = map_while.get_mut();
+    let _: Pin<&mut Iter<_>> = Pin::new(&mut map_while).get_pin_mut();
+
+    assert_eq!(map_while.next().await, Some(101));
+
+    let mut inner = map_while.into_inner();
+    assert_eq!(inner.next().await, Some(2));
+}
+
+#[tokio::test]
+async fn take_while_accessors() {
+    let mut take_while = base().take_while(|&x| x < 3);
+
+    let _: &Iter<_> = take_while.get_ref();
+    let _: &mut Iter<_> = take_while.get_mut();
+    let _: Pin<&mut Iter<_>> = Pin::new(&mut take_while).get_pin_mut();
+
+    assert_eq!(take_while.next().await, Some(1));
+
+    let mut inner = take_while.into_inner();
+    assert_eq!(inner.next().await, Some(2));
+}
+
+#[tokio::test]
+async fn skip_while_accessors() {
+    let mut skip_while = base().skip_while(|&x| x < 2);
+
+    let _: &Iter<_> = skip_while.get_ref();
+    let _: &mut Iter<_> = skip_while.get_mut();
+    let _: Pin<&mut Iter<_>> = Pin::new(&mut skip_while).get_pin_mut();
+
+    assert_eq!(skip_while.next().await, Some(2));
+
+    let mut inner = skip_while.into_inner();
+    assert_eq!(inner.next().await, Some(3));
+}
+
+#[tokio::test]
+async fn then_accessors() {
+    let mut then = base().then(|x| std::future::ready(x + 1));
+
+    let _: &Iter<_> = then.get_ref();
+    let _: &mut Iter<_> = then.get_mut();
+    let _: Pin<&mut Iter<_>> = Pin::new(&mut then).get_pin_mut();
+
+    assert_eq!(then.next().await, Some(2));
+
+    let mut inner = then.into_inner();
+    assert_eq!(inner.next().await, Some(2));
+}
+
+#[tokio::test]
+async fn get_mut_can_mutate_inner() {
+    // Mutating through `get_mut` is observed by the combinator.
+    let mut map = base().map(|x| x * 2);
+
+    // Skip one item directly on the inner stream.
+    let inner = map.get_mut();
+    assert_eq!(inner.next().await, Some(1));
+
+    // The combinator now sees the stream from item 2 onwards.
+    assert_eq!(map.next().await, Some(4));
+}