Skip to main content

datex_core/utils/
maybe_async.rs

1use core::{
2    future::Future,
3    pin::Pin,
4    task::{Context, Poll},
5};
6
7pub struct Ready<T>(Option<T>);
8pub fn ready<T>(t: T) -> Ready<T> {
9    Ready(Some(t))
10}
11impl<T> Future for Ready<T> {
12    type Output = T;
13
14    fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<T> {
15        let this = unsafe { self.get_unchecked_mut() };
16        Poll::Ready(this.0.take().unwrap())
17    }
18}
19pub enum MaybeAsync<T, F: Future<Output = T>> {
20    Sync(T),
21    Async(F),
22}
23
24pub enum Either<L, R> {
25    Left(L),
26    Right(R),
27}
28impl<L, R> Future for Either<L, R>
29where
30    L: Future,
31    R: Future<Output = L::Output>,
32{
33    type Output = L::Output;
34
35    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
36        unsafe {
37            match self.get_unchecked_mut() {
38                Either::Left(l) => Pin::new_unchecked(l).poll(cx),
39                Either::Right(r) => Pin::new_unchecked(r).poll(cx),
40            }
41        }
42    }
43}
44
45pub type MaybeAsyncResult<T, E, F> = MaybeAsync<Result<T, E>, F>;
46
47impl<T, F: Future<Output = T>> MaybeAsync<T, F> {
48    pub fn into_future(self) -> Either<Ready<T>, F> {
49        match self {
50            MaybeAsync::Sync(v) => Either::Left(ready(v)),
51            MaybeAsync::Async(f) => Either::Right(f),
52        }
53    }
54
55    pub fn map<U, Func>(self, func: Func) -> MaybeAsync<U, MapFuture<F, Func>>
56    where
57        Func: FnOnce(T) -> U,
58    {
59        match self {
60            MaybeAsync::Sync(v) => MaybeAsync::Sync(func(v)),
61            MaybeAsync::Async(fut) => MaybeAsync::Async(MapFuture {
62                fut,
63                func: Some(func),
64            }),
65        }
66    }
67}
68
69pub struct MapFuture<Fut, Func> {
70    fut: Fut,
71    func: Option<Func>,
72}
73
74impl<T, U, Fut, Func> Future for MapFuture<Fut, Func>
75where
76    Fut: Future<Output = T>,
77    Func: FnOnce(T) -> U,
78{
79    type Output = U;
80
81    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<U> {
82        unsafe {
83            let this = self.get_unchecked_mut();
84            let t = match Pin::new_unchecked(&mut this.fut).poll(cx) {
85                Poll::Pending => return Poll::Pending,
86                Poll::Ready(t) => t,
87            };
88            let f = this.func.take().unwrap();
89            Poll::Ready(f(t))
90        }
91    }
92}
93
94pub enum SyncOrAsync<SyncValue, AsyncValue, F: Future<Output = AsyncValue>> {
95    Sync(SyncValue),
96    Async(F),
97}
98
99pub enum SyncOrAsyncResolved<SyncValue, AsyncValue> {
100    Sync(SyncValue),
101    Async(AsyncValue),
102}
103
104use core::marker::PhantomData;
105
106pub struct MapToResolvedAsync<SyncValue, Fut> {
107    fut: Fut,
108    _pd: PhantomData<SyncValue>,
109}
110
111impl<SyncValue, AsyncValue, Fut> core::future::Future
112    for MapToResolvedAsync<SyncValue, Fut>
113where
114    Fut: core::future::Future<Output = AsyncValue>,
115{
116    type Output = SyncOrAsyncResolved<SyncValue, AsyncValue>;
117
118    fn poll(
119        self: core::pin::Pin<&mut Self>,
120        cx: &mut core::task::Context<'_>,
121    ) -> core::task::Poll<Self::Output> {
122        unsafe {
123            let this = self.get_unchecked_mut();
124            match core::pin::Pin::new_unchecked(&mut this.fut).poll(cx) {
125                core::task::Poll::Pending => core::task::Poll::Pending,
126                core::task::Poll::Ready(v) => {
127                    core::task::Poll::Ready(SyncOrAsyncResolved::Async(v))
128                }
129            }
130        }
131    }
132}
133
134pub type SyncOrAsyncResult<SyncOkValue, AsyncOkValue, E, F> =
135    SyncOrAsync<Result<SyncOkValue, E>, Result<AsyncOkValue, E>, F>;
136
137pub struct ErrOptionFuture<Fut>(Fut);
138
139impl<AsyncOkValue, E, Fut> Future for ErrOptionFuture<Fut>
140where
141    Fut: Future<Output = Result<AsyncOkValue, E>>,
142{
143    type Output = Option<E>;
144
145    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<E>> {
146        unsafe {
147            let this = self.get_unchecked_mut();
148            match Pin::new_unchecked(&mut this.0).poll(cx) {
149                Poll::Pending => Poll::Pending,
150                Poll::Ready(r) => Poll::Ready(r.err()),
151            }
152        }
153    }
154}
155
156pub struct OkOptionFuture<SyncOkValue, AsyncOkValue, E, Fut> {
157    fut: Fut,
158    _phantom: core::marker::PhantomData<(SyncOkValue, AsyncOkValue, E)>,
159}
160
161impl<SyncOkValue, AsyncOkValue, E, Fut> Future
162    for OkOptionFuture<SyncOkValue, AsyncOkValue, E, Fut>
163where
164    Fut: Future<Output = Result<AsyncOkValue, E>>,
165{
166    type Output = Option<SyncOrAsyncResolved<SyncOkValue, AsyncOkValue>>;
167
168    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
169        unsafe {
170            let this = self.get_unchecked_mut();
171            match Pin::new_unchecked(&mut this.fut).poll(cx) {
172                Poll::Pending => Poll::Pending,
173                Poll::Ready(r) => {
174                    Poll::Ready(r.ok().map(SyncOrAsyncResolved::Async))
175                }
176            }
177        }
178    }
179}
180
181impl<SyncOkValue, AsyncOkValue, E, F>
182    SyncOrAsyncResult<SyncOkValue, AsyncOkValue, E, F>
183where
184    F: Future<Output = Result<AsyncOkValue, E>>,
185{
186    pub fn into_error_future(
187        self,
188    ) -> Either<Ready<Option<E>>, ErrOptionFuture<F>> {
189        match self {
190            SyncOrAsync::Sync(r) => Either::Left(ready(r.err())),
191            SyncOrAsync::Async(fut) => Either::Right(ErrOptionFuture(fut)),
192        }
193    }
194
195    pub fn into_ok_future(
196        self,
197    ) -> Either<
198        Ready<Option<SyncOrAsyncResolved<SyncOkValue, AsyncOkValue>>>,
199        OkOptionFuture<SyncOkValue, AsyncOkValue, E, F>,
200    > {
201        match self {
202            SyncOrAsync::Sync(r) => {
203                let v = r.ok().map(SyncOrAsyncResolved::Sync);
204                Either::Left(ready(v))
205            }
206            SyncOrAsync::Async(fut) => Either::Right(OkOptionFuture {
207                fut,
208                _phantom: core::marker::PhantomData,
209            }),
210        }
211    }
212
213    pub fn into_result(self) -> Either<Ready<Result<AsyncOkValue, E>>, F>
214    where
215        SyncOkValue: Into<AsyncOkValue>,
216    {
217        match self {
218            SyncOrAsync::Sync(r) => Either::Left(ready(r.map(Into::into))),
219            SyncOrAsync::Async(fut) => Either::Right(fut),
220        }
221    }
222}
223
224pub struct FlattenFuture<OuterF, InnerF, T> {
225    state: FlattenState<OuterF, InnerF, T>,
226}
227
228enum FlattenState<OuterF, InnerF, T> {
229    Outer(OuterF),
230    Inner(InnerF),
231    Done(Option<T>),
232}
233
234impl<OuterF, InnerF, T> Future for FlattenFuture<OuterF, InnerF, T>
235where
236    OuterF: Future<Output = MaybeAsync<T, InnerF>>,
237    InnerF: Future<Output = T>,
238{
239    type Output = T;
240
241    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
242        unsafe {
243            let this = self.get_unchecked_mut();
244            loop {
245                match &mut this.state {
246                    FlattenState::Outer(outer) => {
247                        let inner = match Pin::new_unchecked(outer).poll(cx) {
248                            Poll::Pending => return Poll::Pending,
249                            Poll::Ready(v) => v,
250                        };
251                        match inner {
252                            MaybeAsync::Sync(v) => {
253                                this.state = FlattenState::Done(Some(v))
254                            }
255                            MaybeAsync::Async(f) => {
256                                this.state = FlattenState::Inner(f)
257                            }
258                        }
259                    }
260                    FlattenState::Inner(inner) => {
261                        let v = match Pin::new_unchecked(inner).poll(cx) {
262                            Poll::Pending => return Poll::Pending,
263                            Poll::Ready(v) => v,
264                        };
265                        this.state = FlattenState::Done(Some(v));
266                    }
267                    FlattenState::Done(v) => {
268                        return Poll::Ready(v.take().unwrap());
269                    }
270                }
271            }
272        }
273    }
274}
275
276impl<T, InnerF, OuterF> MaybeAsync<MaybeAsync<T, InnerF>, OuterF>
277where
278    InnerF: Future<Output = T>,
279    OuterF: Future<Output = MaybeAsync<T, InnerF>>,
280{
281    pub fn flatten(self) -> MaybeAsync<T, FlattenFuture<OuterF, InnerF, T>> {
282        match self {
283            MaybeAsync::Sync(inner) => match inner {
284                MaybeAsync::Sync(v) => MaybeAsync::Sync(v),
285                MaybeAsync::Async(inner_fut) => {
286                    MaybeAsync::Async(FlattenFuture {
287                        state: FlattenState::Inner(inner_fut),
288                    })
289                }
290            },
291            MaybeAsync::Async(outer_fut) => MaybeAsync::Async(FlattenFuture {
292                state: FlattenState::Outer(outer_fut),
293            }),
294        }
295    }
296}
297
298pub struct MapAsyncResultToResolved<SyncOkValue, AsyncOkValue, E, Fut> {
299    fut: Fut,
300    _pd: PhantomData<(SyncOkValue, AsyncOkValue, E)>,
301}
302
303impl<SyncOkValue, AsyncOkValue, E, Fut> Future
304    for MapAsyncResultToResolved<SyncOkValue, AsyncOkValue, E, Fut>
305where
306    Fut: Future<Output = Result<AsyncOkValue, E>>,
307{
308    type Output =
309        SyncOrAsyncResolved<Result<SyncOkValue, E>, Result<AsyncOkValue, E>>;
310
311    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
312        unsafe {
313            let this = self.get_unchecked_mut();
314            match Pin::new_unchecked(&mut this.fut).poll(cx) {
315                Poll::Pending => Poll::Pending,
316                Poll::Ready(r) => Poll::Ready(SyncOrAsyncResolved::Async(r)),
317            }
318        }
319    }
320}
321
322impl<SyncOkValue, AsyncOkValue, E, F>
323    SyncOrAsyncResult<SyncOkValue, AsyncOkValue, E, F>
324where
325    F: Future<Output = Result<AsyncOkValue, E>>,
326{
327    pub fn into_future(
328        self,
329    ) -> Either<
330        Ready<
331            SyncOrAsyncResolved<
332                Result<SyncOkValue, E>,
333                Result<AsyncOkValue, E>,
334            >,
335        >,
336        MapAsyncResultToResolved<SyncOkValue, AsyncOkValue, E, F>,
337    > {
338        match self {
339            SyncOrAsync::Sync(r_sync) => {
340                Either::Left(ready(SyncOrAsyncResolved::Sync(r_sync)))
341            }
342            SyncOrAsync::Async(fut) => {
343                Either::Right(MapAsyncResultToResolved {
344                    fut,
345                    _pd: PhantomData,
346                })
347            }
348        }
349    }
350}