implement ic-kit-macros (#13)

Cargo.toml

@@ -2,6 +2,7 @@
 members = [
     "ic-kit",
     "ic-kit-sys",
+    "ic-kit-macros",
 ]
 
 [profile.release]

ic-kit-macros/Cargo.toml

@@ -0,0 +1,25 @@
+[package]
+name = "ic-kit-macros"
+version = "0.1.0"
+edition = "2021"
+authors = ["Parsa Ghadimi <[email protected]>"]
+description = "IC-Kit's macros for canister development"
+license = "GPL-3.0"
+repository = "https://github.com/Psychedelic/ic-kit"
+documentation = "https://docs.rs/ic-kit-macros"
+categories = ["api-bindings", "development-tools::testing"]
+keywords = ["internet-computer", "canister", "cdk", "fleek"]
+include = ["src", "Cargo.toml", "README.md"]
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+quote = "1.0"
+proc-macro2 = "1.0"
+syn = "1.0"
+serde = "1.0"
+serde_tokenstream = "0.1"
+lazy_static = "1.4"
+
+[lib]
+proc-macro = true

ic-kit-macros/src/entry.rs

@@ -0,0 +1,202 @@
+//! Generate the Rust code for Internet Computer's [entry points] [1]
+//!
+//! [1]: <https://internetcomputer.org/docs/current/references/ic-interface-spec/#entry-points>
+
+use proc_macro2::{Ident, Span, TokenStream};
+use quote::quote;
+use serde::Deserialize;
+use serde_tokenstream::from_tokenstream;
+use std::fmt::Formatter;
+use syn::{
+    parse2, spanned::Spanned, Error, FnArg, ItemFn, Pat, PatIdent, PatType, ReturnType, Signature,
+    Type,
+};
+
+#[derive(Copy, Clone)]
+pub enum EntryPoint {
+    Init,
+    PreUpgrade,
+    PostUpgrade,
+    InspectMessage,
+    Heartbeat,
+    Update,
+    Query,
+}
+
+impl std::fmt::Display for EntryPoint {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            EntryPoint::Init => f.write_str("init"),
+            EntryPoint::PreUpgrade => f.write_str("pre_upgrade"),
+            EntryPoint::PostUpgrade => f.write_str("post_upgrade"),
+            EntryPoint::InspectMessage => f.write_str("inspect_message"),
+            EntryPoint::Heartbeat => f.write_str("heartbeat"),
+            EntryPoint::Update => f.write_str("update"),
+            EntryPoint::Query => f.write_str("query"),
+        }
+    }
+}
+
+impl EntryPoint {
+    pub fn is_lifecycle(&self) -> bool {
+        match &self {
+            EntryPoint::Update | EntryPoint::Query => false,
+            _ => true,
+        }
+    }
+}
+
+#[derive(Deserialize)]
+struct Config {
+    name: Option<String>,
+    guard: Option<String>,
+}
+
+fn collect_args(entry_point: EntryPoint, signature: &Signature) -> Result<Vec<Ident>, Error> {
+    let mut args = Vec::new();
+
+    for (id, arg) in signature.inputs.iter().enumerate() {
+        let ident = match arg {
+            FnArg::Receiver(r) => {
+                return Err(Error::new(
+                    r.span(),
+                    format!(
+                        "#[{}] macro can not be used on a function with `self` as a parameter.",
+                        entry_point
+                    ),
+                ))
+            }
+            FnArg::Typed(PatType { pat, .. }) => {
+                if let Pat::Ident(PatIdent { ident, .. }) = pat.as_ref() {
+                    ident.clone()
+                } else {
+                    Ident::new(&format!("arg_{}", id), pat.span())
+                }
+            }
+        };
+
+        args.push(ident)
+    }
+
+    Ok(args)
+}
+
+/// Process a rust syntax and generate the code for processing it.
+pub fn gen_entry_point_code(
+    entry_point: EntryPoint,
+    attr: TokenStream,
+    item: TokenStream,
+) -> Result<TokenStream, Error> {
+    let attrs = from_tokenstream::<Config>(&attr)?;
+    let fun: ItemFn = parse2::<ItemFn>(item.clone()).map_err(|e| {
+        Error::new(
+            item.span(),
+            format!("#[{0}] must be above a function. \n{1}", entry_point, e),
+        )
+    })?;
+    let signature = &fun.sig;
+    let generics = &signature.generics;
+
+    if !generics.params.is_empty() {
+        return Err(Error::new(
+            generics.span(),
+            format!(
+                "#[{}] must be above a function with no generic parameters.",
+                entry_point
+            ),
+        ));
+    }
+
+    let is_async = signature.asyncness.is_some();
+
+    let return_length = match &signature.output {
+        ReturnType::Default => 0,
+        ReturnType::Type(_, ty) => match ty.as_ref() {
+            Type::Tuple(tuple) => tuple.elems.len(),
+            _ => 1,
+        },
+    };
+
+    if entry_point.is_lifecycle() && return_length > 0 {
+        return Err(Error::new(
+            Span::call_site(),
+            format!("#[{}] function cannot have a return value.", entry_point),
+        ));
+    }
+
+    let arg_tuple: Vec<Ident> = collect_args(entry_point, signature)?;
+    let name = &signature.ident;
+
+    let outer_function_ident = Ident::new(
+        &format!("canister_{}_{}_", entry_point, name),
+        Span::call_site(),
+    );
+
+    let export_name = if entry_point.is_lifecycle() {
+        format!("canister_{}", entry_point)
+    } else {
+        format!(
+            "canister_{0} {1}",
+            entry_point,
+            attrs.name.unwrap_or_else(|| name.to_string())
+        )
+    };
+
+    let function_call = if is_async {
+        quote! { #name ( #(#arg_tuple),* ) .await }
+    } else {
+        quote! { #name ( #(#arg_tuple),* ) }
+    };
+
+    let arg_count = arg_tuple.len();
+
+    let return_encode = if entry_point.is_lifecycle() {
+        quote! {}
+    } else {
+        match return_length {
+            0 => quote! { ic_kit::ic_call_api_v0_::reply(()) },
+            1 => quote! { ic_kit::ic_call_api_v0_::reply((result,)) },
+            _ => quote! { ic_kit::ic_call_api_v0_::reply(result) },
+        }
+    };
+
+    // On initialization we can actually not receive any input and it's okay, only if
+    // we don't have any arguments either.
+    // If the data we receive is not empty, then try to unwrap it as if it's DID.
+    let arg_decode = if entry_point.is_lifecycle() && arg_count == 0 {
+        quote! {}
+    } else {
+        quote! { let ( #( #arg_tuple, )* ) = ic_kit::ic_call_api_v0_::arg_data(); }
+    };
+
+    let guard = if let Some(guard_name) = attrs.guard {
+        let guard_ident = Ident::new(&guard_name, Span::call_site());
+
+        quote! {
+            let r: Result<(), String> = #guard_ident ();
+            if let Err(e) = r {
+                ic_kit::ic_call_api_v0_::reject(&e);
+                return;
+            }
+        }
+    } else {
+        quote! {}
+    };
+
+    Ok(quote! {
+        #[export_name = #export_name]
+        fn #outer_function_ident() {
+            ic_kit::setup();
+
+            #guard
+
+            ic_kit::ic::spawn(async {
+                #arg_decode
+                let result = #function_call;
+                #return_encode
+            });
+        }
+
+        #item
+    })
+}

ic-kit-macros/src/lib.rs

@@ -0,0 +1,56 @@
+mod entry;
+
+use entry::{gen_entry_point_code, EntryPoint};
+use proc_macro::TokenStream;
+
+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)
+}

ic-kit/Cargo.toml

@@ -13,8 +13,8 @@ keywords = ["internet-computer", "canister", "cdk", "fleek"]
 include = ["src", "Cargo.toml", "README.md"]
 
 [dependencies]
+ic-kit-macros = {path="../ic-kit-macros", version="0.1.0"}
 ic-cdk = "0.5"
-ic-cdk-macros = "0.5"
 candid="0.7"
 serde = { version="1.0.130", features = ["derive"] }
 serde_bytes = "0.11.5"

ic-kit/src/ic.rs

@@ -200,7 +200,7 @@ pub fn stable_bytes() -> Vec<u8> {
 ///
 /// impl Counter {
 ///     fn get(&self) -> u64 {
-///         self.count
+///         *self.count
 ///     }
 /// }
 ///
@@ -242,7 +242,7 @@ pub fn maybe_with<T: 'static, U, F: FnOnce(&T) -> U>(callback: F) -> Option<U> {
 /// impl Counter {
 ///     fn increment(&mut self) -> u64 {
 ///         self.count += 1;
-///         self.count
+///         *self.count
 ///     }
 /// }
 ///

ic-kit/src/lib.rs

@@ -1,11 +1,13 @@
 pub use handler::*;
 pub use interface::*;
 pub use mock::*;
+pub use setup::*;
 
 mod handler;
 mod inject;
 mod interface;
 mod mock;
+mod setup;
 #[cfg(target_family = "wasm")]
 mod wasm;
 
@@ -61,4 +63,12 @@ pub use async_std::test as async_test;
 pub use ic_cdk::api::call::{CallResult, RejectionCode};
 pub use ic_cdk::export::candid;
 pub use ic_cdk::export::Principal;
-pub use ic_cdk_macros as macros;
+pub use ic_kit_macros as macros;
+
+/// ic_cdk APIs to be used with ic-kit-macros only, please don't use this directly
+/// we may decide to change it anytime and break compatability.
+pub mod ic_call_api_v0_ {
+    pub use ic_cdk::api::call::arg_data;
+    pub use ic_cdk::api::call::reject;
+    pub use ic_cdk::api::call::reply;
+}

ic-kit/src/setup.rs

@@ -0,0 +1,36 @@
+use crate::ic;
+use std::panic;
+
+static mut DONE: bool = false;
+
+pub fn setup() {
+    unsafe {
+        if DONE {
+            return;
+        }
+        DONE = true;
+    }
+
+    set_panic_hook()
+}
+
+/// Sets a custom panic hook, uses debug.trace
+pub fn set_panic_hook() {
+    panic::set_hook(Box::new(|info| {
+        let file = info.location().unwrap().file();
+        let line = info.location().unwrap().line();
+        let col = info.location().unwrap().column();
+
+        let msg = match info.payload().downcast_ref::<&'static str>() {
+            Some(s) => *s,
+            None => match info.payload().downcast_ref::<String>() {
+                Some(s) => &s[..],
+                None => "Box<Any>",
+            },
+        };
+
+        let err_info = format!("Panicked at '{}', {}:{}:{}", msg, file, line, col);
+        ic::print(&err_info);
+        ic::trap(&err_info);
+    }));
+}