1use core::fmt::Display;
7use gprimitives::ActorId;
8use proc_macro::TokenStream;
9use proc_macro2::Ident;
10use quote::{ToTokens, quote};
11use std::{collections::BTreeSet, str::FromStr};
12use syn::{
13 Path, Token,
14 parse::{Parse, ParseStream},
15 punctuated::Punctuated,
16};
17
18mod utils;
19
20static mut HANDLE_REPLY_FLAG: Flag = Flag(false);
22
23#[cfg(not(feature = "ethexe"))]
25static mut HANDLE_SIGNAL_FLAG: Flag = Flag(false);
26
27#[cfg(feature = "ethexe")]
28static mut HANDLE_SIGNAL_FLAG: Flag = Flag(true);
29
30fn literal_to_actor_id(literal: syn::LitStr) -> syn::Result<TokenStream> {
31 let actor_id: [u8; 32] = ActorId::from_str(&literal.value())
32 .map_err(|err| syn::Error::new_spanned(literal, err))?
33 .into();
34
35 let actor_id_array = format!("{actor_id:?}")
36 .parse::<proc_macro2::TokenStream>()
37 .expect("failed to parse token stream");
38
39 Ok(quote! { gstd::ActorId::new(#actor_id_array) }.into())
40}
41
42#[proc_macro]
62pub fn actor_id(input: TokenStream) -> TokenStream {
63 literal_to_actor_id(syn::parse_macro_input!(input as syn::LitStr))
64 .unwrap_or_else(|err| err.to_compile_error().into())
65}
66
67struct Flag(bool);
68
69impl Flag {
70 fn get_and_set(&mut self) -> bool {
71 let ret = self.0;
72 self.0 = true;
73 ret
74 }
75}
76
77struct MainAttrs {
78 handle_reply: Option<Path>,
79 handle_signal: Option<Path>,
80}
81
82impl MainAttrs {
83 fn check_attrs_not_exist(&self) -> Result<(), TokenStream> {
84 let Self {
85 handle_reply,
86 handle_signal,
87 } = self;
88
89 for (path, flag) in unsafe {
90 [
91 (handle_reply, HANDLE_REPLY_FLAG.0),
92 (handle_signal, HANDLE_SIGNAL_FLAG.0),
93 ]
94 } {
95 if let (Some(path), true) = (path, flag) {
96 return Err(compile_error(path, "parameter already defined"));
97 }
98 }
99
100 Ok(())
101 }
102}
103
104impl Parse for MainAttrs {
105 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
106 let punctuated: Punctuated<MainAttr, Token![,]> = Punctuated::parse_terminated(input)?;
107 let mut attrs = MainAttrs {
108 handle_reply: None,
109 handle_signal: None,
110 };
111 let mut existing_attrs = BTreeSet::new();
112
113 for MainAttr { name, path, .. } in punctuated {
114 let name = name.to_string();
115 if existing_attrs.contains(&name) {
116 return Err(syn::Error::new_spanned(name, "parameter already defined"));
117 }
118
119 match &*name {
120 "handle_reply" => {
121 attrs.handle_reply = Some(path);
122 }
123 #[cfg(not(feature = "ethexe"))]
124 "handle_signal" => {
125 attrs.handle_signal = Some(path);
126 }
127 #[cfg(feature = "ethexe")]
128 "handle_signal" => {
129 return Err(syn::Error::new_spanned(
130 name,
131 "`handle_signal` is forbidden with `ethexe` feature on",
132 ));
133 }
134 _ => return Err(syn::Error::new_spanned(name, "unknown parameter")),
135 }
136
137 existing_attrs.insert(name);
138 }
139
140 Ok(attrs)
141 }
142}
143
144struct MainAttr {
145 name: Ident,
146 path: Path,
147}
148
149impl Parse for MainAttr {
150 fn parse(input: ParseStream<'_>) -> syn::Result<Self> {
151 let name: Ident = input.parse()?;
152 let _: Token![=] = input.parse()?;
153 let path: Path = input.parse()?;
154
155 Ok(Self { name, path })
156 }
157}
158
159fn compile_error<T: ToTokens, U: Display>(tokens: T, msg: U) -> TokenStream {
160 syn::Error::new_spanned(tokens, msg)
161 .to_compile_error()
162 .into()
163}
164
165fn check_signature(name: &str, function: &syn::ItemFn) -> Result<(), TokenStream> {
166 if function.sig.ident != name {
167 return Err(compile_error(
168 &function.sig.ident,
169 format!("function must be called `{name}`"),
170 ));
171 }
172
173 if !function.sig.inputs.is_empty() {
174 return Err(compile_error(
175 &function.sig.ident,
176 "function must have no arguments",
177 ));
178 }
179
180 if function.sig.asyncness.is_none() {
181 return Err(compile_error(
182 function.sig.fn_token,
183 "function must be async",
184 ));
185 }
186
187 Ok(())
188}
189
190fn generate_handle_reply_if_required(mut code: TokenStream, attr: Option<Path>) -> TokenStream {
191 #[allow(clippy::deref_addrof)] let reply_generated = unsafe { (*&raw mut HANDLE_REPLY_FLAG).get_and_set() };
193 if !reply_generated {
194 let handle_reply: TokenStream = quote!(
195 #[unsafe(no_mangle)]
196 extern "C" fn handle_reply() {
197 gstd::handle_reply_with_hook();
198 #attr ();
199 }
200 )
201 .into();
202 code.extend([handle_reply]);
203 }
204
205 code
206}
207
208fn generate_handle_signal_if_required(mut code: TokenStream, attr: Option<Path>) -> TokenStream {
209 #[allow(clippy::deref_addrof)] let signal_generated = unsafe { (*&raw mut HANDLE_SIGNAL_FLAG).get_and_set() };
211 if !signal_generated {
212 let handle_signal: TokenStream = quote!(
213 #[unsafe(no_mangle)]
214 extern "C" fn handle_signal() {
215 gstd::handle_signal();
216 #attr ();
217 }
218 )
219 .into();
220 code.extend([handle_signal]);
221 }
222
223 code
224}
225
226fn generate_if_required(code: TokenStream, attrs: MainAttrs) -> TokenStream {
227 let code = generate_handle_reply_if_required(code, attrs.handle_reply);
228 generate_handle_signal_if_required(code, attrs.handle_signal)
229}
230
231#[proc_macro_attribute]
267pub fn async_main(attr: TokenStream, item: TokenStream) -> TokenStream {
268 let function = syn::parse_macro_input!(item as syn::ItemFn);
269 if let Err(tokens) = check_signature("main", &function) {
270 return tokens;
271 }
272
273 let attrs = syn::parse_macro_input!(attr as MainAttrs);
274 if let Err(tokens) = attrs.check_attrs_not_exist() {
275 return tokens;
276 }
277
278 let body = &function.block;
279 let code: TokenStream = quote!(
280
281 fn __main_safe() {
282 gstd::message_loop(async #body);
283 }
284
285 #[unsafe(no_mangle)]
286 extern "C" fn handle() {
287 __main_safe();
288 }
289 )
290 .into();
291
292 generate_if_required(code, attrs)
293}
294
295#[proc_macro_attribute]
329pub fn async_init(attr: TokenStream, item: TokenStream) -> TokenStream {
330 let function = syn::parse_macro_input!(item as syn::ItemFn);
331 if let Err(tokens) = check_signature("init", &function) {
332 return tokens;
333 }
334
335 let attrs = syn::parse_macro_input!(attr as MainAttrs);
336 if let Err(tokens) = attrs.check_attrs_not_exist() {
337 return tokens;
338 }
339
340 let body = &function.block;
341 let code: TokenStream = quote!(
342 #[unsafe(no_mangle)]
343 extern "C" fn init() {
344 gstd::message_loop(async #body);
345 }
346 )
347 .into();
348
349 generate_if_required(code, attrs)
350}
351
352#[proc_macro_attribute]
434pub fn wait_for_reply(attr: TokenStream, item: TokenStream) -> TokenStream {
435 let function = syn::parse_macro_input!(item as syn::ItemFn);
436 let ident = &function.sig.ident;
437
438 let (for_reply, for_reply_as) = (
440 utils::with_suffix(ident, "_for_reply"),
441 utils::with_suffix(ident, "_for_reply_as"),
442 );
443
444 let style = if !attr.is_empty() {
446 utils::DocumentationStyle::Method
447 } else {
448 utils::DocumentationStyle::Function
449 };
450
451 let (for_reply_docs, for_reply_as_docs) = utils::wait_for_reply_docs(ident.to_string(), style);
452
453 #[cfg_attr(feature = "ethexe", allow(unused_mut))]
455 let (mut inputs, variadic) = (function.sig.inputs.clone(), function.sig.variadic.clone());
456 let args = utils::get_args(&inputs);
457
458 #[cfg(not(feature = "ethexe"))]
460 inputs.push(syn::parse_quote!(reply_deposit: u64));
461
462 let decodable_ty = utils::ident("D");
464 let decodable_traits = vec![syn::parse_quote!(crate::codec::Decode)];
465 let (for_reply_generics, for_reply_as_generics) = (
466 function.sig.generics.clone(),
467 utils::append_generic(
468 function.sig.generics.clone(),
469 decodable_ty,
470 decodable_traits,
471 ),
472 );
473
474 let ident = if !attr.is_empty() {
475 assert_eq!(
476 attr.to_string(),
477 "self",
478 "Proc macro attribute should be used only to specify self source of the function"
479 );
480
481 quote! { self.#ident }
482 } else {
483 quote! { #ident }
484 };
485
486 match () {
487 #[cfg(not(feature = "ethexe"))]
488 () => quote! {
489 #function
490
491 #[doc = #for_reply_docs]
492 pub fn #for_reply #for_reply_generics ( #inputs #variadic ) -> Result<crate::msg::MessageFuture> {
493 let waiting_reply_to = #ident #args ?;
495
496 if reply_deposit != 0 {
498 crate::exec::reply_deposit(waiting_reply_to, reply_deposit)?;
499 }
500
501 crate::async_runtime::signals().register_signal(waiting_reply_to);
503
504 Ok(crate::msg::MessageFuture { waiting_reply_to, reply_deposit })
505 }
506
507 #[doc = #for_reply_as_docs]
508 pub fn #for_reply_as #for_reply_as_generics ( #inputs #variadic ) -> Result<crate::msg::CodecMessageFuture<D>> {
509 let waiting_reply_to = #ident #args ?;
511
512 if reply_deposit != 0 {
514 crate::exec::reply_deposit(waiting_reply_to, reply_deposit)?;
515 }
516
517 crate::async_runtime::signals().register_signal(waiting_reply_to);
519
520 Ok(crate::msg::CodecMessageFuture::<D> { waiting_reply_to, reply_deposit, _marker: Default::default() })
521 }
522 },
523 #[cfg(feature = "ethexe")]
524 () => quote! {
525 #function
526
527 #[doc = #for_reply_docs]
528 pub fn #for_reply #for_reply_generics ( #inputs #variadic ) -> Result<crate::msg::MessageFuture> {
529 let waiting_reply_to = #ident #args ?;
531
532 crate::async_runtime::signals().register_signal(waiting_reply_to);
534
535 Ok(crate::msg::MessageFuture { waiting_reply_to, reply_deposit: 0 })
536 }
537
538 #[doc = #for_reply_as_docs]
539 pub fn #for_reply_as #for_reply_as_generics ( #inputs #variadic ) -> Result<crate::msg::CodecMessageFuture<D>> {
540 let waiting_reply_to = #ident #args ?;
542
543 crate::async_runtime::signals().register_signal(waiting_reply_to);
545
546 Ok(crate::msg::CodecMessageFuture::<D> { waiting_reply_to, reply_deposit: 0, _marker: Default::default() })
547 }
548 },
549 }.into()
550}
551
552#[proc_macro_attribute]
555pub fn wait_create_program_for_reply(attr: TokenStream, item: TokenStream) -> TokenStream {
556 let function = syn::parse_macro_input!(item as syn::ItemFn);
557
558 let function_ident = &function.sig.ident;
559
560 let ident = if !attr.is_empty() {
561 assert_eq!(
562 attr.to_string(),
563 "Self",
564 "Proc macro attribute should be used only to specify Self source of the function"
565 );
566
567 quote! { Self::#function_ident }
568 } else {
569 quote! { #function_ident }
570 };
571
572 let (for_reply, for_reply_as) = (
574 utils::with_suffix(&function.sig.ident, "_for_reply"),
575 utils::with_suffix(&function.sig.ident, "_for_reply_as"),
576 );
577
578 let style = if !attr.is_empty() {
580 utils::DocumentationStyle::Method
581 } else {
582 utils::DocumentationStyle::Function
583 };
584
585 let (for_reply_docs, for_reply_as_docs) =
586 utils::wait_for_reply_docs(function_ident.to_string(), style);
587
588 #[cfg_attr(feature = "ethexe", allow(unused_mut))]
590 let (mut inputs, variadic) = (function.sig.inputs.clone(), function.sig.variadic.clone());
591 let args = utils::get_args(&inputs);
592
593 #[cfg(not(feature = "ethexe"))]
595 inputs.push(syn::parse_quote!(reply_deposit: u64));
596
597 let decodable_ty = utils::ident("D");
599 let decodable_traits = vec![syn::parse_quote!(crate::codec::Decode)];
600 let (for_reply_generics, for_reply_as_generics) = (
601 function.sig.generics.clone(),
602 utils::append_generic(
603 function.sig.generics.clone(),
604 decodable_ty,
605 decodable_traits,
606 ),
607 );
608
609 match () {
610 #[cfg(not(feature = "ethexe"))]
611 () => quote! {
612 #function
613
614 #[doc = #for_reply_docs]
615 pub fn #for_reply #for_reply_generics ( #inputs #variadic ) -> Result<crate::msg::CreateProgramFuture> {
616 let (waiting_reply_to, program_id) = #ident #args ?;
618
619 if reply_deposit != 0 {
621 crate::exec::reply_deposit(waiting_reply_to, reply_deposit)?;
622 }
623
624 crate::async_runtime::signals().register_signal(waiting_reply_to);
626
627 Ok(crate::msg::CreateProgramFuture { waiting_reply_to, program_id, reply_deposit })
628 }
629
630 #[doc = #for_reply_as_docs]
631 pub fn #for_reply_as #for_reply_as_generics ( #inputs #variadic ) -> Result<crate::msg::CodecCreateProgramFuture<D>> {
632 let (waiting_reply_to, program_id) = #ident #args ?;
634
635 if reply_deposit != 0 {
637 crate::exec::reply_deposit(waiting_reply_to, reply_deposit)?;
638 }
639
640 crate::async_runtime::signals().register_signal(waiting_reply_to);
642
643 Ok(crate::msg::CodecCreateProgramFuture::<D> { waiting_reply_to, program_id, reply_deposit, _marker: Default::default() })
644 }
645 },
646 #[cfg(feature = "ethexe")]
647 () => quote! {
648 #function
649
650 #[doc = #for_reply_docs]
651 pub fn #for_reply #for_reply_generics ( #inputs #variadic ) -> Result<crate::msg::CreateProgramFuture> {
652 let (waiting_reply_to, program_id) = #ident #args ?;
654
655 crate::async_runtime::signals().register_signal(waiting_reply_to);
657
658 Ok(crate::msg::CreateProgramFuture { waiting_reply_to, program_id, reply_deposit: 0 })
659 }
660
661 #[doc = #for_reply_as_docs]
662 pub fn #for_reply_as #for_reply_as_generics ( #inputs #variadic ) -> Result<crate::msg::CodecCreateProgramFuture<D>> {
663 let (waiting_reply_to, program_id) = #ident #args ?;
665
666 crate::async_runtime::signals().register_signal(waiting_reply_to);
668
669 Ok(crate::msg::CodecCreateProgramFuture::<D> { waiting_reply_to, program_id, reply_deposit: 0, _marker: Default::default() })
670 }
671 },
672 }.into()
673}
674
675#[cfg(test)]
676mod tests {
677 #[test]
678 fn ui() {
679 let t = trybuild::TestCases::new();
680
681 #[cfg(not(feature = "ethexe"))]
682 {
683 t.pass("tests/ui/async_init_works.rs");
684 t.pass("tests/ui/async_main_works.rs");
685 t.compile_fail("tests/ui/signal_double_definition_not_work.rs");
686 t.compile_fail("tests/ui/reply_double_definition_not_work.rs");
687 }
688
689 #[cfg(feature = "ethexe")]
690 {
691 t.compile_fail("tests/ui/signal_doesnt_work_with_ethexe.rs");
692 }
693 }
694}