forked from mirror/async-std
commit
9a8805678e
@ -0,0 +1,52 @@
|
||||
use futures_core::ready;
|
||||
use std::pin::Pin;
|
||||
|
||||
use crate::future::Future;
|
||||
use crate::future::IntoFuture;
|
||||
use crate::task::{Context, Poll};
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct FlattenFuture<Fut1, Fut2> {
|
||||
state: State<Fut1, Fut2>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum State<Fut1, Fut2> {
|
||||
First(Fut1),
|
||||
Second(Fut2),
|
||||
Empty,
|
||||
}
|
||||
|
||||
impl<Fut1, Fut2> FlattenFuture<Fut1, Fut2> {
|
||||
pub fn new(future: Fut1) -> FlattenFuture<Fut1, Fut2> {
|
||||
FlattenFuture {
|
||||
state: State::First(future),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl<Fut1> Future for FlattenFuture<Fut1, <Fut1::Output as IntoFuture>::Future>
|
||||
where
|
||||
Fut1: Future,
|
||||
Fut1::Output: IntoFuture,
|
||||
{
|
||||
type Output = <Fut1::Output as IntoFuture>::Output;
|
||||
|
||||
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
||||
let Self { state } = unsafe { self.get_unchecked_mut() };
|
||||
loop {
|
||||
match state {
|
||||
State::First(fut1) => {
|
||||
let fut2 = ready!(unsafe { Pin::new_unchecked(fut1) }.poll(cx)).into_future();
|
||||
*state = State::Second(fut2);
|
||||
}
|
||||
State::Second(fut2) => {
|
||||
let v = ready!(unsafe { Pin::new_unchecked(fut2) }.poll(cx));
|
||||
*state = State::Empty;
|
||||
return Poll::Ready(v);
|
||||
}
|
||||
State::Empty => panic!("polled a completed future"),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
Loading…
Reference in New Issue