async_hofs/
option.rs

1use core::future::Future;
2use core::pin::Pin;
3use core::task::{Context, Poll};
4use pin_project::pin_project;
5
6pub trait AsyncMapExt<T> {
7    /// Basically same as [`Option::map`], but it accepts closure that returns [`Future`]
8    ///
9    /// [`Option::map`]: core::option::Option::map
10    /// [`Future`]: core::future::Future
11    ///
12    /// # Examples
13    ///
14    /// ```
15    /// # #[tokio::main]
16    /// # async fn main() {
17    /// use async_hofs::prelude::*;
18    ///
19    /// assert_eq!(
20    ///     Some(1).async_map(|x: i32| async move { x + 1 }).await,
21    ///     Some(2),
22    /// );
23    ///
24    /// assert_eq!(
25    ///     None.async_map(|x: i32| async move { x + 1 }).await,
26    ///     None
27    /// );
28    /// # }
29    /// ```
30    fn async_map<TFn, TFuture>(self, f: TFn) -> AsyncMap<T, TFn, TFuture>
31    where
32        TFn: FnOnce(T) -> TFuture,
33        TFuture: Future;
34
35    /// Basically same as [`Option::and_then`], but it accepts closure that returns [`Future`]
36    ///
37    /// [`Option::and_then`]: core::option::Option::and_then
38    /// [`Future`]: core::future::Future
39    ///
40    /// # Examples
41    ///
42    /// ```
43    /// # #[tokio::main]
44    /// # async fn main() {
45    /// use async_hofs::prelude::*;
46    ///
47    /// assert_eq!(
48    ///     Some(1)
49    ///         .async_and_then(|x: i32| async move { Some(x + 1) })
50    ///         .await,
51    ///     Some(2),
52    /// );
53    ///
54    /// assert_eq!(
55    ///     Some(1)
56    ///         .async_and_then(|x: i32| async move { Option::<i32>::None })
57    ///         .await,
58    ///     None
59    /// );
60    ///
61    /// assert_eq!(
62    ///     None.async_and_then(|x: i32| async move { Some(x + 1) })
63    ///         .await,
64    ///     None
65    /// );
66    /// # }
67    /// ```
68    fn async_and_then<U, TFn, TFuture>(self, f: TFn) -> AsyncAndThen<T, TFn, TFuture>
69    where
70        TFn: FnOnce(T) -> TFuture,
71        TFuture: Future<Output = Option<U>>;
72}
73
74#[doc(hidden)]
75#[pin_project(project = AsyncMapProj)]
76pub enum AsyncMap<T, TFn, TFuture> {
77    None,
78    Pending(Option<(T, TFn)>),
79    Polling(#[pin] TFuture),
80}
81
82impl<T, U, TFn, TFuture> Future for AsyncMap<T, TFn, TFuture>
83where
84    TFn: FnOnce(T) -> TFuture,
85    TFuture: Future<Output = U>,
86{
87    type Output = Option<U>;
88
89    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
90        use AsyncMapProj::*;
91
92        match self.as_mut().project() {
93            None => Poll::Ready(Option::None),
94
95            Pending(payload) => {
96                let (x, f) = payload.take().expect("AsyncMap::Pending polled twice");
97                let future = f(x);
98                self.set(AsyncMap::Polling(future));
99                self.poll(cx)
100            }
101
102            Polling(future) => future.poll(cx).map(Some),
103        }
104    }
105}
106
107#[doc(hidden)]
108#[pin_project(project = AsyncAndThenProj)]
109pub enum AsyncAndThen<T, TFn, TFuture> {
110    None,
111    Pending(Option<(T, TFn)>),
112    Polling(#[pin] TFuture),
113}
114
115impl<T, U, TFn, TFuture> Future for AsyncAndThen<T, TFn, TFuture>
116where
117    TFn: FnOnce(T) -> TFuture,
118    TFuture: Future<Output = Option<U>>,
119{
120    type Output = Option<U>;
121
122    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
123        use AsyncAndThenProj::*;
124
125        match self.as_mut().project() {
126            None => Poll::Ready(Option::None),
127
128            Pending(payload) => {
129                let (x, f) = payload.take().expect("AsyncMap::Pending polled twice");
130                let future = f(x);
131                self.set(AsyncAndThen::Polling(future));
132                self.poll(cx)
133            }
134
135            Polling(future) => match future.poll(cx) {
136                Poll::Ready(Option::Some(d)) => Poll::Ready(Option::Some(d)),
137                Poll::Ready(Option::None) => Poll::Ready(Option::None),
138                Poll::Pending => Poll::Pending,
139            },
140        }
141    }
142}
143
144impl<T> AsyncMapExt<T> for Option<T> {
145    fn async_map<TFn, TFuture>(self, f: TFn) -> AsyncMap<T, TFn, TFuture>
146    where
147        TFn: FnOnce(T) -> TFuture,
148        TFuture: Future,
149    {
150        match self {
151            Some(v) => AsyncMap::Pending(Some((v, f))),
152            None => AsyncMap::None,
153        }
154    }
155
156    fn async_and_then<U, TFn, TFuture>(self, f: TFn) -> AsyncAndThen<T, TFn, TFuture>
157    where
158        TFn: FnOnce(T) -> TFuture,
159        TFuture: Future<Output = Option<U>>,
160    {
161        match self {
162            Some(v) => AsyncAndThen::Pending(Some((v, f))),
163            None => AsyncAndThen::None,
164        }
165    }
166}
167
168#[cfg(test)]
169mod test {
170    use super::AsyncMapExt;
171
172    #[tokio::test]
173    async fn map() {
174        assert_eq!(
175            Some(1).async_map(|x: i32| async move { x + 1 }).await,
176            Some(2),
177        );
178
179        assert_eq!(None.async_map(|x: i32| async move { x + 1 }).await, None);
180    }
181
182    #[tokio::test]
183    async fn and_then() {
184        assert_eq!(
185            Some(1)
186                .async_and_then(|x: i32| async move { Some(x + 1) })
187                .await,
188            Some(2),
189        );
190
191        assert_eq!(
192            Some(1)
193                .async_and_then(|_: i32| async move { Option::<i32>::None })
194                .await,
195            None
196        );
197
198        assert_eq!(
199            None.async_and_then(|x: i32| async move { Some(x + 1) })
200                .await,
201            None
202        );
203    }
204}