tokio-stream: add `StreamExt::map_while` #4351

BraulioVM · 2021-12-27T23:41:41Z

Rust 1.57 stabilized the Iterator::map_while API. This PR adds the same functionality to the StreamExt trait, to keep parity.

Motivation

Keep parity with the standard iterator trait. There are many APIs that are available in std::iter::Iterator that aren't in StreamExt, but I guess it wouldn't harm to add this one?

Solution

The change is straight-forward, and is heavily inspired by the implementations of map and take_while.

Darksonn · 2021-12-28T09:22:30Z

tokio-stream/src/stream_ext/map_while.rs

+        #[pin]
+        stream: St,
+        f: F,
+        done: bool


This lets the Option optimization make the struct smaller.

Suggested change

#[pin]

stream: St,

f: F,

done: bool

#[pin]

stream: Option<St>,

f: F,

How do you know the option optimization applies here? Wouldn't that depend on the monomorphized St type? Is it common for the optimization to apply to Stream implementations?

Yes, it depends on the underlying Stream, but it includes any stream containing pretty much any kind of pointer, and that's very common.

Thanks!

A perhaps interesting aspect of this change is that the St object will be dropped as soon as the f returns None. I believe this is an acceptable aspect of this API, but should that be documented explicitly? The same does not apply to take_while, for example, where the St is not dropped until the TakeWhile is dropped, even if the TakeWhile predicate returns false.

I don't think it's necessary to document that it happens.

Darksonn · 2021-12-28T09:34:37Z

tokio-stream/src/stream_ext/map_while.rs

+        if !*self.as_mut().project().done {
+            self.as_mut().project().stream.poll_next(cx).map(|ready| {
+                let ready = ready.and_then(self.as_mut().project().f);
+
+                if ready.is_none() {
+                    *self.as_mut().project().done = true;
+                }
+
+                ready
+            })
+        } else {
+            Poll::Ready(None)
+        }


The code here is quite noisy with all the self.as_mut().project() repetitions. Consider something like this:

Suggested change

if !*self.as_mut().project().done {

self.as_mut().project().stream.poll_next(cx).map(|ready| {

let ready = ready.and_then(self.as_mut().project().f);

if ready.is_none() {

*self.as_mut().project().done = true;

}

ready

})

} else {

Poll::Ready(None)

}

let me = self.project();

match me.stream.as_pin_mut() {

Some(stream) => match stream.poll_next(cx).map(|opt| opt.and_then(me.f)) {

Poll::Ready(Some(item)) => Poll::Ready(Some(item)),

Poll::Ready(None) => {

me.stream.set(None);

Poll::Ready(None)

},

Poll::Pending => Poll::Pending,

},

None => Poll::Ready(None),

}

I agree with reducing the noise in the implemetnation, but I'm not sure this will work. In particular, me.stream.as_pin_mut will consume the Pin<&mut Option<St>> and then we won't be able to call set on it. I can't find any other way of accessing the Pin<&mut St> we need to call poll_next on the underlying structure.

On the other hand, if we leave St as it is but then make the F optional, (f: Option<F>), because pinning is not structural to f, we can empty out the Option<F> whenever we want. I am not sure if the Option optimization is less likely to apply to F than it is to the St though. For what it's worth, the skip_while combinator also makes the associated closure optional, to the St.

Would it make sense to make the F optional instead of the St?

I'm not super comfortable with pinning so maybe there's something obvious I'm missing.

You can do me.stream.as_mut().as_pin_mut().

!! Thanks, that's useful.

Darksonn

Thanks.

taiki-e · 2021-12-30T17:35:17Z

tokio-stream/src/stream_ext/map_while.rs

