nuclei/macros/
select.rs

1//! `select!`: waits on several futures at once and runs the handler of the
2//! first one to finish whose output matches its pattern.
3//!
4//! The expansion keeps the futures pinned in a tuple and polls them from one
5//! `poll_fn`; a future that finishes with an output its pattern rejects
6//! disables its branch and the others go on. The handlers run after the
7//! futures are dropped and outside any closure or loop of the expansion, so
8//! `return`, `?`, `break` and `continue` in a handler act on the code around
9//! the `select!`, and a handler may use again what the futures borrowed.
10
11/// Waits on several branches at once and runs the handler of the first to
12/// finish, dropping the others.
13///
14/// Each branch is `<pattern> = <future> (, if <precondition>)? => <handler>`,
15/// with commas between the branches (optional after a block handler), and
16/// an `else => <expression>` branch may come last. `select!`:
17///
18/// 1. Evaluates every precondition; a branch whose precondition is `false`
19///    is disabled.
20/// 2. Evaluates every future expression (of disabled branches too), each
21///    an [`IntoFuture`](core::future::IntoFuture).
22/// 3. Polls the futures of the enabled branches, starting at a random one so
23///    no branch is favoured, or in the order written after `biased;`.
24/// 4. When a future finishes, its branch is disabled. If the output matches
25///    the branch's pattern, the other futures are dropped and the handler
26///    runs with the pattern's bindings; if not, the output is dropped and
27///    the other branches go on.
28/// 5. When every branch is disabled, the `else` expression is the value.
29///
30/// The value of `select!` is the value of the handler that ran, so every
31/// handler and `else` has the same type. A `select!` takes at most 64
32/// branches.
33///
34/// # Panics
35///
36/// When every branch is disabled and there is no `else` branch.
37///
38/// # Examples
39///
40/// ```
41/// use nuclei::sync::mpsc;
42///
43/// # nuclei::block_on(async {
44/// let (tx, mut numbers) = mpsc::unbounded_channel::<u32>();
45/// let (tick_tx, mut ticks) = mpsc::unbounded_channel::<()>();
46/// tx.send(1).unwrap();
47/// tx.send(2).unwrap();
48/// tick_tx.send(()).unwrap();
49/// drop((tx, tick_tx));
50/// let (mut seen, mut ticked) = (Vec::new(), 0);
51/// // Each closed channel gives `None`, which disables its branch; once
52/// // both are closed, `else` ends the loop.
53/// loop {
54///     nuclei::select! {
55///         Some(n) = numbers.recv() => seen.push(n),
56///         Some(()) = ticks.recv() => ticked += 1,
57///         else => break,
58///     }
59/// }
60/// assert_eq!(seen, [1, 2]);
61/// assert_eq!(ticked, 1);
62/// # });
63/// ```
64///
65/// Patterns bind as written, `mut`, `ref` and `&` included, in every
66/// edition:
67///
68/// ```edition2021
69/// # nuclei::block_on(async {
70/// let label = String::from("one");
71/// let n = nuclei::select! {
72///     (mut n, ref name) = async { (1u32, &label) } => {
73///         n += name.len() as u32;
74///         n
75///     }
76/// };
77/// assert_eq!(n, 4);
78/// # });
79/// ```
80#[macro_export]
81macro_rules! select {
82    (biased;) => {
83        ::core::compile_error!("`select!` needs at least one branch")
84    };
85    (biased; $($branches:tt)+) => {
86        $crate::__select!(@branch [biased] [] [] $($branches)+)
87    };
88    () => {
89        ::core::compile_error!("`select!` needs at least one branch")
90    };
91    ($($branches:tt)+) => {
92        $crate::__select!(@branch [fair] [] [] $($branches)+)
93    };
94}
95
96/// The branch muncher of [`select!`]. The state is
97/// `[mode] [parsed branches] [one _ per parsed branch]`; a parsed branch is
98/// `([its _ list] [its pattern] future, precondition => handler)`.
99#[doc(hidden)]
100#[macro_export]
101macro_rules! __select {
102    // The `else` branch, last.
103    (@branch $mode:tt [$($done:tt)*] [$($all:tt)*] else => $else:expr $(,)?) => {
104        $crate::__select!(@emit $mode [$($done)*] [$($all)*] [$else])
105    };
106    // No branch left and no `else`.
107    (@branch $mode:tt [$($done:tt)*] [$($all:tt)*]) => {
108        $crate::__select!(
109            @emit $mode [$($done)*] [$($all)*]
110            [::core::panic!("all branches of `select!` are disabled and there is no `else` branch")]
111        )
112    };
113    // A branch: its pattern runs up to the first `=` outside a group.
114    (@branch $mode:tt [$($done:tt)*] [$($all:tt)*] $($rest:tt)+) => {
115        $crate::__select!(@pat $mode [$($done)*] [$($all)*] [] $($rest)+)
116    };
117
118    (@pat $mode:tt [$($done:tt)*] [$($all:tt)*] [$($pat:tt)+] = $future:expr, if $cond:expr => $($rest:tt)+) => {
119        $crate::__select!(@handler $mode [$($done)*] [$($all)*] [$($pat)+] $future, $cond => $($rest)+)
120    };
121    (@pat $mode:tt [$($done:tt)*] [$($all:tt)*] [$($pat:tt)+] = $future:expr => $($rest:tt)+) => {
122        $crate::__select!(@handler $mode [$($done)*] [$($all)*] [$($pat)+] $future, true => $($rest)+)
123    };
124    (@pat $mode:tt [$($done:tt)*] [$($all:tt)*] [$($pat:tt)*] $next:tt $($rest:tt)*) => {
125        $crate::__select!(@pat $mode [$($done)*] [$($all)*] [$($pat)* $next] $($rest)*)
126    };
127
128    (@handler $mode:tt [$($done:tt)*] [$($all:tt)*] [$($pat:tt)+] $future:expr, $cond:expr => $handler:block, $($rest:tt)*) => {
129        $crate::__select!(@branch $mode [$($done)* ([$($all)*] [$($pat)+] $future, $cond => $handler)] [$($all)* _] $($rest)*)
130    };
131    (@handler $mode:tt [$($done:tt)*] [$($all:tt)*] [$($pat:tt)+] $future:expr, $cond:expr => $handler:block $($rest:tt)*) => {
132        $crate::__select!(@branch $mode [$($done)* ([$($all)*] [$($pat)+] $future, $cond => $handler)] [$($all)* _] $($rest)*)
133    };
134    (@handler $mode:tt [$($done:tt)*] [$($all:tt)*] [$($pat:tt)+] $future:expr, $cond:expr => $handler:expr, $($rest:tt)*) => {
135        $crate::__select!(@branch $mode [$($done)* ([$($all)*] [$($pat)+] $future, $cond => { $handler })] [$($all)* _] $($rest)*)
136    };
137    (@handler $mode:tt [$($done:tt)*] [$($all:tt)*] [$($pat:tt)+] $future:expr, $cond:expr => $handler:expr) => {
138        $crate::__select!(@branch $mode [$($done)* ([$($all)*] [$($pat)+] $future, $cond => { $handler })] [$($all)* _])
139    };
140
141    (@emit [$mode:ident]
142        [$( ([$($skip:tt)*] [$($pat:tt)+] $future:expr, $cond:expr => $handler:tt) )+]
143        [$($all:tt)*]
144        [$else:expr]
145    ) => {{
146        const _: () = ::core::assert!(
147            $crate::__count!($($all)*) <= 64,
148            "`select!` takes at most 64 branches"
149        );
150        // The futures live in this block only: they are dropped before the
151        // handler runs, so the handler may use what they borrowed.
152        let __output = {
153            let mut __disabled: u64 = 0;
154            $(
155                if !$cond {
156                    __disabled |= 1u64 << $crate::__count!($($skip)*);
157                }
158            )+
159            let mut __futures = ($(
160                ::core::pin::pin!(::core::future::IntoFuture::into_future($future)),
161            )+);
162            $crate::macros::support::poll_fn(|__cx: &mut ::core::task::Context<'_>| {
163                if !$crate::macros::support::proceed(__cx) {
164                    return ::core::task::Poll::Pending;
165                }
166                let __branches: u32 = $crate::__count!($($all)*);
167                let __start: u32 = $crate::__select_start!($mode, __branches);
168                let mut __pending = false;
169                for __turn in 0..__branches {
170                    let __branch = (__start + __turn) % __branches;
171                    $(
172                        if __branch == $crate::__count!($($skip)*) {
173                            if __disabled & (1u64 << __branch) == 0 {
174                                let ($($skip,)* __future, ..) = &mut __futures;
175                                match ::core::future::Future::poll(::core::pin::Pin::as_mut(__future), __cx) {
176                                    ::core::task::Poll::Ready(__value) => {
177                                        __disabled |= 1u64 << __branch;
178                                        let __matches = match &__value {
179                                            #[allow(unused_variables)]
180                                            $crate::__select_clean!($($pat)+) => true,
181                                            #[allow(unreachable_patterns)]
182                                            _ => false,
183                                        };
184                                        if __matches {
185                                            return ::core::task::Poll::Ready($crate::__select_wrap!(
186                                                [$($skip)*]
187                                                $crate::macros::support::Branch::Taken(__value)
188                                            ));
189                                        }
190                                    }
191                                    ::core::task::Poll::Pending => __pending = true,
192                                }
193                            }
194                            continue;
195                        }
196                    )+
197                }
198                if __pending {
199                    ::core::task::Poll::Pending
200                } else {
201                    ::core::task::Poll::Ready($crate::__select_wrap!(
202                        [$($all)*]
203                        $crate::macros::support::AllDisabled
204                    ))
205                }
206            })
207            .await
208        };
209        match __output {
210            $(
211                $crate::__select_wrap!(
212                    [$($skip)*]
213                    $crate::macros::support::Branch::Taken($($pat)+)
214                ) => $handler,
215            )+
216            $crate::__select_wrap!([$($all)*] $crate::macros::support::AllDisabled) => $else,
217            // An output reaches here only through a pattern it matched.
218            #[allow(unreachable_patterns)]
219            _ => ::core::unreachable!("a `select!` output its pattern rejected"),
220        }
221    }};
222}
223
224/// The first branch a `select!` polls.
225#[doc(hidden)]
226#[macro_export]
227macro_rules! __select_start {
228    (biased, $branches:expr) => {
229        0
230    };
231    (fair, $branches:expr) => {
232        $crate::macros::support::start($branches)
233    };
234}
235
236/// Wraps a value or a pattern in one `Branch::Next` per `_`.
237#[doc(hidden)]
238#[macro_export]
239macro_rules! __select_wrap {
240    ([] $($inner:tt)+) => {
241        $($inner)+
242    };
243    ([_ $($skip:tt)*] $($inner:tt)+) => {
244        $crate::macros::support::Branch::Next($crate::__select_wrap!([$($skip)*] $($inner)+))
245    };
246}
247
248/// A branch's pattern as a test against a reference to the output: the same
249/// pattern with its `ref`, `mut` and `&` taken out, so its bindings borrow
250/// and match ergonomics apply in every edition. The state is
251/// `@ [cleaned so far] [enclosing groups]`; a group is cleaned on its own
252/// with the outer state saved on the stack, then put back in its delimiters.
253#[doc(hidden)]
254#[macro_export]
255macro_rules! __select_clean {
256    (@ [$($out:tt)*] []) => {
257        $($out)*
258    };
259    (@ [$($inner:tt)*] [(paren [$($out:tt)*] [$($rest:tt)*]) $($stack:tt)*]) => {
260        $crate::__select_clean!(@ [$($out)* ($($inner)*)] [$($stack)*] $($rest)*)
261    };
262    (@ [$($inner:tt)*] [(bracket [$($out:tt)*] [$($rest:tt)*]) $($stack:tt)*]) => {
263        $crate::__select_clean!(@ [$($out)* [$($inner)*]] [$($stack)*] $($rest)*)
264    };
265    (@ [$($inner:tt)*] [(brace [$($out:tt)*] [$($rest:tt)*]) $($stack:tt)*]) => {
266        $crate::__select_clean!(@ [$($out)* {$($inner)*}] [$($stack)*] $($rest)*)
267    };
268    (@ [$($out:tt)*] [$($stack:tt)*] ref $($rest:tt)*) => {
269        $crate::__select_clean!(@ [$($out)*] [$($stack)*] $($rest)*)
270    };
271    (@ [$($out:tt)*] [$($stack:tt)*] mut $($rest:tt)*) => {
272        $crate::__select_clean!(@ [$($out)*] [$($stack)*] $($rest)*)
273    };
274    (@ [$($out:tt)*] [$($stack:tt)*] & $($rest:tt)*) => {
275        $crate::__select_clean!(@ [$($out)*] [$($stack)*] $($rest)*)
276    };
277    (@ [$($out:tt)*] [$($stack:tt)*] && $($rest:tt)*) => {
278        $crate::__select_clean!(@ [$($out)*] [$($stack)*] $($rest)*)
279    };
280    (@ [$($out:tt)*] [$($stack:tt)*] ($($group:tt)*) $($rest:tt)*) => {
281        $crate::__select_clean!(@ [] [(paren [$($out)*] [$($rest)*]) $($stack)*] $($group)*)
282    };
283    (@ [$($out:tt)*] [$($stack:tt)*] [$($group:tt)*] $($rest:tt)*) => {
284        $crate::__select_clean!(@ [] [(bracket [$($out)*] [$($rest)*]) $($stack)*] $($group)*)
285    };
286    (@ [$($out:tt)*] [$($stack:tt)*] {$($group:tt)*} $($rest:tt)*) => {
287        $crate::__select_clean!(@ [] [(brace [$($out)*] [$($rest)*]) $($stack)*] $($group)*)
288    };
289    (@ [$($out:tt)*] [$($stack:tt)*] $next:tt $($rest:tt)*) => {
290        $crate::__select_clean!(@ [$($out)* $next] [$($stack)*] $($rest)*)
291    };
292    ($($pat:tt)+) => {
293        $crate::__select_clean!(@ [] [] $($pat)+)
294    };
295}
nuclei/macros/select.rs

nuclei/macros/
select.rs
