Repository: CosmWasm/sylvia
Pull request #24 (2022-08-25 12:52:17)

Allow defining #exec and #query msgs in contract impl.


Patch:
diff --git a/contracts/cw1-whitelist/src/contract.rs b/contracts/cw1-whitelist/src/contract.rs
index 74f62cf..5e6fe4e 100644
--- a/contracts/cw1-whitelist/src/contract.rs
+++ b/contracts/cw1-whitelist/src/contract.rs
@@ -1,11 +1,17 @@
 use crate::error::ContractError;
-use cosmwasm_std::{Addr, DepsMut, Empty, Env, MessageInfo, Response};
+use cosmwasm_std::{Addr, Deps, DepsMut, Empty, Env, MessageInfo, Order, Response};
 
-use cw1::Cw1;
-use cw1::*;
+use cw1::{Cw1, FindMemberResponse};
 use cw_storage_plus::Map;
+use schemars::JsonSchema;
+use serde::{Deserialize, Serialize};
 use sylvia::contract;
 
+#[derive(Serialize, Deserialize, Clone, PartialEq, Eq, JsonSchema, Debug)]
+pub struct MembersCntResponse {
+    pub cnt: usize,
+}
+
 pub struct Cw1WhitelistContract {
     members: Map<'static, Addr, Empty>,
 }
@@ -62,6 +68,29 @@ impl Cw1WhitelistContract {
 
         Ok(Response::new())
     }
+
+    #[msg(exec)]
+    fn remove_member(
+        &self,
+        (deps, _env, msg): (DepsMut, Env, MessageInfo),
+        member: String,
+    ) -> Result<Response, ContractError> {
+        self.members
+            .remove(deps.storage, deps.api.addr_validate(&member)?);
+
+        Ok(Response::new()
+            .add_attribute("action", "remove_member")
+            .add_attribute("sender", msg.sender))
+    }
+
+    #[msg(query)]
+    fn members_cnt(&self, (deps, _env): (Deps, Env)) -> Result<MembersCntResponse, ContractError> {
+        let cnt = self
+            .members
+            .keys(deps.storage, None, None, Order::Ascending)
+            .count();
+        Ok(MembersCntResponse { cnt })
+    }
 }
 
 #[cfg(test)]
@@ -100,6 +129,49 @@ mod tests {
         );
     }
 
