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
clean_pattern_macro(input: TokenStream) -> TokenStream45 pub(crate) fn clean_pattern_macro(input: TokenStream) -> TokenStream {
46 // If this isn't a pattern, we return the token stream as-is. The select!
47 // macro is using it in a location requiring a pattern, so an error will be
48 // emitted there.
49 let mut input: syn::Pat = match syn::parse(input.clone()) {
50 Ok(it) => it,
51 Err(_) => return input,
52 };
53
54 clean_pattern(&mut input);
55 quote::ToTokens::into_token_stream(input).into()
56 }
57
58 // Removes any occurrences of ref or mut in the provided pattern.
clean_pattern(pat: &mut syn::Pat)59 fn clean_pattern(pat: &mut syn::Pat) {
60 match pat {
61 syn::Pat::Box(_box) => {}
62 syn::Pat::Lit(_literal) => {}
63 syn::Pat::Macro(_macro) => {}
64 syn::Pat::Path(_path) => {}
65 syn::Pat::Range(_range) => {}
66 syn::Pat::Rest(_rest) => {}
67 syn::Pat::Verbatim(_tokens) => {}
68 syn::Pat::Wild(_underscore) => {}
69 syn::Pat::Ident(ident) => {
70 ident.by_ref = None;
71 ident.mutability = None;
72 if let Some((_at, pat)) = &mut ident.subpat {
73 clean_pattern(&mut *pat);
74 }
75 }
76 syn::Pat::Or(or) => {
77 for case in or.cases.iter_mut() {
78 clean_pattern(case);
79 }
80 }
81 syn::Pat::Slice(slice) => {
82 for elem in slice.elems.iter_mut() {
83 clean_pattern(elem);
84 }
85 }
86 syn::Pat::Struct(struct_pat) => {
87 for field in struct_pat.fields.iter_mut() {
88 clean_pattern(&mut field.pat);
89 }
90 }
91 syn::Pat::Tuple(tuple) => {
92 for elem in tuple.elems.iter_mut() {
93 clean_pattern(elem);
94 }
95 }
96 syn::Pat::TupleStruct(tuple) => {
97 for elem in tuple.pat.elems.iter_mut() {
98 clean_pattern(elem);
99 }
100 }
101 syn::Pat::Reference(reference) => {
102 reference.mutability = None;
103 clean_pattern(&mut *reference.pat);
104 }
105 syn::Pat::Type(type_pat) => {
106 clean_pattern(&mut *type_pat.pat);
107 }
108 _ => {}
109 }
110 }
111