yield_closures/
lib.rs

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}