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