forked from mirror/async-std
fix: Split FlattenCompat logic to Flatten and FlatMap
This commit is contained in:
parent
a42ae2f3d9
commit
688976203e
3 changed files with 72 additions and 106 deletions
src/stream/stream
62
src/stream/stream/flat_map.rs
Normal file
62
src/stream/stream/flat_map.rs
Normal 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())),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
|
@ -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),
|
||||
stream,
|
||||
inner_stream: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<S, U> Stream for FlattenCompat<S, U>
|
||||
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(())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
@ -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! {
|
||||
|
|
Loading…
Reference in a new issue