1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
use core::future::Future;
use core::pin::Pin;
use core::task::{Poll, Waker};
mod poll_fn;
pub use poll_fn::{PollFn, poll_fn};
mod poll_state;
pub use poll_state::{PollState, poll_state};
mod stolen_from_lite;
pub async fn pending() {
poll_fn(|_| Poll::Pending).await
}
pub async fn ready<T>(val: T) -> T {
poll_state(Some(val), |val, _| Poll::Ready(val.take().unwrap())).await
}
pub async fn waker() -> Waker {
poll_fn(|ctx| Poll::Ready(ctx.waker().clone())).await
}
pub async fn sleep() {
poll_state(false, |done, _| {
if *done { Poll::Ready(()) }
else {
*done = true;
Poll::Pending
}
}).await
}
pub async fn next_poll<F: Future>(mut f: F) -> Result<F::Output, F> {
{
let mut pin = unsafe { Pin::new_unchecked(&mut f) };
poll_fn(|ctx| {
match pin.as_mut().poll(ctx) {
Poll::Ready(val) => Poll::Ready(Ok(val)),
Poll::Pending => Poll::Ready(Err(())),
}
}).await
}.map_err(|_| f)
}
pub async fn yield_once() {
poll_state(false, |done, ctx| {
if *done { Poll::Ready(()) }
else {
*done = true;
ctx.waker().wake_by_ref();
Poll::Pending
}
}).await
}
pub async fn or<T>(a: impl Future<Output=T>, b: impl Future<Output=T>) -> T {
pin!(a, b);
poll_fn(move |ctx|{
match a.as_mut().poll(ctx) {
Poll::Ready(val) => Poll::Ready(val),
Poll::Pending => {
match b.as_mut().poll(ctx) {
Poll::Ready(val) => Poll::Ready(val),
Poll::Pending => Poll::Pending,
}
}
}
}).await
}
pub async fn zip<A, B>(a: impl Future<Output=A>, b: impl Future<Output=B>) -> (A, B) {
pin!(a, b);
poll_state((None, None), move |(c, d), ctx|{
if c.is_none() {
if let Poll::Ready(val) = a.as_mut().poll(ctx) {
*c = Some(val);
}
}
if d.is_none() {
if let Poll::Ready(val) = b.as_mut().poll(ctx) {
*d = Some(val);
}
}
if c.is_some() && d.is_some() {
Poll::Ready((c.take().unwrap(), d.take().unwrap()))
} else {
Poll::Pending
}
}).await
}
#[macro_export]
macro_rules! zips {
($($es:expr),+ $(,)?) => {
$crate::poll_state(
$crate::__internal_fold_with!($crate::zip, $($es),+),
|zips, ctx| {
use ::core::future::Future;
use ::core::pin::Pin;
use ::core::task::Poll;
let zips = unsafe { Pin::new_unchecked(zips) };
if let Poll::Ready(val) = zips.poll(ctx) {
Poll::Ready($crate::zips!(@flatten; ; val; $($es),+))
} else {
Poll::Pending
}
},
)
};
(@flatten; $($prev:expr,)*; $tuple:expr; $e:expr) => {
($($prev,)* $tuple)
};
(@flatten; $($prev:expr,)*; $tuple:expr; $e:expr, $($es:expr),+) => {
$crate::zips!(@flatten; $($prev,)* $tuple.0,; $tuple.1; $($es),+)
};
}