1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

use proc_macro::TokenStream;
use proc_macro2::{Span, TokenStream as TokenStream2};
use quote::quote;
use syn::{Ident, ItemFn};

/// Proc macro to denote a Transaction
///
/// NOTE: Currently this macro only works on functions with a `Result<T, E>` return value. This is a
/// restriction which will be lifted soon.
///
/// # Example
/// ```ignore
/// use balter::prelude::*;
///
/// #[transaction]
/// fn my_transaction(arg_1: u32, arg_2: &str) -> Result<String, MyError> {
///     ...
/// }
/// ```
#[proc_macro_attribute]
pub fn transaction(attr: TokenStream, item: TokenStream) -> TokenStream {
    transaction_internal(attr, item).into()
}

fn transaction_internal(_attr: TokenStream, item: TokenStream) -> TokenStream2 {
    let input = syn::parse::<ItemFn>(item).unwrap();

    let ItemFn {
        attrs,
        vis,
        sig,
        block,
    } = input;
    let stmts = &block.stmts;

    quote! {
        #(#attrs)* #vis #sig {
            ::balter::transaction::transaction_hook(async move {
                #(#stmts)*
            }).await
        }
    }
}

/// Proc macro to denote a Scenario
///
/// NOTE: Currently this macro only works on functions which take no arguments and with no return value.
/// (void functions). This is a restriction which will be lifted soon.
///
/// See the `Scenario` struct for more information on the methods this macro provides on functions.
///
/// # Example
/// ```ignore
/// use balter::prelude::*;
///
/// #[scenario]
/// fn my_scenario() {
/// }
/// ```
#[proc_macro_attribute]
pub fn scenario(attr: TokenStream, item: TokenStream) -> TokenStream {
    scenario_internal(attr, item).into()
}

fn scenario_internal(_attr: TokenStream, item: TokenStream) -> TokenStream2 {
    let input = syn::parse::<ItemFn>(item).unwrap();

    let ItemFn {
        attrs,
        vis,
        sig,
        block,
    } = input;
    let stmts = &block.stmts;

    let new_name = Ident::new(&format!("__balter_{}", sig.ident), Span::call_site());
    let mut new_sig = sig.clone();
    new_sig.ident = new_name.clone();

    let mut scen_sig = sig.clone();
    let scen_name = sig.ident.clone();
    scen_sig.asyncness = None;
    scen_sig.output = syn::parse(
        quote! {
            -> ::balter::scenario::Scenario
        }
        .into(),
    )
    .unwrap();

    let mut inter_sig = sig.clone();
    let inter_name = Ident::new(&format!("__balter_inter_{}", sig.ident), Span::call_site());
    inter_sig.ident = inter_name.clone();
    inter_sig.asyncness = None;
    inter_sig.output = syn::parse(
        quote! {
            -> ::std::pin::Pin<std::boxed::Box<dyn ::std::future::Future<Output = ()> + Send>>
        }
        .into(),
    )
    .unwrap();

    let static_name = Ident::new(
        &format!("__BALTER_{}", sig.ident.to_string().to_ascii_uppercase()),
        Span::call_site(),
    );

    quote! {
        #[::linkme::distributed_slice(::balter::runtime::BALTER_SCENARIOS)]
        static #static_name: (&'static str, fn() -> ::balter::scenario::Scenario) = (stringify!(#scen_name), #scen_name);

        #(#attrs)* #vis #scen_sig {
            ::balter::scenario::Scenario::new(stringify!(#scen_name), #inter_name)
        }

        #(#attrs)* #vis #inter_sig {
            Box::pin(async {
                #new_name().await
            })
        }

        #(#attrs)* #vis #new_sig {
            #(#stmts)*
        }
    }
}