futures_concurrency/utils/poll_state/
maybe_done.rs

1use core::future::Future;
2use core::mem;
3use core::pin::Pin;
4use core::task::{ready, Context, Poll};
5
6/// A future that may have completed.
7#[derive(Debug)]
8pub(crate) enum MaybeDone<Fut: Future> {
9    /// A not-yet-completed future
10    Future(Fut),
11
12    /// The output of the completed future
13    Done(Fut::Output),
14
15    /// The empty variant after the result of a [`MaybeDone`] has been
16    /// taken using the [`take`](MaybeDone::take) method.
17    Gone,
18}
19
20impl<Fut: Future> MaybeDone<Fut> {
21    /// Create a new instance of `MaybeDone`.
22    pub(crate) fn new(future: Fut) -> MaybeDone<Fut> {
23        Self::Future(future)
24    }
25}
26
27impl<T, E, Fut> MaybeDone<Fut>
28where
29    Fut: Future<Output = Result<T, E>>,
30{
31    /// Attempt to take the `Ok(output)` of a `MaybeDone` without driving it towards completion.
32    /// If the future is done but is an `Err(_)`, this will return `None`.
33    #[inline]
34    pub(crate) fn take_ok(self: Pin<&mut Self>) -> Option<T> {
35        let this = unsafe { self.get_unchecked_mut() };
36        match this {
37            MaybeDone::Done(Ok(_)) => {}
38            MaybeDone::Done(Err(_)) | MaybeDone::Future(_) | MaybeDone::Gone => return None,
39        }
40        match mem::replace(this, MaybeDone::Gone) {
41            MaybeDone::Done(Ok(output)) => Some(output),
42            _ => unreachable!(),
43        }
44    }
45
46    /// Attempt to take the `Err(output)` of a `MaybeDone` without driving it towards completion.
47    /// If the future is done but is an `Ok(_)`, this will return `None`.
48    #[inline]
49    pub(crate) fn take_err(self: Pin<&mut Self>) -> Option<E> {
50        let this = unsafe { self.get_unchecked_mut() };
51        match this {
52            MaybeDone::Done(Err(_)) => {}
53            MaybeDone::Done(Ok(_)) | MaybeDone::Future(_) | MaybeDone::Gone => return None,
54        }
55        match mem::replace(this, MaybeDone::Gone) {
56            MaybeDone::Done(Err(output)) => Some(output),
57            _ => unreachable!(),
58        }
59    }
60}
61
62impl<Fut: Future> Future for MaybeDone<Fut> {
63    type Output = ();
64
65    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
66        let res = unsafe {
67            match Pin::as_mut(&mut self).get_unchecked_mut() {
68                MaybeDone::Future(a) => ready!(Pin::new_unchecked(a).poll(cx)),
69                MaybeDone::Done(_) => return Poll::Ready(()),
70                MaybeDone::Gone => panic!("MaybeDone polled after value taken"),
71            }
72        };
73        self.set(MaybeDone::Done(res));
74        Poll::Ready(())
75    }
76}
futures_concurrency/utils/poll_state/maybe_done.rs

futures_concurrency/utils/poll_state/
maybe_done.rs
