1 | //! Definition of the TryMaybeDone combinator |
2 | |
3 | use super::assert_future; |
4 | use core::mem; |
5 | use core::pin::Pin; |
6 | use futures_core::future::{FusedFuture, Future, TryFuture}; |
7 | use futures_core::ready; |
8 | use futures_core::task::{Context, Poll}; |
9 | |
10 | /// A future that may have completed with an error. |
11 | /// |
12 | /// This is created by the [`try_maybe_done()`] function. |
13 | #[derive (Debug)] |
14 | pub enum TryMaybeDone<Fut: TryFuture> { |
15 | /// A not-yet-completed future |
16 | Future(/* #[pin] */ Fut), |
17 | /// The output of the completed future |
18 | Done(Fut::Ok), |
19 | /// The empty variant after the result of a [`TryMaybeDone`] has been |
20 | /// taken using the [`take_output`](TryMaybeDone::take_output) method, |
21 | /// or if the future returned an error. |
22 | Gone, |
23 | } |
24 | |
25 | impl<Fut: TryFuture + Unpin> Unpin for TryMaybeDone<Fut> {} |
26 | |
27 | /// Wraps a future into a `TryMaybeDone` |
28 | pub fn try_maybe_done<Fut: TryFuture>(future: Fut) -> TryMaybeDone<Fut> { |
29 | assert_future::<Result<(), Fut::Error>, _>(TryMaybeDone::Future(future)) |
30 | } |
31 | |
32 | impl<Fut: TryFuture> TryMaybeDone<Fut> { |
33 | /// Returns an [`Option`] containing a mutable reference to the output of the future. |
34 | /// The output of this method will be [`Some`] if and only if the inner |
35 | /// future has completed successfully and [`take_output`](TryMaybeDone::take_output) |
36 | /// has not yet been called. |
37 | #[inline ] |
38 | pub fn output_mut(self: Pin<&mut Self>) -> Option<&mut Fut::Ok> { |
39 | unsafe { |
40 | match self.get_unchecked_mut() { |
41 | TryMaybeDone::Done(res) => Some(res), |
42 | _ => None, |
43 | } |
44 | } |
45 | } |
46 | |
47 | /// Attempt to take the output of a `TryMaybeDone` without driving it |
48 | /// towards completion. |
49 | #[inline ] |
50 | pub fn take_output(self: Pin<&mut Self>) -> Option<Fut::Ok> { |
51 | match &*self { |
52 | Self::Done(_) => {} |
53 | Self::Future(_) | Self::Gone => return None, |
54 | } |
55 | unsafe { |
56 | match mem::replace(self.get_unchecked_mut(), Self::Gone) { |
57 | TryMaybeDone::Done(output) => Some(output), |
58 | _ => unreachable!(), |
59 | } |
60 | } |
61 | } |
62 | } |
63 | |
64 | impl<Fut: TryFuture> FusedFuture for TryMaybeDone<Fut> { |
65 | fn is_terminated(&self) -> bool { |
66 | match self { |
67 | Self::Future(_) => false, |
68 | Self::Done(_) | Self::Gone => true, |
69 | } |
70 | } |
71 | } |
72 | |
73 | impl<Fut: TryFuture> Future for TryMaybeDone<Fut> { |
74 | type Output = Result<(), Fut::Error>; |
75 | |
76 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
77 | unsafe { |
78 | match self.as_mut().get_unchecked_mut() { |
79 | TryMaybeDone::Future(f: &mut Fut) => match ready!(Pin::new_unchecked(f).try_poll(cx)) { |
80 | Ok(res: ::Ok) => self.set(Self::Done(res)), |
81 | Err(e: ::Error) => { |
82 | self.set(Self::Gone); |
83 | return Poll::Ready(Err(e)); |
84 | } |
85 | }, |
86 | TryMaybeDone::Done(_) => {} |
87 | TryMaybeDone::Gone => panic!("TryMaybeDone polled after value taken" ), |
88 | } |
89 | } |
90 | Poll::Ready(Ok(())) |
91 | } |
92 | } |
93 | |