Skip to main content

async_std/future/future/
flatten.rs

1use pin_project_lite::pin_project;
2use std::future::Future;
3use std::pin::Pin;
4
5use crate::future::IntoFuture;
6use crate::task::{ready, Context, Poll};
7
8pin_project! {
9    #[doc(hidden)]
10    #[allow(missing_debug_implementations)]
11    pub struct FlattenFuture<Fut1, Fut2> {
12        #[pin]
13        state: State<Fut1, Fut2>,
14    }
15}
16
17pin_project! {
18    #[project = StateProj]
19    #[derive(Debug)]
20    enum State<Fut1, Fut2> {
21        First {
22            #[pin]
23            fut1: Fut1,
24        },
25        Second {
26            #[pin]
27            fut2: Fut2,
28        },
29        Empty,
30    }
31}
32
33impl<Fut1, Fut2> FlattenFuture<Fut1, Fut2> {
34    pub(crate) fn new(fut1: Fut1) -> FlattenFuture<Fut1, Fut2> {
35        FlattenFuture {
36            state: State::First { fut1 },
37        }
38    }
39}
40
41impl<Fut1> Future for FlattenFuture<Fut1, <Fut1::Output as IntoFuture>::Future>
42where
43    Fut1: Future,
44    Fut1::Output: IntoFuture,
45{
46    type Output = <Fut1::Output as IntoFuture>::Output;
47
48    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
49        let mut state = self.project().state;
50        loop {
51            match state.as_mut().project() {
52                StateProj::First { fut1 } => {
53                    let fut2 = ready!(fut1.poll(cx)).into_future();
54                    state.set(State::Second { fut2 });
55                }
56                StateProj::Second { fut2 } => {
57                    let v = ready!(fut2.poll(cx));
58                    state.set(State::Empty);
59                    return Poll::Ready(v);
60                }
61                StateProj::Empty => panic!("polled a completed future"),
62            }
63        }
64    }
65}