Skip to main content

gstd_codegen/
lib.rs

1// Copyright (C) Gear Technologies Inc.
2// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
3
4//! Provides macros for async runtime of Gear programs.
5
6use 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
20/// A global flag, determining if `handle_reply` already was generated.
21static mut HANDLE_REPLY_FLAG: Flag = Flag(false);
22
23/// A global flag, determining if `handle_signal` already was generated.
24#[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/// Macro to declare `ActorId` from hexadecimal and ss58 format.
43///
44/// # Example
45/// ```
46/// use gstd::{ActorId, actor_id};
47///
48/// # fn main() {
49/// //polkadot address
50/// let alice_1: ActorId = actor_id!("5GrwvaEF5zXb26Fz9rcQpDWS57CtERHpNehXCPcNoHGKutQY");
51/// //vara address
52/// let alice_2: ActorId = actor_id!("kGkLEU3e3XXkJp2WK4eNpVmSab5xUNL9QtmLPh8QfCL2EgotW");
53/// //hex address
54/// let alice_3: ActorId =
55///     actor_id!("0xd43593c715fdd31c61141abd04a99fd6822c8558854ccde39a5684e7a56da27d");
56///
57/// assert_eq!(alice_1, alice_2);
58/// assert_eq!(alice_2, alice_3);
59/// # }
60/// ```
61#[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)] // https://github.com/rust-lang/rust-clippy/issues/13783
192    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)] // https://github.com/rust-lang/rust-clippy/issues/13783
210    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/// Mark the main async function to be the program entry point.
232///
233/// Can be used together with [`macro@async_init`].
234///
235/// When this macro is used, it’s not possible to specify the `handle` function.
236/// If you need to specify the `handle` function explicitly, don't use this
237/// macro.
238///
239/// # Examples
240///
241/// Simple async handle function:
242///
243/// ```
244/// #[gstd::async_main]
245/// async fn main() {
246///     gstd::debug!("Hello world!");
247/// }
248/// # fn main() {}
249/// ```
250///
251/// Use `handle_reply` and `handle_signal` parameters to specify corresponding
252/// handlers. Note that custom reply and signal handlers derive their default
253/// behavior.
254///
255/// ```
256/// #[gstd::async_main(handle_reply = my_handle_reply)]
257/// async fn main() {
258///     // ...
259/// }
260///
261/// fn my_handle_reply() {
262///     // ...
263/// }
264/// # fn main() {}
265/// ```
266#[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/// Mark async function to be the program initialization method.
296///
297/// Can be used together with [`macro@async_main`].
298///
299/// The `init` function cannot be specified if this macro is used.
300/// If you need to specify the `init` function explicitly, don't use this macro.
301///
302///
303/// # Examples
304///
305/// Simple async init function:
306///
307/// ```
308/// #[gstd::async_init]
309/// async fn init() {
310///     gstd::debug!("Hello world!");
311/// }
312/// ```
313///
314/// Use `handle_reply` and `handle_signal` parameters to specify corresponding
315/// handlers. Note that custom reply and signal handlers derive their default
316/// behavior.
317///
318/// ```ignore
319/// #[gstd::async_init(handle_signal = my_handle_signal)]
320/// async fn init() {
321///     // ...
322/// }
323///
324/// fn my_handle_signal() {
325///     // ...
326/// }
327/// ```
328#[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/// Extends async methods `for_reply` and `for_reply_as` for sending
353/// methods.
354///
355/// # Usage
356///
357/// ```ignore
358/// use gcore::errors::Result;
359///
360/// #[wait_for_reply]
361/// pub fn send_bytes<T: AsRef<[u8]>>(program: ActorId, payload: T, value: u128) -> Result<MessageId> {
362///   gcore::msg::send(program.into(), payload.as_ref(), value)
363/// }
364/// ```
365///
366/// outputs:
367///
368/// ```ignore
369/// /// Same as [`send_bytes`](self::send_bytes), but the program
370/// /// will interrupt until the reply is received.
371/// ///
372/// /// Argument `reply_deposit: u64` used to provide gas for
373/// /// future reply handling (skipped if zero).
374/// ///
375/// /// # See also
376/// ///
377/// /// - [`send_bytes_for_reply_as`](self::send_bytes_for_reply_as)
378/// pub fn send_bytes_for_reply<T: AsRef<[u8]>>(
379///     program: ActorId,
380///     payload: T,
381///     value: u128,
382///     reply_deposit: u64
383/// ) -> Result<crate::msg::MessageFuture> {
384///     // Function call.
385///     let waiting_reply_to = send_bytes(program, payload, value)?;
386///
387///     // Depositing gas for future reply handling if not zero.
388///     if reply_deposit != 0 {
389///         crate::exec::reply_deposit(waiting_reply_to, reply_deposit)?;
390///     }
391///
392///     // Registering signal.
393///     crate::async_runtime::signals().register_signal(waiting_reply_to);
394///
395///     Ok(crate::msg::MessageFuture { waiting_reply_to })
396/// }
397///
398/// /// Same as [`send_bytes`](self::send_bytes), but the program
399/// /// will interrupt until the reply is received.
400/// ///
401/// /// Argument `reply_deposit: u64` used to provide gas for
402/// /// future reply handling (skipped if zero).
403/// ///
404/// /// The output should be decodable via SCALE codec.
405/// ///
406/// /// # See also
407/// ///
408/// /// - [`send_bytes_for_reply`](self::send_bytes_for_reply)
409/// /// - <https://docs.substrate.io/reference/scale-codec>
410/// pub fn send_bytes_for_reply_as<T: AsRef<[u8]>, D: crate::codec::Decode>(
411///     program: ActorId,
412///     payload: T,
413///     value: u128,
414///     reply_deposit: u64,
415/// ) -> Result<crate::msg::CodecMessageFuture<D>> {
416///     // Function call.
417///     let waiting_reply_to = send_bytes(program, payload, value)?;
418///
419///     // Depositing gas for future reply handling if not zero.
420///     if reply_deposit != 0 {
421///         crate::exec::reply_deposit(waiting_reply_to, reply_deposit)?;
422///     }
423///
424///     // Registering signal.
425///     crate::async_runtime::signals().register_signal(waiting_reply_to);
426///
427///     Ok(crate::msg::CodecMessageFuture::<D> {
428///         waiting_reply_to,
429///         _marker: Default::default(),
430///     })
431/// }
432/// ```
433#[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    // Generate functions' idents.
439    let (for_reply, for_reply_as) = (
440        utils::with_suffix(ident, "_for_reply"),
441        utils::with_suffix(ident, "_for_reply_as"),
442    );
443
444    // Generate docs.
445    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    // Generate arguments.
454    #[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    // Add `reply_deposit` argument.
459    #[cfg(not(feature = "ethexe"))]
460    inputs.push(syn::parse_quote!(reply_deposit: u64));
461
462    // Generate generics.
463    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                // Function call.
494                let waiting_reply_to = #ident #args ?;
495
496                // Depositing gas for future reply handling if not zero.
497                if reply_deposit != 0 {
498                    crate::exec::reply_deposit(waiting_reply_to, reply_deposit)?;
499                }
500
501                // Registering signal.
502                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                // Function call.
510                let waiting_reply_to = #ident #args ?;
511
512                // Depositing gas for future reply handling if not zero.
513                if reply_deposit != 0 {
514                    crate::exec::reply_deposit(waiting_reply_to, reply_deposit)?;
515                }
516
517                // Registering signal.
518                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                // Function call.
530                let waiting_reply_to = #ident #args ?;
531
532                // Registering signal.
533                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                // Function call.
541                let waiting_reply_to = #ident #args ?;
542
543                // Registering signal.
544                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/// Similar to [`macro@wait_for_reply`], but works with functions that create
553/// programs: It returns a message id with a newly created program id.
554#[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    // Generate functions' idents.
573    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    // Generate docs.
579    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    // Generate arguments.
589    #[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    // Add `reply_deposit` argument.
594    #[cfg(not(feature = "ethexe"))]
595    inputs.push(syn::parse_quote!(reply_deposit: u64));
596
597    // Generate generics.
598    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                // Function call.
617                let (waiting_reply_to, program_id) = #ident #args ?;
618
619                // Depositing gas for future reply handling if not zero.
620                if reply_deposit != 0 {
621                    crate::exec::reply_deposit(waiting_reply_to, reply_deposit)?;
622                }
623
624                // Registering signal.
625                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                // Function call.
633                let (waiting_reply_to, program_id) = #ident #args ?;
634
635                // Depositing gas for future reply handling if not zero.
636                if reply_deposit != 0 {
637                    crate::exec::reply_deposit(waiting_reply_to, reply_deposit)?;
638                }
639
640                // Registering signal.
641                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                // Function call.
653                let (waiting_reply_to, program_id) = #ident #args ?;
654
655                // Registering signal.
656                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                // Function call.
664                let (waiting_reply_to, program_id) = #ident #args ?;
665
666                // Registering signal.
667                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}