+    #[test]
+    fn binary_serialize_exec() {
+        let original_msg = ExecMsg::Cw1WhitelistContract(ImplExecMsg::RemoveMember {
+            member: "some_member".to_owned(),
+        });
+
+        let serialized_msg = to_binary(&original_msg).unwrap();
+        let serialized_msg: ExecMsg = from_binary(&serialized_msg).unwrap();
+
+        assert_eq!(serialized_msg, original_msg);
+    }
+
+    #[test]
+    fn slice_deserialize_exec() {
+        let deserialized: ExecMsg =
+            from_slice(br#"{"remove_member": {"member": "some_member"}}"#).unwrap();
+        assert_eq!(
+            deserialized,
+            ExecMsg::Cw1WhitelistContract(ImplExecMsg::RemoveMember {
+                member: "some_member".to_owned()
+            })
+        );
+    }
+
+    #[test]
+    fn binary_serialize_query() {
+        let original_msg = QueryMsg::Cw1WhitelistContract(ImplQueryMsg::MembersCnt {});
+
+        let serialized_msg = to_binary(&original_msg).unwrap();
+        let serialized_msg: QueryMsg = from_binary(&serialized_msg).unwrap();
+
+        assert_eq!(serialized_msg, original_msg);
+    }
+
+    #[test]
+    fn slice_deserialize_query() {
+        let deserialized: QueryMsg = from_slice(br#"{"members_cnt": {}}"#).unwrap();
+        assert_eq!(
+            deserialized,
+            QueryMsg::Cw1WhitelistContract(ImplQueryMsg::MembersCnt {})
+        );
+    }
+
     #[test]
     fn instantiate_empty() {
         let contract = Cw1WhitelistContract::new();
diff --git a/contracts/cw1-whitelist/src/lib.rs b/contracts/cw1-whitelist/src/lib.rs
index d5bc132..8b1bfe5 100644
--- a/contracts/cw1-whitelist/src/lib.rs
+++ b/contracts/cw1-whitelist/src/lib.rs
@@ -8,10 +8,8 @@ pub mod entry_points {
         entry_point, from_slice, Binary, Deps, DepsMut, Env, MessageInfo, Response,
     };
 
-    use crate::{
-        contract::{Cw1WhitelistContract, ExecMsg, InstantiateMsg, QueryMsg},
-        error::ContractError,
-    };
+    use crate::contract::{Cw1WhitelistContract, ExecMsg, InstantiateMsg, QueryMsg};
+    use crate::error::ContractError;
 
     const CONTRACT: Cw1WhitelistContract = Cw1WhitelistContract::new();
 
diff --git a/contracts/cw1/src/lib.rs b/contracts/cw1/src/lib.rs
index 3d3c9f9..5d7c757 100644
--- a/contracts/cw1/src/lib.rs
+++ b/contracts/cw1/src/lib.rs
@@ -1,4 +1,3 @@
-#![allow(dead_code)]
 use cosmwasm_std::{Deps, DepsMut, Env, MessageInfo, Response, StdError};
 
 use schemars::JsonSchema;
diff --git a/contracts/cw4/src/lib.rs b/contracts/cw4/src/lib.rs
index 59ad028..bfbac24 100644
--- a/contracts/cw4/src/lib.rs
+++ b/contracts/cw4/src/lib.rs
@@ -1,4 +1,3 @@
-#![allow(dead_code)]
 use cosmwasm_std::{Deps, DepsMut, Env, MessageInfo, Response, StdError};
 
 use schemars::JsonSchema;
diff --git a/sylvia-derive/src/input.rs b/sylvia-derive/src/input.rs
index 83bac3d..d4417e6 100644
--- a/sylvia-derive/src/input.rs
+++ b/sylvia-derive/src/input.rs
@@ -3,7 +3,7 @@ use proc_macro_error::emit_error;
 use quote::quote;
 use syn::{GenericParam, Ident, ItemImpl, ItemTrait, TraitItem};
 
-use crate::message::{EnumMessage, GlueMessage, StructMessage};
+use crate::message::{EnumMessage, GlueMessage, ImplEnumMessage, StructMessage};
 use crate::parser::{ContractArgs, InterfaceArgs, MsgType};
 
 /// Preprocessed `interface` macro input
@@ -111,23 +111,39 @@ impl<'a> ImplInput<'a> {
     }
 
     fn emit_messages(&self) -> TokenStream {
-        let instantiate = self.emit_msg(MsgType::Instantiate);
+        let instantiate = self.emit_struct_msg(MsgType::Instantiate);
+        let exec_impl = self.emit_enum_msg(
+            &Ident::new("ImplExecMsg", Span::mixed_site()),
+            MsgType::Exec,
+        );
+        let query_impl = self.emit_enum_msg(
+            &Ident::new("ImplQueryMsg", Span::mixed_site()),
+            MsgType::Query,
+        );
         let exec = self.emit_glue_msg(&Ident::new("ExecMsg", Span::mixed_site()), MsgType::Exec);
         let query = self.emit_glue_msg(&Ident::new("QueryMsg", Span::mixed_site()), MsgType::Query);
 
         quote! {
             #instantiate
 
+            #exec_impl
+
+            #query_impl
+
             #exec
 
             #query
         }
     }
 
-    fn emit_msg(&self, msg_ty: MsgType) -> TokenStream {
+    fn emit_struct_msg(&self, msg_ty: MsgType) -> TokenStream {
         StructMessage::new(self.item, msg_ty, &self.generics).map_or(quote! {}, |msg| msg.emit())
     }
 
+    fn emit_enum_msg(&self, name: &Ident, msg_ty: MsgType) -> TokenStream {
+        ImplEnumMessage::new(name, self.item, msg_ty, &self.generics, self.attributes).emit()
+    }
+
     fn emit_glue_msg(&self, name: &Ident, msg_ty: MsgType) -> TokenStream {
         GlueMessage::new(name, self.item, msg_ty, &self.attributes.error).emit()
     }
diff --git a/sylvia-derive/src/message.rs b/sylvia-derive/src/message.rs
index 86a327e..fe6cd3d 100644
--- a/sylvia-derive/src/message.rs
+++ b/sylvia-derive/src/message.rs
@@ -1,6 +1,6 @@
 use crate::check_generics::CheckGenerics;
 use crate::crate_module;
-use crate::parser::{ContractMessageAttr, InterfaceArgs, MsgAttr, MsgType};
+use crate::parser::{ContractArgs, ContractMessageAttr, InterfaceArgs, MsgAttr, MsgType};
 use crate::strip_generics::StripGenerics;
 use convert_case::{Case, Casing};
 use proc_macro2::{Span, TokenStream};
@@ -11,8 +11,8 @@ use syn::parse::{Parse, Parser};
 use syn::spanned::Spanned;
 use syn::visit::Visit;
 use syn::{
-    parse_quote, FnArg, GenericParam, Ident, ImplItem, ItemImpl, ItemTrait, Pat, PatType,
-    ReturnType, Signature, TraitItem, TraitItemMethod, Type, WhereClause, WherePredicate,
+    parse_quote, FnArg, GenericParam, Ident, ImplItem, ImplItemMethod, ItemImpl, ItemTrait, Pat,
+    PatType, ReturnType, Signature, TraitItem, TraitItemMethod, Type, WhereClause, WherePredicate,
 };
 
 fn filter_wheres<'a>(
@@ -345,6 +345,179 @@ impl<'a> EnumMessage<'a> {
     }
 }
 
+/// Representation of single enum message
+pub struct ImplEnumMessage<'a> {
+    name: &'a Ident,
+    variants: Vec<ImplMsgVariant<'a>>,
+    msg_ty: MsgType,
+    contract: &'a Type,
+    error: &'a Type,
+}
+
+impl<'a> ImplEnumMessage<'a> {
+    pub fn new(
+        name: &'a Ident,
+        source: &'a ItemImpl,
+        ty: MsgType,
+        generics: &'a [&'a GenericParam],
+        args: &'a ContractArgs,
+    ) -> Self {
+        let mut generics_checker = CheckGenerics::new(generics);
+        let variants: Vec<_> = source
+            .items
+            .iter()
+            .filter_map(|item| match item {
+                ImplItem::Method(method) => {
+                    let msg_attr = method.attrs.iter().find(|attr| attr.path.is_ident("msg"))?;
+                    let attr = match MsgAttr::parse.parse2(msg_attr.tokens.clone()) {
+                        Ok(attr) => attr,
+                        Err(err) => {
+                            emit_error!(method.span(), err);
+                            return None;
+                        }
+                    };
+
+                    if attr == ty {
+                        Some(ImplMsgVariant::new(method, &mut generics_checker))
+                    } else {
+                        None
+                    }
+                }
+                _ => None,
+            })
+            .collect();
+
+        Self {
+            name,
+            variants,
+            msg_ty: ty,
+            contract: &source.self_ty,
+            error: &args.error,
+        }
+    }
+
+    pub fn emit(&self) -> TokenStream {
+        let sylvia = crate_module();
+
+        let Self {
+            name,
+            variants,
+            msg_ty,
+            contract,
+            error,
+        } = self;
+
+        let match_arms = variants
+            .iter()
+            .map(|variant| variant.emit_dispatch_leg(*msg_ty));
+        let msgs: Vec<String> = variants
+            .iter()
+            .map(|var| var.name.to_string().to_case(Case::Snake))
+            .collect();
+        let msgs_cnt = msgs.len();
+        let variants = variants.iter().map(ImplMsgVariant::emit);
+
+        let ctx_type = msg_ty.emit_ctx_type();
+        let contract = StripGenerics.fold_type((*contract).clone());
+        let ret_type = msg_ty.emit_result_type(&None, error);
+
+        quote! {
+            #[derive(#sylvia ::serde::Serialize, #sylvia ::serde::Deserialize, Clone, Debug, PartialEq, #sylvia ::schemars::JsonSchema)]
+            #[serde(rename_all="snake_case")]
+            pub enum #name {
+                #(#variants,)*
+            }
+
+            impl #name {
+                pub fn dispatch(self, contract: &#contract, ctx: #ctx_type) -> #ret_type {
+                    use #name::*;
+
+                    match self {
+                        #(#match_arms,)*
+                    }
+                }
+                pub const fn messages() -> [&'static str; #msgs_cnt] {
+                    [#(#msgs,)*]
+                }
+            }
+        }
+    }
+}
+
+/// Representation of whole message variant
+pub struct ImplMsgVariant<'a> {
+    name: Ident,
+    function_name: &'a Ident,
+    // With https://github.com/rust-lang/rust/issues/63063 this could be just an iterator over
+    // `MsgField<'a>`
+    fields: Vec<MsgField<'a>>,
+}
+
+impl<'a> ImplMsgVariant<'a> {
+    /// Creates new message variant from trait method
+    pub fn new(
+        method: &'a ImplItemMethod,
+        generics_checker: &mut CheckGenerics,
+    ) -> ImplMsgVariant<'a> {
+        let function_name = &method.sig.ident;
+        let name = Ident::new(
+            &function_name.to_string().to_case(Case::UpperCamel),
+            function_name.span(),
+        );
+        let fields = process_fields(&method.sig, generics_checker);
+
+        Self {
+            name,
+            function_name,
+            fields,
+        }
+    }
+
+    /// Emits message variant
+    pub fn emit(&self) -> TokenStream {
+        let Self { name, fields, .. } = self;
+        let fields = fields.iter().map(MsgField::emit);
+
+        quote! {
+            #name {
+                #(#fields,)*
+            }
+        }
+    }
+
+    /// Emits match leg dispatching against this variant. Assumes enum variants are imported into the
+    /// scope. Dispatching is performed by calling the function this variant is build from on the
+    /// `contract` variable, with `ctx` as its first argument - both of them should be in scope.
+    pub fn emit_dispatch_leg(&self, msg_attr: MsgType) -> TokenStream {
+        use MsgType::*;
+
+        let Self {
+            name,
+            fields,
+            function_name,
+        } = self;
+        let args = fields.iter().map(|field| field.name);
+        let fields = fields.iter().map(|field| field.name);
+
+        match msg_attr {
+            Exec => quote! {
+                #name {
+                    #(#fields,)*
+                } => contract.#function_name(ctx.into(), #(#args),*).map_err(Into::into)
+            },
+            Query => quote! {
+                #name {
+                    #(#fields,)*
+                } => cosmwasm_std::to_binary(&contract.#function_name(ctx.into(), #(#args),*)?).map_err(Into::into)
+            },
+            Instantiate => {
+                emit_error!(name.span(), "Instantiation messages not supported on traits, they should be defined on contracts directly");
+                quote! {}
+            }
+        }
+    }
+}
+
 /// Representation of whole message variant
 pub struct MsgVariant<'a> {
     name: Ident,
