simpleos/sys/
join.rs

1use core::future::Future;
2use core::pin::Pin;
3use core::task::{Context, Poll};
4
5use alloc::boxed::Box;
6
7#[allow(unused)]
8pub struct Join2<F1, F2>
9where
10    F1: Future,
11    F2: Future,
12{
13    future1: Pin<Box<F1>>,
14    future2: Pin<Box<F2>>,
15    future1_done: bool,
16    future2_done: bool,
17    output1: Option<F1::Output>,
18    output2: Option<F2::Output>,
19}
20
21impl<F1, F2> Join2<F1, F2>
22where
23    F1: Future,
24    F2: Future,
25{
26    pub fn new(future1: F1, future2: F2) -> Self {
27        Self {
28            future1: Box::pin(future1),
29            future2: Box::pin(future2),
30            future1_done: false,
31            future2_done: false,
32            output1: None,
33            output2: None,
34        }
35    }
36}
37
38impl<F1, F2> Future for Join2<F1, F2>
39where
40    F1: Future,
41    F2: Future,
42{
43    type Output = (F1::Output, F2::Output);
44
45    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
46        let this = unsafe { self.get_unchecked_mut() };
47
48        if !this.future1_done {
49            match this.future1.as_mut().poll(cx) {
50                Poll::Ready(output) => {
51                    this.output1 = Some(output);
52                    this.future1_done = true;
53                }
54                Poll::Pending => {}
55            }
56        }
57
58        if !this.future2_done {
59            match this.future2.as_mut().poll(cx) {
60                Poll::Ready(output) => {
61                    this.output2 = Some(output);
62                    this.future2_done = true;
63                }
64                Poll::Pending => {}
65            }
66        }
67
68        if this.future1_done && this.future2_done {
69            let output1 = this.output1.take().unwrap();
70            let output2 = this.output2.take().unwrap();
71            Poll::Ready((output1, output2))
72        } else {
73            cx.waker().wake_by_ref();
74            Poll::Pending
75        }
76    }
77}
78
79impl<F1, F2> Unpin for Join2<F1, F2>
80where
81    F1: Future,
82    F2: Future,
83{
84}
85
86#[allow(unused)]
87pub fn join<F1, F2>(future1: F1, future2: F2) -> Join2<F1, F2>
88where
89    F1: Future,
90    F2: Future,
91{
92    Join2::new(future1, future2)
93}
94
95// 改进版宏,自动包含 .await
96#[macro_export]
97macro_rules! join {
98    ($fut1:expr, $fut2:expr) => {
99        async {
100            $crate::sys::join($fut1, $fut2).await
101        }.await
102    };
103    ($fut1:expr, $fut2:expr, $fut3:expr) => {
104        async {
105            let ((r1, r2), r3) = $crate::sys::join(
106                $crate::sys::join($fut1, $fut2),
107                $fut3
108            ).await;
109            (r1, r2, r3)
110        }.await
111    };
112    ($fut1:expr, $fut2:expr, $fut3:expr, $fut4:expr) => {
113        async {
114            let ((r1, r2), (r3, r4)) = $crate::sys::join(
115                $crate::sys::join($fut1, $fut2),
116                $crate::sys::join($fut3, $fut4)
117            ).await;
118            (r1, r2, r3, r4)
119        }.await
120    };
121}
simpleos/sys/join.rs

simpleos/sys/
join.rs
