futures_util/future/
try_maybe_done.rs

1//! Definition of the TryMaybeDone combinator
2
3use super::assert_future;
4use core::mem;
5use core::pin::Pin;
6use futures_core::future::{FusedFuture, Future, TryFuture};
7use futures_core::ready;
8use 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)]
14pub 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
25impl<Fut: TryFuture + Unpin> Unpin for TryMaybeDone<Fut> {}
26
27/// Wraps a future into a `TryMaybeDone`
28pub fn try_maybe_done<Fut: TryFuture>(future: Fut) -> TryMaybeDone<Fut> {
29    assert_future::<Result<(), Fut::Error>, _>(TryMaybeDone::Future(future))
30}
31
32impl<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
64impl<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
73impl<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) => match ready!(Pin::new_unchecked(f).try_poll(cx)) {
80                    Ok(res) => self.set(Self::Done(res)),
81                    Err(e) => {
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}