1 use proc_macro::{TokenStream, TokenTree};
2 use proc_macro2::Span;
3 use quote::quote;
4 use syn::Ident;
5
declare_output_enum(input: TokenStream) -> TokenStream6 pub(crate) fn declare_output_enum(input: TokenStream) -> TokenStream {
7 // passed in is: `(_ _ _)` with one `_` per branch
8 let branches = match input.into_iter().next() {
9 Some(TokenTree::Group(group)) => group.stream().into_iter().count(),
10 _ => panic!("unexpected macro input"),
11 };
12
13 let variants = (0..branches)
14 .map(|num| Ident::new(&format!("_{}", num), Span::call_site()))
15 .collect::<Vec<_>>();
16
17 // Use a bitfield to track which futures completed
18 let mask = Ident::new(
19 if branches <= 8 {
20 "u8"
21 } else if branches <= 16 {
22 "u16"
23 } else if branches <= 32 {
24 "u32"
25 } else if branches <= 64 {
26 "u64"
27 } else {
28 panic!("up to 64 branches supported");
29 },
30 Span::call_site(),
31 );
32
33 TokenStream::from(quote! {
34 pub(super) enum Out<#( #variants ),*> {
35 #( #variants(#variants), )*
36 // Include a `Disabled` variant signifying that all select branches
37 // failed to resolve.
38 Disabled,
39 }
40
41 pub(super) type Mask = #mask;
42 })
43 }
44