eater/async-std
1
0
Fork 0
forked from mirror/async-std
Code Issues Pull requests Releases Wiki Activity

Fixed flatten

Browse source
This commit is contained in:
Oleg Nosov 2020-02-08 15:41:33 +03:00
parent 85c32ef9d2
commit 32068942a6
No known key found for this signature in database
GPG key ID: DE90B83800644E24
3 changed files with 75 additions and 58 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
src/stream/stream/flatten.rs
View file

@ -1,5 +1,5 @@
use std::fmt;
use std::pin::Pin;
use core::fmt;
use core::pin::Pin;
use pin_project_lite::pin_project;
@ -52,14 +52,21 @@ where
let mut this = self.project();
loop {
if let Some(inner) = this.inner_stream.as_mut().as_pin_mut() {
if let item @ Some(_) = futures_core::ready!(inner.poll_next(cx)) {
return Poll::Ready(item);
let next_item = futures_core::ready!(inner.poll_next(cx));
if next_item.is_some() {
return Poll::Ready(next_item);
} else {
this.inner_stream.set(None);
}
}
match futures_core::ready!(this.stream.as_mut().poll_next(cx)) {
None => return Poll::Ready(None),
Some(inner) => this.inner_stream.set(Some(inner.into_stream())),
let inner = futures_core::ready!(this.stream.as_mut().poll_next(cx));
if inner.is_some() {
this.inner_stream.set(inner.map(IntoStream::into_stream));
} else {
return Poll::Ready(None);
}
}
}

110
tests/stream.rs
View file

@ -1,3 +1,5 @@
use std::convert::identity;
use std::marker::Unpin;
use std::pin::Pin;
use std::task::{Context, Poll};
@ -99,58 +101,52 @@ fn merge_works_with_unfused_streams() {
assert_eq!(xs, vec![92, 92]);
}
struct S<T>(T);
impl<T: Stream + Unpin> Stream for S<T> {
type Item = T::Item;
fn poll_next(mut self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Option<Self::Item>> {
unsafe { Pin::new_unchecked(&mut self.0) }.poll_next(ctx)
}
}
struct StrictOnce {
polled: bool,
}
impl Stream for StrictOnce {
type Item = ();
fn poll_next(mut self: Pin<&mut Self>, _: &mut Context) -> Poll<Option<Self::Item>> {
assert!(!self.polled, "Polled after completion!");
self.polled = true;
Poll::Ready(None)
}
}
struct Interchanger {
polled: bool,
}
impl Stream for Interchanger {
type Item = S<Box<dyn Stream<Item = ()> + Unpin>>;
fn poll_next(mut self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Option<Self::Item>> {
if self.polled {
self.polled = false;
ctx.waker().wake_by_ref();
Poll::Pending
} else {
self.polled = true;
Poll::Ready(Some(S(Box::new(StrictOnce { polled: false }))))
}
}
}
#[test]
fn flat_map_doesnt_poll_completed_inner_stream() {
async_std::task::block_on(async {
use async_std::prelude::*;
use async_std::task::*;
use std::convert::identity;
use std::marker::Unpin;
use std::pin::Pin;
struct S<T>(T);
impl<T: Stream + Unpin> Stream for S<T> {
type Item = T::Item;
fn poll_next(mut self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Option<Self::Item>> {
unsafe { Pin::new_unchecked(&mut self.0) }.poll_next(ctx)
}
}
struct StrictOnce {
polled: bool,
};
impl Stream for StrictOnce {
type Item = ();
fn poll_next(mut self: Pin<&mut Self>, _: &mut Context) -> Poll<Option<Self::Item>> {
assert!(!self.polled, "Polled after completion!");
self.polled = true;
Poll::Ready(None)
}
}
struct Interchanger {
polled: bool,
};
impl Stream for Interchanger {
type Item = S<Box<dyn Stream<Item = ()> + Unpin>>;
fn poll_next(mut self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Option<Self::Item>> {
if self.polled {
self.polled = false;
ctx.waker().wake_by_ref();
Poll::Pending
} else {
self.polled = true;
Poll::Ready(Some(S(Box::new(StrictOnce { polled: false }))))
}
}
}
task::block_on(async {
assert_eq!(
Interchanger { polled: false }
.take(2)
@ -161,3 +157,17 @@ fn flat_map_doesnt_poll_completed_inner_stream() {
);
});
}
#[test]
fn flatten_doesnt_poll_completed_inner_stream() {
task::block_on(async {
assert_eq!(
Interchanger { polled: false }
.take(2)
.flatten()
.count()
.await,
0
);
});
}