ic-platform-macros 0.1.16

// Derived from ic-cdk-macros 0.20.3; see README.md and LICENSE.
// Darling's generated FromMeta impl spells out the multiple field `guard: guard`.
#![allow(clippy::redundant_field_names)]

use darling::FromMeta;
use darling::ast::NestedMeta;
use proc_macro2::{Ident, Span, TokenStream};
use quote::{ToTokens, format_ident, quote};
use std::fmt::Formatter;
use syn::punctuated::Punctuated;
use syn::spanned::Spanned;
use syn::{
    Error, FnArg, ItemFn, Pat, PatIdent, PatType, Path, ReturnType, Signature, Type, parse_str,
};

#[derive(Default, FromMeta)]
struct ExportAttributes {
    pub name: Option<String>,
    #[darling(multiple)]
    pub guard: Vec<String>,
    /// The name of the function to use for decoding arguments.
    /// If not provided, the arguments are decoded as Candid.
    ///
    /// Even if the argument type is empty, the specified function will still be executed.
    pub decode_with: Option<String>,
    /// The name of the function to use for encoding the return value.
    /// If not provided, the return value is encoded as Candid.
    ///
    /// If the method returns a tuple, this custom encoder function should take the tuple as an argument.
    pub encode_with: Option<String>,
    #[darling(default)]
    pub manual_reply: bool,
    #[darling(default)]
    pub composite: bool,
    #[darling(default)]
    pub hidden: bool,
    #[darling(rename = "crate")]
    pub cratename: Option<String>,
    /// Path to the Candid crate or an equivalent explicitly supported facade.
    pub candid: Option<String>,
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
enum MethodType {
    Init,
    PreUpgrade,
    PostUpgrade,
    Update,
    Query,
    Heartbeat,
    InspectMessage,
    OnLowWasmMemory,
}

impl MethodType {
    /// A lifecycle method is a method that is called by the system and not by the user.
    /// So far, `update` and `query` are the only methods that are not lifecycle methods.
    ///
    /// We have a few assumptions for lifecycle methods:
    /// - They cannot have a return value.
    /// - The export name is prefixed with `canister_`, e.g. `init` => `canister_init`.
    pub fn is_lifecycle(&self) -> bool {
        match self {
            MethodType::Init
            | MethodType::PreUpgrade
            | MethodType::PostUpgrade
            | MethodType::Heartbeat
            | MethodType::InspectMessage
            | MethodType::OnLowWasmMemory => true,
            MethodType::Update | MethodType::Query => false,
        }
    }

    /// `init`, `post_upgrade`, `update`, `query` can have arguments.
    pub fn can_have_args(&self) -> bool {
        match self {
            MethodType::Init | MethodType::PostUpgrade | MethodType::Update | MethodType::Query => {
                true
            }
            MethodType::PreUpgrade
            | MethodType::Heartbeat
            | MethodType::InspectMessage
            | MethodType::OnLowWasmMemory => false,
        }
    }

    pub fn is_state_persistent(&self) -> bool {
        match self {
            Self::Query | Self::InspectMessage => false,
            Self::Update
            | Self::Heartbeat
            | Self::Init
            | Self::PreUpgrade
            | Self::PostUpgrade
            | Self::OnLowWasmMemory => true,
        }
    }
}

impl std::fmt::Display for MethodType {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            MethodType::Init => f.write_str("init"),
            MethodType::PreUpgrade => f.write_str("pre_upgrade"),
            MethodType::PostUpgrade => f.write_str("post_upgrade"),
            MethodType::Query => f.write_str("query"),
            MethodType::Update => f.write_str("update"),
            MethodType::Heartbeat => f.write_str("heartbeat"),
            MethodType::InspectMessage => f.write_str("inspect_message"),
            MethodType::OnLowWasmMemory => f.write_str("on_low_wasm_memory"),
        }
    }
}

fn get_args(method: MethodType, signature: &Signature) -> Result<Vec<(Ident, Box<Type>)>, Error> {
    // We only need the tuple of arguments, not their types. Magic of type inference.
    let mut args = vec![];
    for (i, arg) in signature.inputs.iter().enumerate() {
        let (ident, ty) = match arg {
            FnArg::Receiver(r) => {
                return Err(Error::new(
                    r.span(),
                    format!("#[{method}] cannot be above functions with `self` as a parameter."),
                ));
            }
            FnArg::Typed(PatType { pat, ty, .. }) => {
                let ident = if let Pat::Ident(PatIdent { ident, .. }) = pat.as_ref() {
                    // If the argument is named the same as the function, we need to rename it.
                    if ident == &signature.ident {
                        format_ident!("__arg_{}", ident, span = pat.span())
                    } else {
                        ident.clone()
                    }
                } else {
                    format_ident!("__unnamed_arg_{i}", span = pat.span())
                };
                (ident, ty.clone())
            }
        };

        args.push((ident, ty));
    }

    Ok(args)
}

