1use itertools::Itertools;
2use proc_macro::TokenStream;
3use proc_macro2::TokenTree;
4use quote::quote;
5use syn::{parse_macro_input, punctuated::Punctuated, DeriveInput};
6
7macro_rules! match_kv_attr {
8 ($key:expr, $value_type:tt) => {
9 |tt| {
10 if let [TokenTree::Ident(key), TokenTree::Punct(eq), TokenTree::$value_type(value)] =
11 &tt[..]
12 {
13 if (key == $key) && (eq.as_char() == '=') {
14 Some(quote!(#value))
15 } else {
16 None
17 }
18 } else {
19 None
20 }
21 }
22 };
23}
24
25#[proc_macro_derive(CosmwasmExt, attributes(proto_message, proto_query))]
26pub fn derive_cosmwasm_ext(input: TokenStream) -> TokenStream {
27 let input = parse_macro_input!(input as DeriveInput);
28 let ident = input.ident;
29
30 let type_url = get_type_url(&input.attrs);
31
32 let (query_request_conversion, cosmwasm_query) = if get_attr("proto_query", &input.attrs)
37 .is_some()
38 {
39 let path = get_query_attrs(&input.attrs, match_kv_attr!("path", Literal));
40 let res = get_query_attrs(&input.attrs, match_kv_attr!("response_type", Ident));
41
42 let query_request_conversion = quote! {
43 impl <Q: cosmwasm_std::CustomQuery> From<#ident> for cosmwasm_std::QueryRequest<Q> {
44 fn from(msg: #ident) -> Self {
45 cosmwasm_std::QueryRequest::<Q>::Stargate {
46 path: #path.to_string(),
47 data: msg.into(),
48 }
49 }
50 }
51 };
52
53 let cosmwasm_query = quote! {
54 pub fn query(self, querier: &cosmwasm_std::QuerierWrapper<impl cosmwasm_std::CustomQuery>) -> cosmwasm_std::StdResult<#res> {
55 querier.query::<#res>(&self.into())
56 }
57
58 pub fn mock_response<T: desmos_mock::MockableQuerier>(querier: &mut T, response: #res) {
59 querier.register_custom_query(#path.to_string(), Box::new(move |data| {
60 cosmwasm_std::SystemResult::Ok(cosmwasm_std::ContractResult::Ok(
61 cosmwasm_std::to_binary(&response)
62 .unwrap()))
63 }))
64 }
65
66 pub fn mock_failed_response<T: desmos_mock::MockableQuerier>(querier: &mut T, error: String) {
67 querier.register_custom_query(#path.to_string(), Box::new(move |data| {
68 cosmwasm_std::SystemResult::Err(cosmwasm_std::SystemError::InvalidResponse {
69 error: error.clone(),
70 response: cosmwasm_std::Binary::default(),
71 })
72 }))
73 }
74 };
75
76 (query_request_conversion, cosmwasm_query)
77 } else {
78 (quote!(), quote!())
79 };
80
81 (quote! {
82 impl #ident {
83 pub const TYPE_URL: &'static str = #type_url;
84 #cosmwasm_query
85 }
86
87 #query_request_conversion
88
89 impl From<#ident> for cosmwasm_std::Binary {
90 fn from(msg: #ident) -> Self {
91 let mut bytes = Vec::new();
92 prost::Message::encode(&msg, &mut bytes)
93 .expect("Message encoding must be infallible");
94
95 cosmwasm_std::Binary(bytes)
96 }
97 }
98
99 impl<T> From<#ident> for cosmwasm_std::CosmosMsg<T> {
100 fn from(msg: #ident) -> Self {
101 cosmwasm_std::CosmosMsg::<T>::Stargate {
102 type_url: #type_url.to_string(),
103 value: msg.into(),
104 }
105 }
106 }
107
108 impl TryFrom<cosmwasm_std::Binary> for #ident {
109 type Error = cosmwasm_std::StdError;
110
111 fn try_from(binary: cosmwasm_std::Binary) -> std::result::Result<Self, Self::Error> {
112 use ::prost::Message;
113 Self::decode(&binary[..]).map_err(|e| {
114 cosmwasm_std::StdError::parse_err(stringify!(#ident).to_string(),
115 format!(
116 "Unable to decode binary: \n - base64: {}\n - bytes array: {:?}\n\n{:?}",
117 binary,
118 binary.to_vec(),
119 e
120 ))
121 })
122 }
123 }
124 })
125 .into()
126}
127
128fn get_type_url(attrs: &Vec<syn::Attribute>) -> proc_macro2::TokenStream {
129 let proto_message = get_attr("proto_message", attrs).and_then(|a| Some(a.meta.clone()));
130
131 if let Some(syn::Meta::List(meta)) = proto_message.clone() {
132 let nested = meta
133 .parse_args_with(Punctuated::<syn::Meta, syn::Token![,]>::parse_terminated)
134 .unwrap();
135 if let syn::Meta::NameValue(meta) = nested[0].clone() {
136 if !meta.path.is_ident("type_url") {
137 return proto_message_attr_error(meta.value);
138 }
139
140 if let syn::Expr::Lit(expr_lit) = &meta.value {
141 if let syn::Lit::Str(lit_str) = &expr_lit.lit {
142 return quote!(#lit_str);
143 }
144 }
145
146 return proto_message_attr_error(meta.value);
147 }
148 }
149
150 proto_message_attr_error(proto_message)
151}
152
153fn get_query_attrs<F>(attrs: &Vec<syn::Attribute>, f: F) -> proc_macro2::TokenStream
154where
155 F: FnMut(&Vec<TokenTree>) -> Option<proc_macro2::TokenStream>,
156{
157 let proto_query = get_attr("proto_query", attrs);
158
159 if let Some(attr) = proto_query {
160 let list = attr.meta.require_list().unwrap();
161
162 let kv_groups = list.tokens.clone().into_iter().group_by(|t| {
163 if let TokenTree::Punct(punct) = t {
164 punct.as_char() != ','
165 } else {
166 true
167 }
168 });
169 let mut key_values: Vec<Vec<TokenTree>> = vec![];
170
171 for (non_sep, g) in &kv_groups {
172 if non_sep {
173 key_values.push(g.collect());
174 }
175 }
176
177 return key_values
178 .iter()
179 .find_map(f)
180 .unwrap_or_else(|| proto_query_attr_error(proto_query));
181 }
182
183 proto_query_attr_error(proto_query)
184}
185
186fn get_attr<'a>(attr_ident: &str, attrs: &'a Vec<syn::Attribute>) -> Option<&'a syn::Attribute> {
187 for attr in attrs {
188 if attr.path().segments.len() == 1 && attr.path().segments[0].ident == attr_ident {
189 return Some(attr);
190 }
191 }
192 None
193}
194
195fn proto_message_attr_error<T: quote::ToTokens>(tokens: T) -> proc_macro2::TokenStream {
196 syn::Error::new_spanned(tokens, "expected `proto_message(type_url = \"...\")`")
197 .to_compile_error()
198}
199
200fn proto_query_attr_error<T: quote::ToTokens>(tokens: T) -> proc_macro2::TokenStream {
201 syn::Error::new_spanned(
202 tokens,
203 "expected `proto_query(path = \"...\", response_type = ...)`",
204 )
205 .to_compile_error()
206}