use super::{candid_path, docs::extract_doc_comments, get_lit_str};
use lazy_static::lazy_static;
use proc_macro2::TokenStream;
use quote::{quote, ToTokens};
use std::collections::BTreeMap;
use std::sync::Mutex;
use syn::{Error, ItemFn, Meta, Result, ReturnType, Signature, Type};
type RawArgs = Vec<String>;
type RawRets = Vec<String>;
type ParsedArgs = Vec<Type>;
type ParsedRets = Vec<Type>;
struct Method {
args: RawArgs,
rets: RawRets,
modes: String,
docs: Vec<String>,
}
lazy_static! {
static ref METHODS: Mutex<Option<BTreeMap<String, Method>>> =
Mutex::new(Some(BTreeMap::default()));
static ref INIT: Mutex<Option<Option<RawArgs>>> = Mutex::new(Some(Option::default()));
}
pub(crate) fn candid_method(attrs: Vec<Meta>, fun: ItemFn) -> Result<TokenStream> {
let attrs = get_candid_attribute(attrs)?;
let sig = &fun.sig;
if !sig.generics.params.is_empty() {
return Err(Error::new_spanned(
&sig.generics,
"candid_method doesn't support generic parameters",
));
}
let ident = sig.ident.to_string();
let name = attrs.rename.as_ref().unwrap_or(&ident).clone();
let modes = attrs.method_type.unwrap_or_else(|| "update".to_string());
let (args, rets) = get_args(sig)?;
let args: RawArgs = args
.iter()
.map(|t| format!("{}", t.to_token_stream()))
.collect();
let rets: RawRets = rets
.iter()
.map(|t| format!("{}", t.to_token_stream()))
.collect();
if modes == "oneway" && !rets.is_empty() {
return Err(Error::new_spanned(
&sig.output,
"oneway function should have no return value",
));
}
if attrs.is_init {
match (rets.len(), rets.first().map(|x| x.as_str())) {
(0, None) | (1, Some("Self")) => {
if let Some(init) = INIT.lock().unwrap().as_mut() {
*init = Some(args);
}
}
_ => {
return Err(Error::new_spanned(
&sig.output,
"init method should have no return value or return Self",
))
}
}
} else if let Some(map) = METHODS.lock().unwrap().as_mut() {
let docs = extract_doc_comments(&fun.attrs);
map.insert(
name.clone(),
Method {
args,
rets,
modes,
docs,
},
);
}
Ok(quote! { #fun })
}
pub(crate) fn export_service(path: Option<TokenStream>) -> TokenStream {
if let Some(meths) = METHODS.lock().unwrap().as_mut() {
let candid = candid_path(&path);
let init = if let Some(opt_args) = INIT.lock().unwrap().as_mut() {
let res = opt_args.as_ref().map(|args| {
let args = args
.iter()
.map(|t| generate_arg(quote! { init_args }, t))
.collect::<Vec<_>>();
quote! {
let mut init_args: Vec<Type> = Vec::new();
#(#args)*
}
});
*opt_args = None;
res
} else {
unreachable!();
};
let gen_tys = meths.iter().map(
|(
name,
Method {
args,
rets,
modes,
docs: method_docs,
},
)| {
let args = args
.iter()
.map(|t| generate_arg(quote! { args }, t))
.collect::<Vec<_>>();
let rets = rets
.iter()
.map(|t| generate_ret(quote! { rets }, t))
.collect::<Vec<_>>();
let modes = match modes.as_ref() {
"query" => quote! { vec![#candid::types::FuncMode::Query] },
"composite_query" => quote! { vec![#candid::types::FuncMode::CompositeQuery] },
"oneway" => quote! { vec![#candid::types::FuncMode::Oneway] },
"update" => quote! { vec![] },
_ => unreachable!(),
};
let doc_storage = if !method_docs.is_empty() {
quote! {
docs.add_service_method(#name.to_string(), vec![#(#method_docs.to_string(),)*]);
}
} else {
quote! {}
};
quote! {
{
#doc_storage
let mut args: Vec<Type> = Vec::new();
#(#args)*
let mut rets: Vec<Type> = Vec::new();
#(#rets)*
let func = Function { args, rets, modes: #modes };
service.push((#name.to_string(), TypeInner::Func(func).into()));
}
}
},
);
let service = quote! {
use #candid::types::{CandidType, Function, Type, TypeInner};
let mut service = Vec::<(String, Type)>::new();
let mut env = #candid::types::internal::TypeContainer::new();
let mut docs = #candid::pretty::candid::DocComments::empty();
#(#gen_tys)*
service.sort_unstable_by_key(|(name, _)| name.clone());
let ty = TypeInner::Service(service).into();
};
let actor = if let Some(init) = init {
quote! {
#init
let actor = Some(TypeInner::Class(init_args, ty).into());
}
} else {
quote! { let actor = Some(ty); }
};
let res = quote! {
fn __export_service() -> String {
#service
#actor
docs.extend_types(env.docs);
let result = #candid::pretty::candid::compile_with_docs(&env.env, &actor, &docs);
format!("{}", result)
}
};
meths.clear();
res
} else {
unreachable!()
}
}
fn generate_arg(name: TokenStream, ty: &str) -> TokenStream {
let ty = syn::parse_str::<Type>(ty).unwrap();
quote! {
#name.push(env.add::<#ty>());
}
}
fn generate_ret(name: TokenStream, ty: &str) -> TokenStream {
generate_arg(name, ty)
}
fn get_args(sig: &Signature) -> Result<(ParsedArgs, ParsedRets)> {
let mut args = Vec::new();
for arg in &sig.inputs {
match arg {
syn::FnArg::Receiver(r) => {
if r.reference.is_none() {
return Err(Error::new_spanned(arg, "only works for borrowed self"));
}
}
syn::FnArg::Typed(syn::PatType { ty, .. }) => args.push(ty.as_ref().clone()),
}
}
let rets = match &sig.output {
ReturnType::Default => Vec::new(),
ReturnType::Type(_, ty) => match ty.as_ref() {
Type::Tuple(tuple) => tuple.elems.iter().cloned().collect(),
_ => vec![ty.as_ref().clone()],
},
};
Ok((args, rets))
}
struct CandidAttribute {
rename: Option<String>,
method_type: Option<String>,
is_init: bool,
}
fn get_candid_attribute(attrs: Vec<Meta>) -> Result<CandidAttribute> {
let mut res = CandidAttribute {
rename: None,
method_type: None,
is_init: false,
};
for attr in attrs {
match &attr {
Meta::NameValue(m) if m.path.is_ident("rename") && res.rename.is_none() => {
if let Ok(lit) = get_lit_str(&m.value) {
res.rename = Some(lit.value());
}
}
Meta::Path(p) if res.method_type.is_none() => {
let mode = p.get_ident().unwrap().to_string();
match mode.as_ref() {
"query" | "composite_query" | "update" | "oneway" => {
res.method_type = Some(mode);
}
"init" => res.is_init = true,
_ => return Err(Error::new_spanned(p, "unknown mode")),
}
}
_ => return Err(Error::new_spanned(attr, "unknown or conflicting attribute")),
}
}
Ok(res)
}