futures_util/future/
try_maybe_done.rs1use 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#[derive(Debug)]
14pub enum TryMaybeDone<Fut: TryFuture> {
15 Future(Fut),
17 Done(Fut::Ok),
19 Gone,
23}
24
25impl<Fut: TryFuture + Unpin> Unpin for TryMaybeDone<Fut> {}
26
27pub 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 #[inline]
38 pub fn output_mut(self: Pin<&mut Self>) -> Option<&mut Fut::Ok> {
39 unsafe {
40 match self.get_unchecked_mut() {
41 Self::Done(res) => Some(res),
42 _ => None,
43 }
44 }
45 }
46
47 #[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 Self::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 Self::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 Self::Done(_) => {}
87 Self::Gone => panic!("TryMaybeDone polled after value taken"),
88 }
89 }
90 Poll::Ready(Ok(()))
91 }
92}