balter_macros/
lib.rs

1use proc_macro::TokenStream;
2use proc_macro2::{Span, TokenStream as TokenStream2};
3use quote::quote;
4use syn::{Ident, ItemFn};
5
6/// Proc macro to denote a Transaction
7///
8/// NOTE: Currently this macro only works on functions with a `Result<T, E>` return value. This is a
9/// restriction which will be lifted soon.
10///
11/// # Example
12/// ```ignore
13/// use balter::prelude::*;
14///
15/// #[transaction]
16/// fn my_transaction(arg_1: u32, arg_2: &str) -> Result<String, MyError> {
17///     ...
18/// }
19/// ```
20#[proc_macro_attribute]
21pub fn transaction(attr: TokenStream, item: TokenStream) -> TokenStream {
22    transaction_internal(attr, item).into()
23}
24
25fn transaction_internal(_attr: TokenStream, item: TokenStream) -> TokenStream2 {
26    let input = syn::parse::<ItemFn>(item).unwrap();
27
28    let ItemFn {
29        attrs,
30        vis,
31        sig,
32        block,
33    } = input;
34    let stmts = &block.stmts;
35
36    let ident = &sig.ident;
37    quote! {
38        #(#attrs)* #vis #sig {
39            ::balter::transaction::transaction_hook(::balter::core::generate_labels!(#ident), async move {
40                #(#stmts)*
41            }).await
42        }
43    }
44}
45
46/// Proc macro to denote a Scenario
47///
48/// NOTE: Currently this macro only works on functions which take no arguments and with no return value.
49/// (void functions). This is a restriction which will be lifted soon.
50///
51/// See the `Scenario` struct for more information on the methods this macro provides on functions.
52///
53/// # Example
54/// ```ignore
55/// use balter::prelude::*;
56///
57/// #[scenario]
58/// fn my_scenario() {
59/// }
60/// ```
61#[proc_macro_attribute]
62pub fn scenario(attr: TokenStream, item: TokenStream) -> TokenStream {
63    scenario_internal(attr, item, false).into()
64}
65
66/// Proc macro to denote a Scenario
67///
68/// NOTE: Currently this macro only works on functions which take no arguments and with no return value.
69/// (void functions). This is a restriction which will be lifted soon.
70///
71/// See the `Scenario` struct for more information on the methods this macro provides on functions.
72///
73/// # Example
74/// ```ignore
75/// use balter::prelude::*;
76///
77/// #[scenario]
78/// fn my_scenario() {
79/// }
80/// ```
81#[proc_macro_attribute]
82pub fn scenario_linkme(attr: TokenStream, item: TokenStream) -> TokenStream {
83    scenario_internal(attr, item, true).into()
84}
85
86fn scenario_internal(_attr: TokenStream, item: TokenStream, linkme: bool) -> TokenStream2 {
87    let input = syn::parse::<ItemFn>(item).expect("Macro only works on fn() items");
88
89    let ItemFn {
90        attrs,
91        vis,
92        sig,
93        block,
94    } = input;
95    let stmts = &block.stmts;
96
97    let new_name = Ident::new(&format!("__balter_{}", sig.ident), Span::call_site());
98    let mut new_sig = sig.clone();
99    new_sig.ident = new_name.clone();
100
101    let mut scen_sig = sig.clone();
102    let scen_name = sig.ident.clone();
103    scen_sig.asyncness = None;
104    scen_sig.output = syn::parse(
105        quote! {
106            -> impl ::balter::scenario::ConfigurableScenario<::balter::prelude::RunStatistics>
107        }
108        .into(),
109    )
110    .expect("Scenario signature is invalid");
111
112    let res = quote! {
113        #(#attrs)* #vis #scen_sig {
114            ::balter::scenario::Scenario::new(stringify!(#scen_name), #new_name)
115        }
116
117        #(#attrs)* #vis #new_sig {
118            #(#stmts)*
119        }
120    };
121
122    if linkme {
123        let mut linkme_sig = sig.clone();
124        let linkme_name = Ident::new(&format!("__balter_distr_{}", sig.ident), Span::call_site());
125        linkme_sig.ident = linkme_name.clone();
126        linkme_sig.asyncness = None;
127        linkme_sig.output = syn::parse(
128            quote! {
129                -> ::core::pin::Pin<Box<dyn ::balter::prelude::DistributedScenario<Output=::balter::prelude::RunStatistics>>>
130            }
131            .into(),
132        )
133        .expect("Scenario signature is invalid");
134
135        let static_name = Ident::new(
136            &format!("__BALTER_{}", sig.ident.to_string().to_ascii_uppercase()),
137            Span::call_site(),
138        );
139
140        let mut linkme = quote! {
141            #[::balter::runtime::distributed_slice(::balter::runtime::BALTER_SCENARIOS)]
142            static #static_name: (&'static str, fn() -> ::core::pin::Pin<Box<dyn ::balter::prelude::DistributedScenario<Output=::balter::prelude::RunStatistics>>>) = (stringify!(#scen_name), #linkme_name);
143
144            // TODO: This definition can almost certainly merge with the #scen_sig definition
145            #(#attrs)* #vis #linkme_sig {
146                Box::pin(::balter::scenario::Scenario::new(stringify!(#scen_name), #new_name))
147            }
148        };
149
150        linkme.extend(res);
151        linkme
152    } else {
153        res
154    }
155}
balter_macros/lib.rs

balter_macros/
lib.rs
