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

mod entry;
mod export_service;
mod test;

use entry::{gen_entry_point_code, EntryPoint};
use proc_macro::TokenStream;
use syn::parse_macro_input;
use test::gen_test_code;

fn process_entry_point(
    entry_point: EntryPoint,
    attr: TokenStream,
    item: TokenStream,
) -> TokenStream {
    gen_entry_point_code(entry_point, attr.into(), item.into())
        .unwrap_or_else(|error| error.to_compile_error())
        .into()
}

/// Export the function as the init hook of the canister.
#[proc_macro_attribute]
pub fn init(attr: TokenStream, item: TokenStream) -> TokenStream {
    process_entry_point(EntryPoint::Init, attr, item)
}

/// Export the function as the pre_upgrade hook of the canister.
#[proc_macro_attribute]
pub fn pre_upgrade(attr: TokenStream, item: TokenStream) -> TokenStream {
    process_entry_point(EntryPoint::PreUpgrade, attr, item)
}

/// Export the function as the post_upgrade hook of the canister.
#[proc_macro_attribute]
pub fn post_upgrade(attr: TokenStream, item: TokenStream) -> TokenStream {
    process_entry_point(EntryPoint::PostUpgrade, attr, item)
}

/// Export the function as the inspect_message hook of the canister.
#[proc_macro_attribute]
pub fn inspect_message(attr: TokenStream, item: TokenStream) -> TokenStream {
    process_entry_point(EntryPoint::InspectMessage, attr, item)
}

/// Export the function as the heartbeat hook of the canister.
#[proc_macro_attribute]
pub fn heartbeat(attr: TokenStream, item: TokenStream) -> TokenStream {
    process_entry_point(EntryPoint::Heartbeat, attr, item)
}

/// Export an update method for the canister.
#[proc_macro_attribute]
pub fn update(attr: TokenStream, item: TokenStream) -> TokenStream {
    process_entry_point(EntryPoint::Update, attr, item)
}

/// Export a query method for the canister.
#[proc_macro_attribute]
pub fn query(attr: TokenStream, item: TokenStream) -> TokenStream {
    process_entry_point(EntryPoint::Query, attr, item)
}

/// A macro to generate IC-Kit tests.
#[proc_macro_attribute]
pub fn kit_test(attr: TokenStream, item: TokenStream) -> TokenStream {
    gen_test_code(attr.into(), item.into())
        .unwrap_or_else(|error| error.to_compile_error())
        .into()
}

#[proc_macro_derive(KitCanister, attributes(candid_path))]
pub fn kit_export(input: TokenStream) -> TokenStream {
    let input = parse_macro_input!(input as syn::DeriveInput);
    let save_candid_path_result = get_save_candid_path(&input);

    match save_candid_path_result {
        Ok(save_candid_path) => export_service::export_service(input, save_candid_path).into(),
        Err(e) => e.to_compile_error().into(),
    }
}

fn get_save_candid_path(input: &syn::DeriveInput) -> syn::Result<Option<syn::LitStr>> {
    let candid_path_helper_attribute_option = input
        .attrs
        .iter()
        .find(|attr| attr.path.is_ident("candid_path"));

    match candid_path_helper_attribute_option {
        Some(candid_path_helper_attribute) => {
            let custom_candid_path_lit: syn::LitStr = candid_path_helper_attribute.parse_args()?;
            Ok(Some(custom_candid_path_lit))
        }
        None => Ok(None),
    }
}