diff options
Diffstat (limited to 'tokio-macros/src/select.rs')
-rw-r--r-- | tokio-macros/src/select.rs | 43 |
1 files changed, 43 insertions, 0 deletions
diff --git a/tokio-macros/src/select.rs b/tokio-macros/src/select.rs new file mode 100644 index 00000000..ddb2e6a0 --- /dev/null +++ b/tokio-macros/src/select.rs @@ -0,0 +1,43 @@ +use proc_macro::{TokenStream, TokenTree}; +use proc_macro2::Span; +use quote::quote; +use syn::Ident; + +pub(crate) fn declare_output_enum(input: TokenStream) -> TokenStream { + // passed in is: `(_ _ _)` with one `_` per branch + let branches = match input.into_iter().next() { + Some(TokenTree::Group(group)) => group.stream().into_iter().count(), + _ => panic!("unexpected macro input"), + }; + + let variants = (0..branches) + .map(|num| Ident::new(&format!("_{}", num), Span::call_site())) + .collect::<Vec<_>>(); + + // Use a bitfield to track which futures completed + let mask = Ident::new( + if branches <= 8 { + "u8" + } else if branches <= 16 { + "u16" + } else if branches <= 32 { + "u32" + } else if branches <= 64 { + "u64" + } else { + panic!("up to 64 branches supported"); + }, + Span::call_site(), + ); + + TokenStream::from(quote! { + pub(super) enum Out<#( #variants ),*> { + #( #variants(#variants), )* + // Include a `Disabled` variant signifying that all select branches + // failed to resolve. + Disabled, + } + + pub(super) type Mask = #mask; + }) +} |