From d4569bb550ef8adc73d5612b2daeb7f3f2f7dbbe Mon Sep 17 00:00:00 2001 From: Rachit2323 <76684704+Rachit2323@users.noreply.github.com> Date: Sun, 9 Aug 2026 21:41:51 +0530 Subject: [PATCH] stream: add inner stream accessors to adaptors (#8272) --- tokio-stream/src/stream_ext/filter.rs | 26 ++++ tokio-stream/src/stream_ext/filter_map.rs | 26 ++++ tokio-stream/src/stream_ext/map.rs | 26 ++++ tokio-stream/src/stream_ext/map_while.rs | 26 ++++ tokio-stream/src/stream_ext/skip.rs | 26 ++++ tokio-stream/src/stream_ext/skip_while.rs | 26 ++++ tokio-stream/src/stream_ext/take.rs | 26 ++++ tokio-stream/src/stream_ext/take_while.rs | 26 ++++ tokio-stream/src/stream_ext/then.rs | 26 ++++ tokio-stream/tests/stream_accessors.rs | 152 ++++++++++++++++++++++ 10 files changed, 386 insertions(+) create mode 100644 tokio-stream/tests/stream_accessors.rs diff --git a/tokio-stream/src/stream_ext/filter.rs b/tokio-stream/src/stream_ext/filter.rs index f4a5eeef4..a946a10bb 100644 --- a/tokio-stream/src/stream_ext/filter.rs +++ b/tokio-stream/src/stream_ext/filter.rs @@ -31,6 +31,32 @@ impl Filter { 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 Stream for Filter diff --git a/tokio-stream/src/stream_ext/filter_map.rs b/tokio-stream/src/stream_ext/filter_map.rs index 01244860c..0b530aca9 100644 --- a/tokio-stream/src/stream_ext/filter_map.rs +++ b/tokio-stream/src/stream_ext/filter_map.rs @@ -31,6 +31,32 @@ impl FilterMap { 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 Stream for FilterMap diff --git a/tokio-stream/src/stream_ext/map.rs b/tokio-stream/src/stream_ext/map.rs index c95603af4..4a89f8de5 100644 --- a/tokio-stream/src/stream_ext/map.rs +++ b/tokio-stream/src/stream_ext/map.rs @@ -29,6 +29,32 @@ impl Map { 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 Stream for Map diff --git a/tokio-stream/src/stream_ext/map_while.rs b/tokio-stream/src/stream_ext/map_while.rs index b8906b66f..8835ac13f 100644 --- a/tokio-stream/src/stream_ext/map_while.rs +++ b/tokio-stream/src/stream_ext/map_while.rs @@ -37,6 +37,32 @@ impl MapWhile { 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 Stream for MapWhile diff --git a/tokio-stream/src/stream_ext/skip.rs b/tokio-stream/src/stream_ext/skip.rs index 0e052f14c..fd28cd9d4 100644 --- a/tokio-stream/src/stream_ext/skip.rs +++ b/tokio-stream/src/stream_ext/skip.rs @@ -31,6 +31,32 @@ impl Skip { 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 Stream for Skip diff --git a/tokio-stream/src/stream_ext/skip_while.rs b/tokio-stream/src/stream_ext/skip_while.rs index 03d2a0d2f..a18bbd103 100644 --- a/tokio-stream/src/stream_ext/skip_while.rs +++ b/tokio-stream/src/stream_ext/skip_while.rs @@ -34,6 +34,32 @@ impl SkipWhile { 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 Stream for SkipWhile diff --git a/tokio-stream/src/stream_ext/take.rs b/tokio-stream/src/stream_ext/take.rs index 3bd08ca0e..eab011a49 100644 --- a/tokio-stream/src/stream_ext/take.rs +++ b/tokio-stream/src/stream_ext/take.rs @@ -32,6 +32,32 @@ impl Take { 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 Stream for Take diff --git a/tokio-stream/src/stream_ext/take_while.rs b/tokio-stream/src/stream_ext/take_while.rs index c228a292c..d10fa1a53 100644 --- a/tokio-stream/src/stream_ext/take_while.rs +++ b/tokio-stream/src/stream_ext/take_while.rs @@ -37,6 +37,32 @@ impl TakeWhile { 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 Stream for TakeWhile diff --git a/tokio-stream/src/stream_ext/then.rs b/tokio-stream/src/stream_ext/then.rs index c4314e6d0..466dc64eb 100644 --- a/tokio-stream/src/stream_ext/then.rs +++ b/tokio-stream/src/stream_ext/then.rs @@ -38,6 +38,32 @@ impl Then { 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 Stream for Then diff --git a/tokio-stream/tests/stream_accessors.rs b/tokio-stream/tests/stream_accessors.rs new file mode 100644 index 000000000..7b844a371 --- /dev/null +++ b/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> { + 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)); +}