@@ -552,11 +725,25 @@ impl<'a> GlueMessage<'a> {
             quote! { #variant(#module :: #name<#(#generics,)*>) }
         });
 
+        let impl_msg_name = match msg_ty {
+            MsgType::Exec => Ident::new("ImplExecMsg", Span::call_site()),
+            MsgType::Query => Ident::new("ImplQueryMsg", Span::call_site()),
+            MsgType::Instantiate => {
+                emit_error!(
+                    name.span(),
+                    "Do not emit InstantiateMsg using GlueMessage. Use StructMessage instead!"
+                );
+                Ident::new("", Span::call_site())
+            }
+        };
+        let impl_msg_name = quote! {#contract ( #impl_msg_name)};
+
         let dispatch_arms = interfaces.iter().map(|interface| {
             let ContractMessageAttr { variant, .. } = interface;
 
             quote! { #name :: #variant(msg) => msg.dispatch(contract, ctx) }
         });
+        let impl_dispatch_arm = quote! {#name :: #contract (msg) =>msg.dispatch(contract, ctx)};
 
         let deserialization_attempts = interfaces.iter().map(|interface| {
             let ContractMessageAttr { variant, .. } = interface;
@@ -570,6 +757,14 @@ impl<'a> GlueMessage<'a> {
             }
         });
 
+        let impl_var = Ident::new("contract", Span::mixed_site());
+        let impl_deserialization_attempt = quote! {
+            let #impl_var = match val.clone().deserialize_into() {
+                Ok(msg) => return Ok(Self:: #contract (msg)),
+                Err(err) => err,
+            };
+        };
+
         let deser_errors = interfaces.iter().map(|interface| {
             let ContractMessageAttr { variant, .. } = interface;
             let var = Ident::new(&variant.to_string().to_case(Case::Snake), variant.span());
@@ -577,6 +772,9 @@ impl<'a> GlueMessage<'a> {
             quote! { format!("As {}: {}", stringify!(#variant), #var) }
         });
 
+        let impl_deser_errors =
+            quote! { format!("As {}: {}", stringify!(#impl_msg_name), #impl_var) };
+
         let messages_type = interfaces.iter().map(|interface| &interface.variant);
 
         let ctx_type = msg_ty.emit_ctx_type();
@@ -584,9 +782,10 @@ impl<'a> GlueMessage<'a> {
 
         quote! {
             #[derive(#sylvia ::serde::Serialize, Clone, Debug, PartialEq, #sylvia ::schemars::JsonSchema)]
-            #[serde(rename_all="snake_case")]
+            #[serde(rename_all="snake_case", untagged)]
             pub enum #name {
                 #(#variants,)*
+                #impl_msg_name
             }
 
             impl #name {
@@ -597,6 +796,7 @@ impl<'a> GlueMessage<'a> {
                 ) -> #ret_type {
                     match self {
                         #(#dispatch_arms,)*
+                        #impl_dispatch_arm
                     }
                 }
             }
@@ -609,11 +809,13 @@ impl<'a> GlueMessage<'a> {
                     let val = #sylvia ::serde_value::Value::deserialize(deserializer)?;
 
                     #(#deserialization_attempts)*
+                    #impl_deserialization_attempt
 
                     Err(D::Error::custom(format!(
-                        "Expected any of {} messages, but cannot deserialize to neither of those\n{}",
+                        "Expected any of {} or {} messages, but cannot deserialize to neither of those\n{}",
                         stringify!(#(#messages_type),*),
-                        [#(#deser_errors,)*].join("\n")
+                        stringify!(#impl_msg_name),
+                        [#(#deser_errors,)* #impl_deser_errors].join("\n")
                     )))
                 }
             }
diff --git a/sylvia/examples/basic.rs b/sylvia/examples/basic.rs
index dee40c7..85b90d9 100644
--- a/sylvia/examples/basic.rs
+++ b/sylvia/examples/basic.rs
@@ -9,7 +9,8 @@ use sylvia::{contract, interface};
     sylvia::serde::Deserialize,
     Clone,
     Debug,
-    PartialEq, Eq,
+    PartialEq,
+    Eq,
     sylvia::schemars::JsonSchema,
 )]
 pub struct Member {
@@ -22,7 +23,8 @@ pub struct Member {
     sylvia::serde::Deserialize,
     Clone,
     Debug,
-    PartialEq, Eq,
+    PartialEq,
+    Eq,
     sylvia::schemars::JsonSchema,
 )]
 pub struct MemberResp {


Test patch:
diff --git a/contracts/cw1-whitelist/src/multitest.rs b/contracts/cw1-whitelist/src/multitest.rs
index 6caa149..c119b39 100644
--- a/contracts/cw1-whitelist/src/multitest.rs
+++ b/contracts/cw1-whitelist/src/multitest.rs
@@ -1,12 +1,12 @@
 #[cfg(test)]
 mod test {
-    use crate::contract::{Cw1WhitelistContract, InstantiateMsg};
+    use crate::contract::{Cw1WhitelistContract, ExecMsg, ImplExecMsg, InstantiateMsg};
     use anyhow::{bail, Result as AnyResult};
     use cosmwasm_std::{
         from_slice, Addr, Binary, DepsMut, Empty, Env, MessageInfo, Reply, Response,
     };
-    use cw1::{ExecMsg, FindMemberResponse, QueryMsg};
-    use cw_multi_test::{AppBuilder, Contract, Executor};
+    use cw1::{FindMemberResponse, QueryMsg};
+    use cw_multi_test::{App, Contract, Executor};
 
     impl Contract<Empty> for Cw1WhitelistContract {
         fn instantiate(
@@ -54,7 +54,7 @@ mod test {
 
     #[test]
     fn entry_points() {
-        let mut app = AppBuilder::new().build(|_, _, _| ());
+        let mut app = App::default();
         let contract_id = app.store_code(Box::new(Cw1WhitelistContract::new()));
 
         let contract = app
@@ -73,7 +73,7 @@ mod test {
         app.execute_contract(
             Addr::unchecked("owner"),
             contract.clone(),
-            &ExecMsg::AddMember {
+            &cw1::ExecMsg::AddMember {
                 member: "other_member".to_owned(),
             },
             &[],
@@ -92,4 +92,45 @@ mod test {
 
         assert_eq!(resp, FindMemberResponse { is_present: true });
     }
+
+    #[test]
+    fn contract_exec() {
+        let mut app = App::default();
+        let contract_id = app.store_code(Box::new(Cw1WhitelistContract::new()));
+
+        let contract = app
+            .instantiate_contract(
+                contract_id,
+                Addr::unchecked("Owner"),
+                &InstantiateMsg {
+                    members: vec!["member".to_owned(), "other_member".to_owned()],
+                },
+                &[],
+                "contract",
+                None,
+            )
+            .unwrap();
+
+        app.execute_contract(
+            Addr::unchecked("owner"),
+            contract.clone(),
+            &ImplExecMsg::RemoveMember {
+                member: "other_member".to_owned(),
+            },
+            &[],
+        )
+        .unwrap();
+
+        let resp: FindMemberResponse = app
+            .wrap()
+            .query_wasm_smart(
+                contract,
+                &QueryMsg::FindMember {
+                    member: "other_member".to_owned(),
+                },
+            )
+            .unwrap();
+
+        assert_eq!(resp, FindMemberResponse { is_present: false });
+    }
 }
diff --git a/sylvia/tests/messages_generation.rs b/sylvia/tests/messages_generation.rs
index 4dbef17..a2d14e6 100644
--- a/sylvia/tests/messages_generation.rs
+++ b/sylvia/tests/messages_generation.rs
@@ -2,7 +2,9 @@ use cosmwasm_std::{Addr, Decimal, Deps, DepsMut, Env, MessageInfo, Response, Std
 
 use sylvia::interface;
 
-#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq, schemars::JsonSchema)]
+#[derive(
+    serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq, schemars::JsonSchema,
+)]
 pub struct QueryResult;
 
 #[interface]
diff --git a/sylvia/tests/separate_module.rs b/sylvia/tests/separate_module.rs
index 0aaddbe..d07260b 100644
--- a/sylvia/tests/separate_module.rs
+++ b/sylvia/tests/separate_module.rs
@@ -2,7 +2,9 @@ use cosmwasm_std::{Addr, Decimal, Deps, DepsMut, Env, MessageInfo, Response, Std
 
 use sylvia::interface;
 
-#[derive(serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq, schemars::JsonSchema)]
+#[derive(
+    serde::Serialize, serde::Deserialize, Clone, Debug, PartialEq, Eq, schemars::JsonSchema,
+)]
 pub struct QueryResult;
 
 #[interface(module=msg)]