1use core::future::Future;
2use core::pin::Pin;
3use core::task::{Context, Poll};
4use pin_project::pin_project;
5
6pub trait AsyncMapExt<T> {
7 fn async_map<TFn, TFuture>(self, f: TFn) -> AsyncMap<T, TFn, TFuture>
31 where
32 TFn: FnOnce(T) -> TFuture,
33 TFuture: Future;
34
35 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}