1extern crate proc_macro;
2
3use proc_macro::TokenStream;
4
5use quote::format_ident;
6use quote::quote;
7
8use syn::parse::Parse;
9use syn::parse::ParseStream;
10use syn::parse::Result;
11use syn::*;
12
13enum Item {
14 Struct(ItemStruct),
15 Impl(ItemImpl),
16}
17
18impl Parse for Item {
19 fn parse(input: ParseStream) -> Result<Self> {
20 let lookahead = input.lookahead1();
21
22 Ok(if lookahead.peek(Token![impl]) {
23 let item: ItemImpl = input.parse()?;
24 Item::Impl(item)
25 } else {
26 let item: ItemStruct = input.parse()?;
27 Item::Struct(item)
28 })
29 }
30}
31
32#[proc_macro_attribute]
33pub fn actor(_args: TokenStream, input: TokenStream) -> TokenStream {
34 let input = syn::parse_macro_input!(input as Item);
35 TokenStream::from(match input {
36 Item::Struct(x) => actor_struct(x),
37 Item::Impl(x) => actor_impl(x),
38 })
39}
40
41#[proc_macro_attribute]
42pub fn handler(_: TokenStream, input: TokenStream) -> TokenStream {
43 input
44}
45
46fn actor_struct(item_struct: ItemStruct) -> proc_macro2::TokenStream {
47 let ItemStruct {
48 attrs,
49 ident,
50 generics,
51 fields,
52 semi_token,
53 ..
54 } = item_struct;
55
56 let actor_mod = format_ident!("__Actor{}", ident);
57
58 quote! {
59 #[doc(hidden)]
60 #[allow(non_snake_case)]
61 mod #actor_mod {
62 use super::*;
63 #(#attrs)*
64 pub struct #ident #generics #fields #semi_token
65
66 impl ::xtra::Actor for #ident {}
67 }
68 }
69}
70
71fn actor_impl(item_impl: ItemImpl) -> proc_macro2::TokenStream {
72 let name = get_name(&item_impl);
73 let mod_name = format_ident!("__Actor{}", name);
74
75 let non_handler_methods = item_impl
76 .items
77 .iter()
78 .filter(handler_function)
79 .collect::<Vec<_>>();
80
81 let handler_methods = item_impl
82 .items
83 .iter()
84 .filter(|x| !handler_function(x))
85 .collect::<Vec<_>>();
86
87 let method_new = non_handler_methods.iter().find(|x| {
88 if let ImplItem::Method(x) = x {
89 if x.sig.ident.to_string() == "new" {
90 return true;
91 }
92 }
93 false
94 });
95
96 let method_new = match method_new {
97 Some(ImplItem::Method(x)) => x,
98 _ => panic!("Actor must have a `new` method"),
99 };
100
101 let args = method_new
102 .sig
103 .inputs
104 .iter()
105 .filter(|x| matches!(x, FnArg::Typed(_)))
106 .collect::<Vec<_>>();
107
108 let arglist = args.iter().map(|x| match x {
109 FnArg::Typed(x) => x.pat.clone(),
110 _ => unreachable!(),
111 });
112
113 let arglist_clone = arglist.clone().map(|x| quote!(#x.clone()));
114
115 let actor_creator = quote! {
116 pub fn new<S: ::xtra::spawn::Spawner>(spawner: &mut S, #(#args),*) -> Self {
117 use ::xtra::Actor;
118
119 Self {
120 addr: #mod_name::#name::new(#(#arglist),*).create(None).spawn(spawner),
121 }
122 }
123
124 pub fn cluster<S: ::xtra::spawn::Spawner>(
125 spawner: &mut S,
126 cluster_size: usize,
127 #(#args),*
128 ) -> (::xtra::Context<#mod_name::#name>, Self) {
129 use ::xtra::Actor;
130 let (addr, mut ctx) = ::xtra::Context::<#mod_name::#name>::new(None);
131 for _ in 0..cluster_size {
132 spawner.spawn(ctx.attach(#mod_name::#name::new(#(#arglist_clone),*)));
133 }
134
135 (ctx, Self { addr })
136 }
137 };
138
139 let message_structs = generate_message_structs(&name, handler_methods.clone());
140 let handlers = generate_handlers(&name, handler_methods.clone());
141 let api_methods = generate_api_methods(&name, handler_methods.clone());
142
143 quote! {
144 impl #mod_name::#name {
145 #(#non_handler_methods)*
146
147 #(#handler_methods)*
148 }
149
150 #[derive(Clone)]
151 pub struct #name {
152 addr: ::xtra::Address<#mod_name::#name>,
153 }
154
155 #message_structs
156 #handlers
157
158 impl #name {
159 #actor_creator
160
161 #(#api_methods)*
162 }
163 }
164}
165
166fn generate_api_methods(
167 actor_name: &Ident,
168 items: Vec<&ImplItem>,
169) -> Vec<proc_macro2::TokenStream> {
170 items
171 .iter()
172 .map(|x| generate_api_method(actor_name, x))
173 .collect::<Vec<_>>()
174}
175
176fn generate_api_method(actor_name: &Ident, item: &ImplItem) -> proc_macro2::TokenStream {
177 let item = match item {
178 ImplItem::Method(x) => x,
179 _ => panic!("tried to generate struct for non fn handler"),
180 };
181
182 let fn_name = item.sig.ident.clone();
183 let msg_name = format_ident!("__{}__{}", actor_name, fn_name);
184
185 let args = item
186 .sig
187 .inputs
188 .iter()
189 .filter(|x| matches!(x, FnArg::Typed(_)))
190 .collect::<Vec<_>>();
191
192 let ret_type = &item.sig.output;
193
194 let arglist = args.iter().map(|x| match x {
195 FnArg::Typed(x) => x.pat.clone(),
196 _ => unreachable!(),
197 });
198
199 let arglist_clone = arglist.clone();
200
201 let no_wait_fn_name = format_ident!("{}_no_wait", &fn_name);
202
203 quote! {
204 pub async fn #fn_name(&self, #(#args),*) #ret_type {
205 self.addr.send(#msg_name {
206 #(#arglist),*
207 }).await.expect("Actor has died.")
208 }
209
210 pub fn #no_wait_fn_name(&self, #(#args),*) {
211 self.addr.do_send(#msg_name {
212 #(#arglist_clone),*
213 }).expect("Actor has died.");
214 }
215 }
216}
217
218fn generate_message_structs(actor_name: &Ident, items: Vec<&ImplItem>) -> proc_macro2::TokenStream {
219 let msg_structs = items
220 .iter()
221 .map(|x| generate_msg_struct(actor_name, x))
222 .collect::<Vec<_>>();
223
224 quote! {
225 #(#msg_structs)*
226 }
227}
228
229fn generate_msg_struct(actor_name: &Ident, item: &ImplItem) -> proc_macro2::TokenStream {
230 let item = match item {
231 ImplItem::Method(x) => x,
232 _ => panic!("tried to generate struct for non fn handler"),
233 };
234
235 let fn_name = item.sig.ident.clone();
236 let msg_name = format_ident!("__{}__{}", actor_name, fn_name);
237
238 let args = item
239 .sig
240 .inputs
241 .iter()
242 .filter(|x| matches!(x, FnArg::Typed(_)))
243 .collect::<Vec<_>>();
244
245 let xtra_msg_impl = match &item.sig.output {
246 ReturnType::Default => quote! {
247 impl ::xtra::Message for #msg_name {
248 type Result = ();
249 }
250 },
251 ReturnType::Type(_, t) => quote! {
252 impl ::xtra::Message for #msg_name {
253 type Result = #t;
254 }
255 },
256 };
257
258 quote! {
259 #[doc(hidden)]
260 #[allow(non_camel_case_types)]
261 struct #msg_name {
262 #(#args),*
263 }
264
265 #[doc(hidden)]
266 #[allow(non_snake_case)]
267 #xtra_msg_impl
268 }
269}
270
271fn generate_handlers(actor_name: &Ident, items: Vec<&ImplItem>) -> proc_macro2::TokenStream {
272 let handlers = items
273 .iter()
274 .map(|x| generate_handler(actor_name, x))
275 .collect::<Vec<_>>();
276
277 quote! {
278 #(#handlers)*
279 }
280}
281
282fn generate_handler(actor_name: &Ident, item: &ImplItem) -> proc_macro2::TokenStream {
283 let item = match item {
284 ImplItem::Method(x) => x,
285 _ => panic!("tried to generate struct for non fn handler"),
286 };
287
288 let fn_name = item.sig.ident.clone();
289 let msg_name = format_ident!("__{}__{}", actor_name, fn_name);
290 let mod_name = format_ident!("__Actor{}", actor_name);
291
292 let args = item
293 .sig
294 .inputs
295 .iter()
296 .filter(|x| matches!(x, FnArg::Typed(_)))
297 .collect::<Vec<_>>();
298
299 let arglist = args
300 .iter()
301 .map(|x| match x {
302 FnArg::Typed(x) => x.pat.clone(),
303 _ => unreachable!(),
304 })
305 .map(|x| {
306 quote! {
307 args.#x
308 }
309 });
310
311 quote! {
312 #[async_trait]
313 impl ::xtra::Handler<#msg_name> for #mod_name::#actor_name {
314 async fn handle(&mut self, args: #msg_name, _: &mut ::xtra::Context<Self>)
315 -> <#msg_name as ::xtra::Message>::Result
316 {
317 self.#fn_name(#(#arglist),*).await
318 }
319 }
320 }
321}
322
323fn get_name(block: &ItemImpl) -> proc_macro2::Ident {
324 let self_ty_path = match &*block.self_ty {
325 Type::Path(path) => &path.path,
326 _ => panic!(),
327 };
328
329 self_ty_path.segments.last().unwrap().ident.clone()
330}
331
332fn handler_function(x: &&ImplItem) -> bool {
333 if let ImplItem::Method(x) = x {
334 let ident = x
335 .attrs
336 .iter()
337 .map(|x| x.path.segments.last().unwrap().ident.to_string())
338 .collect::<Vec<_>>();
339
340 if ident.contains(&"handler".to_string()) {
341 return false;
342 }
343 }
344
345 true
346}