1 | use core::pin::Pin; |
2 | use futures_core::future::{FusedFuture, Future, TryFuture}; |
3 | use futures_core::ready; |
4 | use futures_core::task::{Context, Poll}; |
5 | use pin_project_lite::pin_project; |
6 | |
7 | pin_project! { |
8 | #[project = TryFlattenErrProj] |
9 | #[derive(Debug)] |
10 | pub enum TryFlattenErr<Fut1, Fut2> { |
11 | First { #[pin] f: Fut1 }, |
12 | Second { #[pin] f: Fut2 }, |
13 | Empty, |
14 | } |
15 | } |
16 | |
17 | impl<Fut1, Fut2> TryFlattenErr<Fut1, Fut2> { |
18 | pub(crate) fn new(future: Fut1) -> Self { |
19 | Self::First { f: future } |
20 | } |
21 | } |
22 | |
23 | impl<Fut> FusedFuture for TryFlattenErr<Fut, Fut::Error> |
24 | where |
25 | Fut: TryFuture, |
26 | Fut::Error: TryFuture<Ok = Fut::Ok>, |
27 | { |
28 | fn is_terminated(&self) -> bool { |
29 | match self { |
30 | Self::Empty => true, |
31 | _ => false, |
32 | } |
33 | } |
34 | } |
35 | |
36 | impl<Fut> Future for TryFlattenErr<Fut, Fut::Error> |
37 | where |
38 | Fut: TryFuture, |
39 | Fut::Error: TryFuture<Ok = Fut::Ok>, |
40 | { |
41 | type Output = Result<Fut::Ok, <Fut::Error as TryFuture>::Error>; |
42 | |
43 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
44 | Poll::Ready(loop { |
45 | match self.as_mut().project() { |
46 | TryFlattenErrProj::First { f } => match ready!(f.try_poll(cx)) { |
47 | Err(f) => self.set(Self::Second { f }), |
48 | Ok(e) => { |
49 | self.set(Self::Empty); |
50 | break Ok(e); |
51 | } |
52 | }, |
53 | TryFlattenErrProj::Second { f } => { |
54 | let output = ready!(f.try_poll(cx)); |
55 | self.set(Self::Empty); |
56 | break output; |
57 | } |
58 | TryFlattenErrProj::Empty => panic!("TryFlattenErr polled after completion" ), |
59 | } |
60 | }) |
61 | } |
62 | } |
63 | |