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`
try_maybe_done<Fut: TryFuture>(future: Fut) -> TryMaybeDone<Fut>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]
output_mut(self: Pin<&mut Self>) -> Option<&mut Fut::Ok>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]
take_output(self: Pin<&mut Self>) -> Option<Fut::Ok>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> {
is_terminated(&self) -> bool65     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 
poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output>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) => {
80                     match ready!(Pin::new_unchecked(f).try_poll(cx)) {
81                         Ok(res) => self.set(Self::Done(res)),
82                         Err(e) => {
83                             self.set(Self::Gone);
84                             return Poll::Ready(Err(e));
85                         }
86                     }
87                 },
88                 TryMaybeDone::Done(_) => {},
89                 TryMaybeDone::Gone => panic!("TryMaybeDone polled after value taken"),
90             }
91         }
92         Poll::Ready(Ok(()))
93     }
94 }
95