1use std::{
2 future::Future,
3 pin::Pin,
4 sync::mpsc,
5 task::{Context, Poll},
6};
7
8pub use yield_closures_impl::co;
9
10struct PendOnce(bool);
11
12impl Future for PendOnce {
13 type Output = ();
14
15 fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
16 if !self.0 {
17 self.get_mut().0 = true;
18 Poll::Pending
19 } else {
20 Poll::Ready(())
21 }
22 }
23}
24
25pub async fn pend_once() {
26 PendOnce(false).await
27}
28
29pub fn co<F, A, R, T>(f: F) -> impl FnMut(A) -> R
30where
31 F: FnOnce(mpsc::Receiver<A>, mpsc::SyncSender<R>) -> T,
32 T: Future<Output = std::convert::Infallible>,
33{
34 let (arg_tx, arg_rx) = mpsc::sync_channel(1);
35 let (yield_tx, yield_rx) = mpsc::sync_channel(1);
36 let mut future = Box::pin(f(arg_rx, yield_tx));
37 move |arg| {
38 arg_tx.send(arg).unwrap();
39 let waker_fn = waker_fn::waker_fn(|| {});
40 let mut cx = Context::from_waker(&waker_fn);
41 match Pin::new(&mut future).poll(&mut cx) {
42 Poll::Ready(_) => panic!("logic flaw"),
43 Poll::Pending => {}
44 }
45 yield_rx.recv().unwrap()
46 }
47}
48
49pub fn co0<F, R, T>(f: F) -> impl FnMut() -> R
50where
51 F: FnOnce(mpsc::Receiver<()>, mpsc::SyncSender<R>) -> T,
52 T: Future<Output = std::convert::Infallible>,
53{
54 let mut f = co::<F, (), R, T>(f);
55 move || f(())
56}
57
58pub fn co2<F, A1, A2, R, T>(f: F) -> impl FnMut(A1, A2) -> R
59where
60 F: FnOnce(mpsc::Receiver<(A1, A2)>, mpsc::SyncSender<R>) -> T,
61 T: Future<Output = std::convert::Infallible>,
62{
63 let mut f = co::<F, (A1, A2), R, T>(f);
64 move |a1, a2| f((a1, a2))
65}
66
67pub fn co3<F, A1, A2, A3, R, T>(f: F) -> impl FnMut(A1, A2, A3) -> R
68where
69 F: FnOnce(mpsc::Receiver<(A1, A2, A3)>, mpsc::SyncSender<R>) -> T,
70 T: Future<Output = std::convert::Infallible>,
71{
72 let mut f = co::<F, (A1, A2, A3), R, T>(f);
73 move |a1, a2, a3| f((a1, a2, a3))
74}
75
76#[doc(hidden)]
77#[macro_export]
78macro_rules! drop_args {
79 () => {{}};
80 ($x:ident, ) => {{
81 #[allow(clippy::drop_copy)]
82 drop($x)
83 }};
84 ($x0:ident, $x1:ident,) => {{
85 #[allow(clippy::drop_copy)]
86 drop($x0);
87 #[allow(clippy::drop_copy)]
88 drop($x1);
89 }};
90 ($x0:ident, $x1:ident, $x2:ident,) => {{
91 #[allow(clippy::drop_copy)]
92 drop($x0);
93 #[allow(clippy::drop_copy)]
94 drop($x1);
95 #[allow(clippy::drop_copy)]
96 drop($x2);
97 }};
98 ($x0:ident, $x1:ident, $x2:ident, $x3:ident,) => {{
99 #[allow(clippy::drop_copy)]
100 drop($x0);
101 #[allow(clippy::drop_copy)]
102 drop($x1);
103 #[allow(clippy::drop_copy)]
104 drop($x2);
105 #[allow(clippy::drop_copy)]
106 drop($x3);
107 }};
108}
109
110#[doc(hidden)]
111#[macro_export]
112macro_rules! reassign_args {
113 ($rx:ident,) => {{
114 let _ = $rx.recv().unwrap();
115 }};
116 ($rx:ident, $x:ident, ) => {{
117 $x = $rx.recv().unwrap();
118 }};
119 ($rx:ident, $x0:ident, $x1:ident,) => {{
120 let a = $rx.recv().unwrap();
121 $x0 = a.0;
122 $x1 = a.1;
123 }};
124 ($rx:ident, $x0:ident, $x1:ident, $x2:ident,) => {{
125 let a = $rx.recv().unwrap();
126 $x0 = a.0;
127 $x1 = a.1;
128 $x2 = a.2;
129 }};
130 ($rx:ident, $x0:ident, $x1:ident, $x2:ident, $x3:ident,) => {{
131 let a = $rx.recv().unwrap();
132 $x0 = a.0;
133 $x1 = a.1;
134 $x2 = a.2;
135 $x3 = a.3;
136 }};
137}