+    #[must_use = "streams do nothing unless polled"]
+    pub struct MapWhile<St, F> {
+        #[pin]
+        stream: Option<St>,


I don't think this Option is needed. Users who need this behavior can use .fuse() explicitly.

See also implementation of Iterator::map_while

That's a very good point. Will address

There are a few combinators (skip_while and take_while at least) where we keep an extra piece of state (like the Option here) to remember whether we were done. I guess we don't need those either! They can be fixed in a different PR though

Darksonn · 2021-12-30T17:54:58Z

Also, there's a merge conflict preventing this from getting merged.

Fixes tokio-rs#4337 Rust 1.57 stabilized the `Iterator::map_while` API. This PR adds the same functionality to the `StreamExt` trait, to keep parity.

This way we can take leverage the Option optimization where for non-null types we don't need any extra space to represent the `None`-ness. Whether this applies will depend on the specific `St` type, but it's very common.

`MapWhile` doesn't need to remember whether it's done. As per the contract of the `Stream` trait: > Once a stream has finished (returned Ready(None) from poll_next), calling its poll_next method again may panic, block forever, or cause other kinds of problems; the Stream trait places no requirements on the effects of such a call. This allows us to save the space that may be needed by the boolean, but also lets us simplify the implementation. People can use `fuse` if they so need.

taiki-e

LGTM aside from a nit.

tokio-stream/src/stream_ext.rs

Co-authored-by: Taiki Endo <te316e89@gmail.com>

taiki-e · 2021-12-31T09:58:16Z

(CI failure will be fixed in #4361)

Darksonn · 2021-12-31T10:31:40Z

You can merge in master to your branch to fix the CI failure.

Darksonn added the A-tokio-stream Area: The tokio-stream crate label Dec 28, 2021

Darksonn reviewed Dec 28, 2021

View reviewed changes

BraulioVM force-pushed the add/tokio-stream-take-while branch from b2d92f8 to f53b309 Compare December 28, 2021 15:49

Darksonn approved these changes Dec 28, 2021

View reviewed changes

taiki-e reviewed Dec 30, 2021

View reviewed changes

BraulioVM added 3 commits December 30, 2021 19:22

tokio-stream: add StreamExt::map_while`

5579e9c

Fixes tokio-rs#4337 Rust 1.57 stabilized the `Iterator::map_while` API. This PR adds the same functionality to the `StreamExt` trait, to keep parity.

bundle the boolean in an Option with the stream

cd834df

This way we can take leverage the Option optimization where for non-null types we don't need any extra space to represent the `None`-ness. Whether this applies will depend on the specific `St` type, but it's very common.

BraulioVM force-pushed the add/tokio-stream-take-while branch from f53b309 to cc2aac6 Compare December 30, 2021 19:37

taiki-e approved these changes Dec 31, 2021

View reviewed changes

tokio-stream/src/stream_ext.rs Outdated Show resolved Hide resolved

Update tokio-stream/src/stream_ext.rs

b6cb91f

Co-authored-by: Taiki Endo <te316e89@gmail.com>

Merge branch 'master' into add/tokio-stream-take-while

b001b6f

taiki-e approved these changes Dec 31, 2021

View reviewed changes

taiki-e merged commit fb35c83 into tokio-rs:master Dec 31, 2021

Darksonn mentioned this pull request Jun 4, 2022

chore: prepare tokio-stream 0.1.9 #4743

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

tokio-stream: add `StreamExt::map_while` #4351

tokio-stream: add `StreamExt::map_while` #4351

BraulioVM commented Dec 27, 2021

Darksonn Dec 28, 2021

BraulioVM Dec 28, 2021

Darksonn Dec 28, 2021

BraulioVM Dec 28, 2021

Darksonn Dec 28, 2021

Darksonn Dec 28, 2021

BraulioVM Dec 28, 2021

Darksonn Dec 28, 2021

BraulioVM Dec 28, 2021

Darksonn left a comment

taiki-e Dec 30, 2021

BraulioVM Dec 30, 2021

BraulioVM Dec 30, 2021

Darksonn commented Dec 30, 2021

taiki-e left a comment

taiki-e commented Dec 31, 2021

Darksonn commented Dec 31, 2021

tokio-stream: add StreamExt::map_while #4351

tokio-stream: add StreamExt::map_while #4351

Conversation

BraulioVM commented Dec 27, 2021

Motivation

Solution

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Darksonn left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Darksonn commented Dec 30, 2021

taiki-e left a comment

Choose a reason for hiding this comment

taiki-e commented Dec 31, 2021

Darksonn commented Dec 31, 2021

tokio-stream: add `StreamExt::map_while` #4351

tokio-stream: add `StreamExt::map_while` #4351