mirror
/
async-std
mirror of https://github.com/async-rs/async-std.git
2
0
Fork 1
Code Issues Releases Wiki Activity

fix: Split FlattenCompat logic to Flatten and FlatMap

Browse Source
pull/367/head
k-nasa 5 years ago
parent a42ae2f3d9
commit 688976203e
3 changed files with 72 additions and 106 deletions
Show Stats Download Patch File Download Diff File

62
src/stream/stream/flat_map.rs
Unescape Escape View File

@ -0,0 +1,62 @@
use pin_project_lite::pin_project;
use std::pin::Pin;
use crate::prelude::*;
use crate::stream::stream::map::Map;
use crate::stream::{IntoStream, Stream};
use crate::task::{Context, Poll};
pin_project! {
/// This `struct` is created by the [`flat_map`] method on [`Stream`]. See its
/// documentation for more.
///
/// [`flat_map`]: trait.Stream.html#method.flat_map
/// [`Stream`]: trait.Stream.html
#[allow(missing_debug_implementations)]
pub struct FlatMap<S, U, T, F> {
#[pin]
stream: Map<S, F, T, U>,
#[pin]
inner_stream: Option<U>,
}
}
impl<S, U, F> FlatMap<S, U, S::Item, F>
where
S: Stream,
U: IntoStream,
F: FnMut(S::Item) -> U,
{
pub(super) fn new(stream: S, f: F) -> FlatMap<S, U, S::Item, F> {
FlatMap {
stream: stream.map(f),
inner_stream: None,
}
}
}
impl<S, U, F> Stream for FlatMap<S, U, S::Item, F>
where
S: Stream,
S::Item: IntoStream<IntoStream = U, Item = U::Item>,
U: Stream,
F: FnMut(S::Item) -> U,
{
type Item = U::Item;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
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);
}
}
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())),
}
}
}
}

112
src/stream/stream/flatten.rs
Unescape Escape View File

@ -6,46 +6,6 @@ use crate::stream::stream::map::Map;
use crate::stream::{IntoStream, Stream};
use crate::task::{Context, Poll};
pin_project! {
/// This `struct` is created by the [`flat_map`] method on [`Stream`]. See its
/// documentation for more.
///
/// [`flat_map`]: trait.Stream.html#method.flat_map
/// [`Stream`]: trait.Stream.html
#[allow(missing_debug_implementations)]
pub struct FlatMap<S, U, T, F> {
#[pin]
inner: FlattenCompat<Map<S, F, T, U>, U>,
}
}
impl<S, U, F> FlatMap<S, U, S::Item, F>
where
S: Stream,
U: IntoStream,
F: FnMut(S::Item) -> U,
{
pub(super) fn new(stream: S, f: F) -> FlatMap<S, U, S::Item, F> {
FlatMap {
inner: FlattenCompat::new(stream.map(f)),
}
}
}
impl<S, U, F> Stream for FlatMap<S, U, S::Item, F>
where
S: Stream,
S::Item: IntoStream<IntoStream = U, Item = U::Item>,
U: Stream,
F: FnMut(S::Item) -> U,
{
type Item = U::Item;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.project().inner.poll_next(cx)
}
}
pin_project! {
/// This `struct` is created by the [`flatten`] method on [`Stream`]. See its
/// documentation for more.
@ -55,7 +15,9 @@ pin_project! {
#[allow(missing_debug_implementations)]
pub struct Flatten<S, U> {
#[pin]
inner: FlattenCompat<S, U>
stream: S,
#[pin]
inner_stream: Option<U>,
}
}
@ -66,47 +28,13 @@ where
{
pub(super) fn new(stream: S) -> Flatten<S, S::Item> {
Flatten {
inner: FlattenCompat::new(stream),
}
}
}
impl<S, U> Stream for Flatten<S, <S::Item as IntoStream>::IntoStream>
where
S: Stream,
S::Item: IntoStream<IntoStream = U, Item = U::Item>,
U: Stream,
{
type Item = U::Item;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
self.project().inner.poll_next(cx)
}
}
pin_project! {
/// Real logic of both `Flatten` and `FlatMap` which simply delegate to
/// this type.
#[derive(Clone, Debug)]
struct FlattenCompat<S, U> {
#[pin]
stream: S,
#[pin]
frontiter: Option<U>,
}
}
impl<S, U> FlattenCompat<S, U> {
/// Adapts an iterator by flattening it, for use in `flatten()` and `flat_map()`.
fn new(stream: S) -> FlattenCompat<S, U> {
FlattenCompat {
stream,
frontiter: None,
inner_stream: None,
}
}
}
impl<S, U> Stream for FlattenCompat<S, U>
impl<S, U> Stream for Flatten<S, <S::Item as IntoStream>::IntoStream>
where
S: Stream,
S::Item: IntoStream<IntoStream = U, Item = U::Item>,
@ -117,7 +45,7 @@ where
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut this = self.project();
loop {
if let Some(inner) = this.frontiter.as_mut().as_pin_mut() {
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);
}
@ -125,34 +53,8 @@ where
match futures_core::ready!(this.stream.as_mut().poll_next(cx)) {
None => return Poll::Ready(None),
Some(inner) => this.frontiter.set(Some(inner.into_stream())),
Some(inner) => this.inner_stream.set(Some(inner.into_stream())),
}
}
}
}
#[cfg(test)]
mod tests {
use super::FlattenCompat;
use crate::prelude::*;
use crate::task;
use std::collections::VecDeque;
#[test]
fn test_poll_next() -> std::io::Result<()> {
let inner1: VecDeque<u8> = vec![1, 2, 3].into_iter().collect();
let inner2: VecDeque<u8> = vec![4, 5, 6].into_iter().collect();
let s: VecDeque<_> = vec![inner1, inner2].into_iter().collect();
task::block_on(async move {
let flat = FlattenCompat::new(s);
let v: Vec<u8> = flat.collect().await;
assert_eq!(v, vec![1, 2, 3, 4, 5, 6]);
Ok(())
})
}
}

4
src/stream/stream/mod.rs
Unescape Escape View File

@ -99,10 +99,12 @@ cfg_unstable! {
use crate::stream::into_stream::IntoStream;
pub use merge::Merge;
pub use flatten::{FlatMap, Flatten};
pub use flatten::Flatten;
pub use flat_map::FlatMap;
mod merge;
mod flatten;
mod flat_map;
}
extension_trait! {

Write Preview
Loading…
Cancel
Save