|
| 1 | +//! A type that wraps a future to keep track of its completion status. |
| 2 | +//! |
| 3 | +//! This implementation was taken from the original `macro_rules` `join/try_join` |
| 4 | +//! macros in the `futures-preview` crate. |
| 5 | +
|
| 6 | +use std::future::Future; |
| 7 | +use std::mem; |
| 8 | +use std::pin::Pin; |
| 9 | +use std::task::{Context, Poll}; |
| 10 | + |
| 11 | +use futures_core::ready; |
| 12 | + |
| 13 | +/// A future that may have completed. |
| 14 | +#[derive(Debug)] |
| 15 | +pub(crate) enum MaybeDone<Fut: Future> { |
| 16 | + /// A not-yet-completed future |
| 17 | + Future(Fut), |
| 18 | + |
| 19 | + /// The output of the completed future |
| 20 | + Done(Fut::Output), |
| 21 | + |
| 22 | + /// The empty variant after the result of a [`MaybeDone`] has been |
| 23 | + /// taken using the [`take`](MaybeDone::take) method. |
| 24 | + Gone, |
| 25 | +} |
| 26 | + |
| 27 | +impl<Fut: Future> MaybeDone<Fut> { |
| 28 | + /// Create a new instance of `MaybeDone`. |
| 29 | + pub(crate) fn new(future: Fut) -> MaybeDone<Fut> { |
| 30 | + Self::Future(future) |
| 31 | + } |
| 32 | + |
| 33 | + /// Returns an [`Option`] containing a 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 been completed and [`take`](MaybeDone::take) |
| 36 | + /// has not yet been called. |
| 37 | + #[inline] |
| 38 | + pub(crate) fn output(self: Pin<&Self>) -> Option<&Fut::Output> { |
| 39 | + let this = self.get_ref(); |
| 40 | + match this { |
| 41 | + MaybeDone::Done(res) => Some(res), |
| 42 | + _ => None, |
| 43 | + } |
| 44 | + } |
| 45 | + |
| 46 | + /// Attempt to take the output of a `MaybeDone` without driving it |
| 47 | + /// towards completion. |
| 48 | + #[inline] |
| 49 | + pub(crate) fn take(self: Pin<&mut Self>) -> Option<Fut::Output> { |
| 50 | + unsafe { |
| 51 | + let this = self.get_unchecked_mut(); |
| 52 | + match this { |
| 53 | + MaybeDone::Done(_) => {} |
| 54 | + MaybeDone::Future(_) | MaybeDone::Gone => return None, |
| 55 | + }; |
| 56 | + if let MaybeDone::Done(output) = mem::replace(this, MaybeDone::Gone) { |
| 57 | + Some(output) |
| 58 | + } else { |
| 59 | + unreachable!() |
| 60 | + } |
| 61 | + } |
| 62 | + } |
| 63 | +} |
| 64 | + |
| 65 | +impl<Fut: Future> Future for MaybeDone<Fut> { |
| 66 | + type Output = (); |
| 67 | + |
| 68 | + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
| 69 | + let res = unsafe { |
| 70 | + match Pin::as_mut(&mut self).get_unchecked_mut() { |
| 71 | + MaybeDone::Future(a) => ready!(Pin::new_unchecked(a).poll(cx)), |
| 72 | + MaybeDone::Done(_) => return Poll::Ready(()), |
| 73 | + MaybeDone::Gone => panic!("MaybeDone polled after value taken"), |
| 74 | + } |
| 75 | + }; |
| 76 | + self.set(MaybeDone::Done(res)); |
| 77 | + Poll::Ready(()) |
| 78 | + } |
| 79 | +} |
0 commit comments