1use proc_macro::TokenStream;
15use quote::{format_ident, quote};
16use syn::parse::{Parse, ParseStream};
17use syn::punctuated::Punctuated;
18use syn::{Error, Expr, LitInt, Result, Token, Type, parse_macro_input};
19
20struct Spec {
22 count: LitInt,
23 ty: Type,
24 build: Expr,
25}
26
27impl Parse for Spec {
28 fn parse(input: ParseStream) -> Result<Self> {
29 let content;
30 syn::parenthesized!(content in input);
31 let count: Expr = content.parse()?;
32 let Expr::Lit(syn::ExprLit {
33 lit: syn::Lit::Int(count),
34 ..
35 }) = count
36 else {
37 return Err(Error::new_spanned(
38 count,
39 "worker count must be a usize literal (ranges are computed at expansion)",
40 ));
41 };
42 content.parse::<Token![,]>()?;
43 let ty: Type = content.parse()?;
44 content.parse::<Token![,]>()?;
45 let build: Expr = content.parse()?;
46 Ok(Spec { count, ty, build })
47 }
48}
49
50struct Specs(Punctuated<Spec, Token![,]>);
51
52impl Parse for Specs {
53 fn parse(input: ParseStream) -> Result<Self> {
54 Ok(Specs(Punctuated::parse_terminated(input)?))
55 }
56}
57
58#[proc_macro]
63pub fn workers(input: TokenStream) -> TokenStream {
64 let Specs(specs) = parse_macro_input!(input as Specs);
65 let specs: Vec<Spec> = specs.into_iter().collect();
66 if specs.is_empty() {
67 return Error::new(
68 proc_macro2::Span::call_site(),
69 "workers! needs at least one (count, Type, build_fn) spec",
70 )
71 .to_compile_error()
72 .into();
73 }
74
75 let first_ty = &specs[0].ty;
76 let variants: Vec<_> = specs
77 .iter()
78 .enumerate()
79 .map(|(i, _)| format_ident!("V{i}"))
80 .collect();
81 let variant_defs = specs.iter().zip(&variants).map(|(s, v)| {
82 let ty = &s.ty;
83 quote! { #v(#ty) }
84 });
85
86 let mut start = 0_usize;
87 let mut build_arms = Vec::new();
88 for (s, v) in specs.iter().zip(&variants) {
89 let n: usize = match s.count.base10_parse() {
90 Ok(n) => n,
91 Err(e) => return e.to_compile_error().into(),
92 };
93 let end = start + n;
94 let build = &s.build;
95 build_arms.push(quote! { #start..#end => Crew::#v((#build)(i)) });
96 start = end;
97 }
98 let total = start;
99
100 let step_arms = variants
101 .iter()
102 .map(|v| quote! { Crew::#v(b) => b.step(ev).await });
103 let init_arms = variants
104 .iter()
105 .map(|v| quote! { Crew::#v(b) => b.init().await });
106
107 let out = quote! {
108 {
109 enum Crew {
111 #(#variant_defs),*
112 }
113
114 impl ::behavior::Behavior for Crew {
115 type Addr = <#first_ty as ::behavior::Behavior>::Addr;
116 type Msg = <#first_ty as ::behavior::Behavior>::Msg;
117 type Event = <#first_ty as ::behavior::Behavior>::Event;
118 type Sends = <#first_ty as ::behavior::Behavior>::Sends;
119 type Ph = ::behavior::Never;
120 type Error = <#first_ty as ::behavior::Behavior>::Error;
121 type Birth = <#first_ty as ::behavior::Behavior>::Birth;
122
123 async fn init(&mut self) -> ::core::result::Result<
124 ::behavior::Actions<Self::Addr, Self::Ph, Self::Sends, Self::Birth>,
125 Self::Error,
126 > {
127 match self {
128 #(#init_arms),*
129 }
130 }
131
132 async fn step(
133 &mut self,
134 ev: Self::Event,
135 ) -> ::core::result::Result<
136 ::behavior::Actions<Self::Addr, Self::Ph, Self::Sends, Self::Birth>,
137 Self::Error,
138 > {
139 match self {
140 #(#step_arms),*
141 }
142 }
143
144 }
145
146 fn crew_build(i: usize) -> Crew {
147 match i {
148 #(#build_arms,)*
149 _ => unreachable!("workers!: fleet index out of range — driver/behavior desync"),
150 }
151 }
152
153 (#total, crew_build as fn(usize) -> Crew)
154 }
155 };
156 out.into()
157}