pollable_map/futures/
optional.rs

1use futures::future::FusedFuture;
2use std::future::Future;
3use std::pin::Pin;
4use std::task::{Context, Poll, Waker};
5
6/// A reusable future that is the equivalent to an `Option`.
7///
8/// By default, this future will be empty, which would return  [`Poll::Pending`] when polled,
9/// but if a [`Future`] is supplied either upon construction via [`OptionalFuture::new`] or
10/// is set via [`OptionalFuture::replace`], the future would then be polled once [`OptionalFuture`]
11/// is polled. Once the future is polled to completion, the results will be returned, with
12/// [`OptionalFuture`] being empty.
13pub struct OptionalFuture<F> {
14    future: Option<F>,
15    waker: Option<Waker>,
16}
17
18impl<F: Unpin> Unpin for OptionalFuture<F> {}
19
20impl<F> Default for OptionalFuture<F> {
21    fn default() -> Self {
22        Self {
23            future: None,
24            waker: None,
25        }
26    }
27}
28
29impl<F> From<Option<F>> for OptionalFuture<F> {
30    fn from(fut: Option<F>) -> Self {
31        Self {
32            future: fut,
33            waker: None,
34        }
35    }
36}
37
38impl<F: Future> From<F> for OptionalFuture<F> {
39    fn from(fut: F) -> Self {
40        Self {
41            future: Some(fut),
42            waker: None,
43        }
44    }
45}
46
47impl<F> OptionalFuture<F> {
48    /// Construct a new `OptionalFuture` with an existing `Future`.
49    pub fn new(future: F) -> Self {
50        Self {
51            future: Some(future),
52            waker: None,
53        }
54    }
55
56    /// Takes the future out, leaving the OptionalFuture empty.
57    pub fn take(&mut self) -> Option<F> {
58        let fut = self.future.take();
59        if let Some(waker) = self.waker.take() {
60            waker.wake();
61        }
62        fut
63    }
64
65    /// Returns true if future still exist.
66    pub fn is_some(&self) -> bool {
67        self.future.is_some()
68    }
69
70    /// Returns false if future doesnt exist or has been completed.
71    pub fn is_none(&self) -> bool {
72        self.future.is_none()
73    }
74
75    /// Returns reference of the future.
76    pub fn as_ref(&self) -> Option<&F> {
77        self.future.as_ref()
78    }
79
80    /// Returns mutable reference of the future.
81    pub fn as_mut(&mut self) -> Option<&mut F> {
82        self.future.as_mut()
83    }
84
85    /// Replaces the current future with a new one, returning the previous future if present.
86    pub fn replace(&mut self, future: F) -> Option<F> {
87        let fut = self.future.replace(future);
88        if let Some(waker) = self.waker.take() {
89            waker.wake();
90        }
91        fut
92    }
93}
94
95impl<F> Future for OptionalFuture<F>
96where
97    F: Future + Send + Unpin + 'static,
98{
99    type Output = F::Output;
100
101    fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
102        let Some(future) = self.future.as_mut() else {
103            self.waker.replace(cx.waker().clone());
104            return Poll::Pending;
105        };
106
107        match Pin::new(future).poll(cx) {
108            Poll::Ready(output) => {
109                self.future.take();
110                Poll::Ready(output)
111            }
112            Poll::Pending => {
113                self.waker.replace(cx.waker().clone());
114                Poll::Pending
115            }
116        }
117    }
118}
119
120impl<F: Future> FusedFuture for OptionalFuture<F>
121where
122    F: Future + Send + Unpin + 'static,
123{
124    fn is_terminated(&self) -> bool {
125        self.future.is_none()
126    }
127}
128
129#[cfg(test)]
130mod test {
131    use super::*;
132
133    #[test]
134    fn test_optional_future() {
135        let mut future = OptionalFuture::new(futures::future::ready(0));
136        assert!(future.is_some());
137        let waker = futures::task::noop_waker_ref();
138
139        let val = Pin::new(&mut future).poll(&mut Context::from_waker(waker));
140        assert_eq!(val, Poll::Ready(0));
141        assert!(future.is_none());
142    }
143
144    #[test]
145    fn reusable_optional_future() {
146        let mut future = OptionalFuture::new(futures::future::ready(0));
147        assert!(future.is_some());
148        let waker = futures::task::noop_waker_ref();
149
150        let val = Pin::new(&mut future).poll(&mut Context::from_waker(waker));
151        assert_eq!(val, Poll::Ready(0));
152        assert!(future.is_none());
153
154        future.replace(futures::future::ready(1));
155        assert!(future.is_some());
156
157        let val = Pin::new(&mut future).poll(&mut Context::from_waker(waker));
158        assert_eq!(val, Poll::Ready(1));
159        assert!(future.is_none());
160    }
161
162    #[test]
163    fn convert_future_to_optional_future() {
164        let fut = futures::future::ready(0);
165
166        let mut future = OptionalFuture::from(fut);
167        assert!(future.is_some());
168        let waker = futures::task::noop_waker_ref();
169
170        let val = Pin::new(&mut future).poll(&mut Context::from_waker(waker));
171        assert_eq!(val, Poll::Ready(0));
172        assert!(future.is_none());
173    }
174}
pollable_map/futures/optional.rs

pollable_map/futures/
optional.rs