fn dfn_macro(
    method: MethodType,
    attr: TokenStream,
    item: TokenStream,
) -> Result<TokenStream, Error> {
    let attr_span = attr.span();
    let attr_args = NestedMeta::parse_meta_list(attr)?;
    let attrs = ExportAttributes::from_list(&attr_args)?;

    let fun: ItemFn = syn::parse2::<syn::ItemFn>(item.clone()).map_err(|e| {
        Error::new(
            item.span(),
            format!("#[{method}] must be above a function. \n{e}"),
        )
    })?;
    let signature = &fun.sig;
    let generics = &signature.generics;

    if !generics.params.is_empty() {
        return Err(Error::new(
            generics.span(),
            format!("#[{method}] must be above a function with no generic parameters."),
        ));
    }
    let cratename: Path =
        syn::parse_str(attrs.cratename.as_deref().unwrap_or("::ic_platform::cdk"))?;
    let candid: Path = syn::parse_str(
        attrs
            .candid
            .as_deref()
            .unwrap_or("::ic_platform::candid::contract"),
    )?;

    // 1. function name(s)
    let name = &signature.ident;
    let outer_function_ident = format_ident!("__canister_method_{name}");
    let candid_method_name = format_ident!("__candid_method_{name}");
    let function_name = if let Some(custom_name) = attrs.name {
        if method.is_lifecycle() {
            return Err(Error::new(
                attr_span,
                format!("#[{method}] cannot have a custom name"),
            ));
        }
        if custom_name.starts_with("<ic-cdk internal>") {
            return Err(Error::new(
                attr_span,
                "Functions starting with `<ic-cdk internal>` are reserved for CDK internal use.",
            ));
        }
        custom_name
    } else {
        name.to_string()
    };
    let export_name = if method.is_lifecycle() {
        format!("canister_{method}")
    } else if method == MethodType::Query && attrs.composite {
        format!("canister_composite_query {function_name}",)
    } else {
        format!("canister_{method} {function_name}")
    };
    let host_compatible_name = export_name.replace(' ', ".").replace(['-', '<', '>'], "_");

    // 2. guard(s)
    if !attrs.guard.is_empty() && method.is_lifecycle() {
        return Err(Error::new(
            attr_span,
            format!("#[{method}] cannot have guard function(s)."),
        ));
    }
    let guards = attrs
        .guard
        .iter()
        .map(|guard_name| -> Result<_, Error> {
            let guard_path = parse_str::<Path>(guard_name)?;
            Ok(quote! {
                let r: Result<(), String> = #guard_path ();
                if let Err(e) = r {
                    #cratename::api::msg_reject(&e);
                    return;
                }
            })
        })
        .collect::<Result<Vec<_>, Error>>()?;
    let guard = quote! {
        #(#guards)*
    };

    // 3. decode arguments
    let (arg_tuple, _): (Vec<Ident>, Vec<Box<Type>>) =
        get_args(method, signature)?.iter().cloned().unzip();
    if !method.can_have_args() {
        if !arg_tuple.is_empty() {
            return Err(Error::new(
                Span::call_site(),
                format!("#[{method}] function cannot have arguments."),
            ));
        }
        if attrs.decode_with.is_some() {
            return Err(Error::new(
                attr_span,
                format!("#[{method}] function cannot have a decode_with attribute."),
            ));
        }
    }
    let arg_decode = if let Some(decode_with) = &attrs.decode_with {
        let decode_with_ident = parse_str::<Path>(decode_with)?;
        if arg_tuple.len() == 1 {
            let arg_one = &arg_tuple[0];
            quote! {
                let arg_bytes = #cratename::api::msg_arg_data();
                let #arg_one = #decode_with_ident(arg_bytes);
            }
        } else {
            quote! {
            let arg_bytes = #cratename::api::msg_arg_data();
            let ( #( #arg_tuple, )* ) = #decode_with_ident(arg_bytes); }
        }
    } else if arg_tuple.is_empty() {
        quote! {}
    } else {
        quote! {
            let arg_bytes = #cratename::api::msg_arg_data();
            let arg_bytes: &[u8] = if arg_bytes.is_empty() { b"DIDL\x00\x00" } else { &arg_bytes };
            let mut decoder_config = #candid::DecoderConfig::new();
            decoder_config.set_skipping_quota(10000);
            let ( #( #arg_tuple, )* ) = #candid::decode_args_with_config(arg_bytes, &decoder_config).unwrap();
        }
    };

    // 4. function call
    let function_call = if signature.asyncness.is_some() {
        quote! { #name ( #(#arg_tuple),* ) .await }
    } else {
        quote! { #name ( #(#arg_tuple),* ) }
    };

    // 5. return
    let return_length = match &signature.output {
        ReturnType::Default => 0,
        ReturnType::Type(_, ty) => match ty.as_ref() {
            Type::Tuple(tuple) => tuple.elems.len(),
            _ => 1,
        },
    };
    if method.is_lifecycle() {
        if return_length > 0 {
            return Err(Error::new(
                Span::call_site(),
                format!("#[{method}] function cannot have a return value."),
            ));
        }
        if attrs.encode_with.is_some() {
            return Err(Error::new(
                attr_span,
                format!("#[{method}] function cannot have an encode_with attribute."),
            ));
        }
    }
    let return_encode = if method.is_lifecycle() || attrs.manual_reply {
        quote! {}
    } else {
        let return_bytes = if let Some(encode_with) = &attrs.encode_with {
            let encode_with_ident = parse_str::<Path>(encode_with)?;
            match return_length {
                0 => quote! { #encode_with_ident()},
                _ => quote! { #encode_with_ident(result)},
            }
        } else {
            match return_length {
                0 => quote! { #candid::encode_one(()).unwrap() },
                1 => quote! { #candid::encode_one(result).unwrap() },
                _ => quote! { #candid::encode_args(result).unwrap() },
            }
        };
        quote! {
            let bytes: Vec<u8> = #return_bytes;
            #cratename::api::msg_reply(bytes);
        }
    };

    // 6. candid attributes for export_candid!()
    let candid_method_attr = if attrs.hidden {
        quote! {}
    } else {
        let annotation = match method {
            MethodType::Query if attrs.composite => {
                quote! { #[#candid::candid_method(composite_query, rename = #function_name)] }
            }
            MethodType::Query => {
                quote! { #[#candid::candid_method(query, rename = #function_name)] }
            }
            MethodType::Update => {
                quote! { #[#candid::candid_method(update, rename = #function_name)] }
            }
            MethodType::Init => quote! { #[#candid::candid_method(init)] },
            _ => quote! {},
        };
        let mut dummy_fun = fun.clone();
        dummy_fun.sig.ident = candid_method_name;
        dummy_fun.block = Box::new(syn::parse_quote!({
            panic!("candid dummy function called")
        }));
        if attrs.decode_with.is_some() {
            let mut inputs = Punctuated::new();
            inputs.push(syn::parse_quote!(arg_bytes: Vec<u8>));
            dummy_fun.sig.inputs = inputs;
        }
        if attrs.encode_with.is_some() {
            dummy_fun.sig.output = syn::parse_quote!(-> Vec<u8>);
        }
        let dummy_fun = dummy_fun.into_token_stream();
        quote! {
            #annotation
            #[allow(unused_variables)]
            #dummy_fun
        }
    };

    // 7. exported function body
    let async_context_name = if method.is_state_persistent() {
        format_ident!("in_executor_context")
    } else {
        format_ident!("in_query_executor_context")
    };
    let body = if signature.asyncness.is_some() {
        quote! {
            #cratename::futures::internals::#async_context_name(|| {
                #guard
                #[allow(clippy::disallowed_methods)]
                #cratename::futures::spawn(async {
                    #arg_decode
                    let result = #function_call;
                    #return_encode
                });
            });
        }
    } else {
        quote! {
            #guard
            #cratename::futures::internals::#async_context_name(|| {
                #arg_decode
                let result = #function_call;
                #return_encode
            });
        }
    };

    Ok(quote! {
        #[cfg_attr(target_family = "wasm", unsafe(export_name = #export_name))]
        #[cfg_attr(not(target_family = "wasm"), unsafe(export_name = #host_compatible_name))]
        fn #outer_function_ident() {
            #body
        }

        #candid_method_attr

        #item
    })
}

pub(crate) fn ic_query(attr: TokenStream, item: TokenStream) -> Result<TokenStream, Error> {
    dfn_macro(MethodType::Query, attr, item)
}

pub(crate) fn ic_update(attr: TokenStream, item: TokenStream) -> Result<TokenStream, Error> {
    dfn_macro(MethodType::Update, attr, item)
}

pub(crate) fn ic_init(attr: TokenStream, item: TokenStream) -> Result<TokenStream, Error> {
    dfn_macro(MethodType::Init, attr, item)
}

pub(crate) fn ic_pre_upgrade(attr: TokenStream, item: TokenStream) -> Result<TokenStream, Error> {
    dfn_macro(MethodType::PreUpgrade, attr, item)
}

pub(crate) fn ic_post_upgrade(attr: TokenStream, item: TokenStream) -> Result<TokenStream, Error> {
    dfn_macro(MethodType::PostUpgrade, attr, item)
}

pub(crate) fn ic_heartbeat(attr: TokenStream, item: TokenStream) -> Result<TokenStream, Error> {
    dfn_macro(MethodType::Heartbeat, attr, item)
}

pub(crate) fn ic_inspect_message(
    attr: TokenStream,
    item: TokenStream,
) -> Result<TokenStream, Error> {
    dfn_macro(MethodType::InspectMessage, attr, item)
}

pub(crate) fn ic_on_low_wasm_memory(
    attr: TokenStream,
    item: TokenStream,
) -> Result<TokenStream, Error> {
    dfn_macro(MethodType::OnLowWasmMemory, attr, item)
}

#[cfg(test)]
mod tests;