1 // Copyright (c) 2023 Huawei Device Co., Ltd.
2 // Licensed under the Apache License, Version 2.0 (the "License");
3 // you may not use this file except in compliance with the License.
4 // You may obtain a copy of the License at
5 //
6 //     http://www.apache.org/licenses/LICENSE-2.0
7 //
8 // Unless required by applicable law or agreed to in writing, software
9 // distributed under the License is distributed on an "AS IS" BASIS,
10 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11 // See the License for the specific language governing permissions and
12 // limitations under the License.
13 
14 /// Waits on multiple concurrent branches, returning when the **first** branch
15 /// completes, cancelling the remaining branches.
16 ///
17 /// There is **no upper limit** on the number of branches.
18 ///
19 /// The `select!` macro must be used inside of async functions, closures, and
20 /// blocks.
21 ///
22 /// Branchex must comply with the following format:
23 ///
24 /// ```text
25 /// <pattern> = <async expression> (, if <precondition>)? => <handler>,
26 /// ```
27 ///
28 /// The `<async expression>` for all branches will run concurrently. Once the
29 /// first expression is completed, the corresponding `<handler>` will be
30 /// executed. Each branch is optional `if <precondition>`. if the
31 /// `<precondition>` returns false, the corresponding `<handler>` will not be
32 /// executed.
33 ///
34 /// When none of the branches can match, Executes the else expression.
35 ///
36 /// ```text
37 /// else => <expression>
38 /// ```
39 ///
40 /// # Examples
41 ///
42 /// Basic select with two branches.
43 /// ```
44 /// async fn do_async1() {
45 ///     // do task
46 /// }
47 /// async fn do_async2() {
48 ///     // do task
49 /// }
50 /// async fn select_test() {
51 ///     ylong_runtime::select! {
52 ///         _ = do_async1() => {
53 ///             println!("do_async1() completed first");
54 ///         },
55 ///         _ = do_async2() => {
56 ///             println!("do_async2() completed first");
57 ///         }
58 ///     }
59 /// }
60 /// ```
61 ///
62 /// # Examples
63 ///
64 /// Uses if to filter asynchronous tasks
65 /// ```
66 /// async fn do_async1() -> i32 {
67 ///     1
68 /// }
69 /// async fn do_async2() -> i32 {
70 ///     2
71 /// }
72 /// async fn do_async3() -> bool {
73 ///     false
74 /// }
75 /// async fn select_test() {
76 ///     let mut count = 0;
77 ///     ylong_runtime::select! {
78 ///         a = do_async1(), if false => {
79 ///             count += a;
80 ///             println!("do_async1() completed first{:?}", a);
81 ///         },
82 ///         b = do_async2() => {
83 ///             count += b;
84 ///             println!("do_async2() completed first{:?}", b);
85 ///         },
86 ///         c = do_async3(), if false => {
87 ///             if c {
88 ///                 println!("do_async3() completed true");
89 ///             }
90 ///             else {
91 ///                 println!("do_async3() completed false");
92 ///             }
93 ///         }
94 ///     }
95 ///     assert_eq!(count, 2);
96 /// }
97 /// ```
98 /// # Examples
99 ///
100 /// Repeated uses select! until all task return.
101 /// ```
102 /// #[cfg(feature = "sync")]
103 /// async fn select_channel() {
104 ///     let (tx1, mut rx1) = ylong_runtime::sync::oneshot::channel();
105 ///     let (tx2, mut rx2) = ylong_runtime::sync::oneshot::channel();
106 ///
107 ///     ylong_runtime::spawn(async move {
108 ///         tx1.send("first").unwrap();
109 ///     });
110 ///
111 ///     ylong_runtime::spawn(async move {
112 ///         tx2.send("second").unwrap();
113 ///     });
114 ///
115 ///     let mut a = None;
116 ///     let mut b = None;
117 ///
118 ///     while a.is_none() || b.is_none() {
119 ///         ylong_runtime::select! {
120 ///             v1 = (&mut rx1), if a.is_none() => a = Some(v1.unwrap()),
121 ///             v2 = (&mut rx2), if b.is_none() => b = Some(v2.unwrap()),
122 ///         }
123 ///     }
124 ///
125 ///     let res = (a.unwrap(), b.unwrap());
126 ///
127 ///     assert_eq!(res.0, "first");
128 ///     assert_eq!(res.1, "second");
129 /// }
130 /// ```
131 /// # Examples
132 ///
133 /// Uses 'biased' to execute four task in the specified sequence.
134 /// It will poll branches with order, the first branch will be polled first.
135 /// ```
136 /// async fn select_biased() {
137 ///     let mut count = 0u8;
138 ///
139 ///     loop {
140 ///         ylong_runtime::select! {
141 ///             biased;
142 ///             _ = async {}, if count < 1 => {
143 ///                 count += 1;
144 ///                 assert_eq!(count, 1);
145 ///             }
146 ///             _ = async {}, if count < 2 => {
147 ///                 count += 1;
148 ///                 assert_eq!(count, 2);
149 ///             }
150 ///             _ = async {}, if count < 3 => {
151 ///                 count += 1;
152 ///                 assert_eq!(count, 3);
153 ///             }
154 ///             _ = async {}, if count < 4 => {
155 ///                 count += 1;
156 ///                 assert_eq!(count, 4);
157 ///             }
158 ///             else => {
159 ///                 break;
160 ///             }
161 ///         }
162 ///     }
163 ///
164 ///     assert_eq!(count, 4);
165 /// }
166 /// ```
167 #[macro_export]
168 macro_rules! select {
169     ( {
170         // Branch from which the execution starts.
171         random = $bool:expr;
172         // Branch count.
173         ( $count:expr, $($_n:tt)* )
174         // ( index:expr ) Branch's index
175         $( ( $index:expr, $($_i:tt)* ) $bind:pat = $fut:expr, if $c:expr => $handle:expr, )+
176         // When all branches fail, executes this;
177         ; $else:expr
178 
179     }) => {{
180         use std::future::Future;
181         use std::pin::Pin;
182         use std::task::Poll::{Ready, Pending};
183 
184         enum Out<T> {
185             Finish(T),
186             Fail,
187         }
188 
189         let branches_count: usize = $count;
190 
191         // The ith element indicates whether the ith branch is available.
192         let mut match_result = [true; $count];
193         // Handle preconditions, this step cannot be handled within poll_fn()
194         $(
195             if (!$c)
196             {
197                 match_result[$index] = false;
198             }
199         )*
200 
201         // When a branch ready first, modify this variable to
202         // branch's index to ensure that the branch is executed first.
203         use $crate::fastrand::fast_random;
204         let mut random_number = fast_random() as usize;
205 
206         let output = {
207             let mut futures = ( $( $fut , )+ );
208             let futures = &mut futures;
209 
210             $crate::futures::poll_fn(|cx| {
211                 let mut anyone_pending = false;
212                 let random = $bool;
213 
214                 for i in 0..branches_count {
215                     let branch  = match random {
216                         true => {(random_number + i) % branches_count },
217                         false => i
218                     };
219 
220                     $(
221                         if (branch == $index && match_result[branch])
222                         {
223                             let ( $($_i,)* fut, .. ) = &mut *futures;
224                             let fut = unsafe { Pin::new_unchecked(fut) };
225                             match Future::poll(fut, cx) {
226                                 Ready(out) => {
227                                     // Check if the returned value match the user input.
228                                     match &out {
229                                         // If the match is successful, the inner value will
230                                         // never used, so here has to use unused_variables.
231                                         #[allow(unused_variables)]
232                                         $bind => {
233                                             // Change the random_number, ensure when return ready, this branch is executed first.
234                                             random_number = branch;
235                                             return Ready($crate::tuple_form!(($count) with Out::Fail except Out::Finish(out) at ($($_i)*)))
236                                         },
237                                         // If the match fails, this branch set false, and wait for the next branch to complete.
238                                         // When user input is not match type, this patterns is unreachable.
239                                         #[allow(unreachable_patterns)]
240                                         _ => {
241                                             match_result[branch] = false;
242                                             continue;
243                                         }
244                                     }
245                                 },
246                                 Pending => {
247                                     anyone_pending = true;
248                                     continue;
249                                 }
250                             };
251                         }
252                     )*
253                 }
254 
255                 if anyone_pending {
256                     Pending
257                 } else {
258                     Ready($crate::tuple_form!(($count) with Out::Fail except Out::Fail at ($($_n)*)))
259                 }
260             }).await
261         };
262 
263         match output {
264             $(
265                 $crate::tuple_form!(($count) with Out::Fail except Out::Finish($bind) at ($($_i)*)) => $handle,
266             )*
267             $crate::tuple_form!(($count) with Out::Fail except Out::Fail at ($($_n)*)) => $else,
268             // If there is only one branch and the user match for that branch returns `_`,
269             // there will be an unreachable pattern alert.
270             #[allow(unreachable_patterns)]
271             _ => unreachable!("finally match fail"),
272         }
273     }};
274 
275     // if there is no 'else' branch, add the default 'else' branch.
276     ( { random = $bool:expr; $($t:tt)* } ) => {
277         $crate::select!({ random = $bool; $($t)*; panic!("select!: All the branches failed.") })
278     };
279     // if there is an 'else' branch, add the 'else' branch into {}.
280     ( { random = $bool:expr; $($t:tt)* } else => $else:expr $(,)?) => {
281         $crate::select!({ random = $bool; $($t)*; $else })
282     };
283     // Recursively join a branch into {}.
284     // The branch is separated by ',', has 'if' conditions and executes block finally.
285     ( { random = $bool:expr; ( $s:expr, $($_n:tt)* ) $($t:tt)* } $p:pat = $f:expr, if $c:expr => $h:block, $($r:tt)* ) => {
286         $crate::select!({ random = $bool; ( $s + 1, $($_n)*_) $($t)* ($s, $($_n)*) $p = $f, if $c => $h, } $($r)*)
287     };
288     // Recursively join a branch into {}.
289     // The branch is separated by ',', does not has 'if' conditions and executes block finally.
290     ( { random = $bool:expr; ( $s:expr, $($_n:tt)* ) $($t:tt)* } $p:pat = $f:expr => $h:block, $($r:tt)* ) => {
291         $crate::select!({ random = $bool; ( $s + 1, $($_n)*_) $($t)* ($s, $($_n)*) $p = $f, if true => $h, } $($r)*)
292     };
293     // Recursively join a branch into {}.
294     // The branch is separated by ' ', has 'if' conditions and executes block finally.
295     ( { random = $bool:expr; ( $s:expr, $($_n:tt)* ) $($t:tt)* } $p:pat = $f:expr, if $c:expr => $h:block $($r:tt)* ) => {
296         $crate::select!({ random = $bool; ( $s + 1, $($_n)*_) $($t)* ($s, $($_n)*) $p = $f, if $c => $h, } $($r)*)
297     };
298     // Recursively join a branch into {}.
299     // The branch is separated by ' ', does not has 'if' conditions and executes block finally.
300     ( { random = $bool:expr; ( $s:expr, $($_n:tt)* ) $($t:tt)* } $p:pat = $f:expr => $h:block $($r:tt)* ) => {
301         $crate::select!( { random = $bool; ( $s + 1, $($_n)*_) $($t)* ($s, $($_n)*) $p = $f, if true => $h, } $($r)*)
302     };
303     // Recursively join a branch into {}.
304     // The branch is separated by ',', has 'if' conditions and executes expressions finally.
305     ( { random = $bool:expr; ( $s:expr, $($_n:tt)*  ) $($t:tt)* } $p:pat = $f:expr, if $c:expr => $h:expr, $($r:tt)* ) => {
306         $crate::select!({ random = $bool; ( $s + 1, $($_n)*_ ) $($t)* ($s, $($_n)*) $p = $f, if $c => $h, } $($r)*)
307     };
308     // Recursively join a branch into {}.
309     // The branch is separated by ',', does not has 'if' conditions and executes expressions finally.
310     ( { random = $bool:expr; ( $s:expr, $($_n:tt)*  ) $($t:tt)* } $p:pat = $f:expr => $h:expr, $($r:tt)* ) => {
311         $crate::select!({ random = $bool; ( $s + 1, $($_n)*_ ) $($t)* ($s, $($_n)*) $p = $f, if true => $h, } $($r)*)
312     };
313     // Recursively join the last branch into {}.
314     // The branch is separated by ',', has 'if' conditions and executes expressions finally.
315     // If the branch executes expressions finally, it can't separated by ' '.
316     ( { random = $bool:expr; ( $s:expr, $($_n:tt)*  ) $($t:tt)* } $p:pat = $f:expr, if $c:expr => $h:expr ) => {
317         $crate::select!({ random = $bool; ( $s + 1, $($_n)*_ ) $($t)* ($s, $($_n)*) $p = $f, if $c => $h, })
318     };
319     // Recursively join the last branch into {}.
320     // The branch is separated by ',', does not has 'if' conditions and executes expressions finally.
321     // If the branch executes expressions finally, it can't separated by ' '.
322     ( { random = $bool:expr; ( $s:expr, $($_n:tt)*  ) $($t:tt)* } $p:pat = $f:expr => $h:expr ) => {
323         $crate::select!({ random = $bool; ( $s + 1, $($_n)*_ ) $($t)* ($s, $($_n)*) $p = $f, if true => $h, })
324     };
325 
326     // Usage entry. Starts with the first branch.
327     (biased; $p:pat = $($t:tt)* ) => {
328         $crate::select!({ random = false; ( 0,) } $p = $($t)*)
329     };
330     // Usage entry. Starts with random branch.
331     ( $p:pat = $($t:tt)* ) => {
332         $crate::select!({ random = true; ( 0,) } $p = $($t)*)
333     };
334     // There is no branch.
335     () => {
336         compile_error!("select! requires at least one branch.")
337     };
338 }
339