Skip to main content

behavior_macros/
lib.rs

1//! `behavior-macros` — proc-macros for the behavior algebra.
2//!
3//! `workers!` compiles a mixed fleet declaration into the erasure-free sum
4//! a `Supervising` fleet requires (design: actorpass docs, surface talk #2):
5//! `(count, Type, build_fn)` per worker kind → a `Crew` enum with a
6//! delegated `Behavior` impl, a per-variant range `crew_build`, and the
7//! total count. `Crew` is a TYPE — every worker stays its own actor.
8//!
9//! v1 scope: every worker kind shares the SAME protocol (`Event`, `Sends`,
10//! `Done`, `Error`, and `Birth` — taken from the first kind). Mixed
11//! protocols need the hand-written sum (the `CrewMsg` widening is a
12//! deliberate, documented step — not this macro's job yet).
13
14use 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
20/// One `(count, Type, build_fn)` worker-kind spec.
21struct 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/// `workers![(4, WorkerA, build_a), (2, WorkerB, build_b)]` → a block
59/// declaring the `Crew` sum and yielding `(total, crew_build)` for
60/// `Supervising`'s fleet. Slots are contiguous per variant (slot = nonce;
61/// rest-for-one's birth order is the declaration order).
62#[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            /// The macro-generated mixed-fleet sum (see `workers!`).
110            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                type Effect = <#first_ty as ::behavior::Behavior>::Effect;
123                type Done = <#first_ty as ::behavior::Behavior>::Done;
124
125                async fn init(&mut self) -> ::core::result::Result<Self::Effect, Self::Error> {
126                    match self {
127                        #(#init_arms),*
128                    }
129                }
130
131                async fn step(
132                    &mut self,
133                    ev: Self::Event,
134                ) -> ::core::result::Result<Self::Effect, Self::Error> {
135                    match self {
136                        #(#step_arms),*
137                    }
138                }
139
140            }
141
142            fn crew_build(i: usize) -> Crew {
143                match i {
144                    #(#build_arms,)*
145                    _ => unreachable!("workers!: fleet index out of range — driver/behavior desync"),
146                }
147            }
148
149            (#total, crew_build as fn(usize) -> Crew)
150        }
151    };
152    out.into()
153}