Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

inline poll_states and remove Fuse for vec::merge #79

Merged
merged 1 commit into from
Nov 16, 2022
Merged

inline poll_states and remove Fuse for vec::merge #79

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 21 additions & 14 deletions src/stream/merge/vec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
use super::Merge as MergeTrait;
use crate::stream::IntoStream;
use crate::utils::{self, Fuse, RandomGenerator, WakerList};
use crate::utils::{self, PollState, PollStates, RandomGenerator, WakerList};

use core::fmt;
use futures_core::Stream;
Expand All @@ -20,22 +20,29 @@ where
S: Stream,
{
#[pin]
streams: Vec<Fuse<S>>,
streams: Vec<S>,
yoshuawuyts marked this conversation as resolved.
Show resolved Hide resolved
rng: RandomGenerator,
complete: usize,
wakers: WakerList,
state: PollStates,
done: bool,
len: usize,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I assume this is the number of streams we're merging? I might suggest count or num_streams as alternate names, but there's plenty of precedent for using len for this sort of thing, so I'm happy to leave it up to your judgement.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually, I'd suggest adding a len() method that returns self.streams.len() and that way we save 8 bytes off the size of the Merge struct.

Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah, I wish we could set methods! - Because we're dealing with a pin projection it generates a struct for us containing the projection. Andy any self methods would need to point to that.

Because we don't control the projection we can't define methods. And so we basically have to define everything in-line. Which is like, another reason why I'd reallllly like us to have safe pin projections, so we can split manually authored futures into separate internal methods etc. Not being able to do this is another reason authoring futures by hand is hard :/

}

impl<S> Merge<S>
where
S: Stream,
{
pub(crate) fn new(streams: Vec<S>) -> Self {
let len = streams.len();
Self {
wakers: WakerList::new(streams.len()),
streams: streams.into_iter().map(Fuse::new).collect(),
wakers: WakerList::new(len),
state: PollStates::new(len),
streams,
rng: RandomGenerator::new(),
complete: 0,
done: false,
len,
}
}
}
Expand All @@ -58,22 +65,19 @@ where
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut this = self.project();

let mut readiness = this.wakers.readiness().lock().unwrap();
readiness.set_waker(cx.waker());

// Iterate over our streams one-by-one. If a stream yields a value,
// we exit early. By default we'll return `Poll::Ready(None)`, but
// this changes if we encounter a `Poll::Pending`.
let mut index = this.rng.generate(this.streams.len() as u32) as usize;

let mut readiness = this.wakers.readiness().lock().unwrap();
readiness.set_waker(cx.waker());
loop {
let len = *this.len;
let r = this.rng.generate(this.streams.len() as u32) as usize;
for index in (0..len).map(|n| (r + n).wrapping_rem(len)) {
if !readiness.any_ready() {
// Nothing is ready yet
return Poll::Pending;
}

index = (index + 1).wrapping_rem(this.streams.len());

if !readiness.clear_ready(index) {
} else if !readiness.clear_ready(index) || this.state[index].is_consumed() {
continue;
}

Expand All @@ -92,6 +96,7 @@ where
}
Poll::Ready(None) => {
*this.complete += 1;
this.state[index] = PollState::Consumed;
yoshuawuyts marked this conversation as resolved.
Show resolved Hide resolved
if *this.complete == this.streams.len() {
return Poll::Ready(None);
}
Expand All @@ -102,6 +107,8 @@ where
// Lock readiness so we can use it again
readiness = this.wakers.readiness().lock().unwrap();
}

Poll::Pending
}
}

Expand Down