From 34e537e286a6f03d5cf8bc1afb7c0d32c31bcece Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20Wo=C5=BAniak?= <jnwznk@gmail.com>
Date: Mon, 18 Sep 2023 17:53:53 +0200
Subject: [PATCH] feat: Generate migrate entry point if message defined on
 contract

---
 sylvia-derive/src/message.rs        | 40 +++++++++++++++++++++++------
 sylvia/tests/messages_generation.rs | 24 +++++++++++++++--
 2 files changed, 54 insertions(+), 10 deletions(-)

diff --git a/sylvia-derive/src/message.rs b/sylvia-derive/src/message.rs
index 9990d57f..3bb14c52 100644
--- a/sylvia-derive/src/message.rs
+++ b/sylvia-derive/src/message.rs
@@ -1041,9 +1041,9 @@ impl<'a> GlueMessage<'a> {
 pub struct EntryPoints<'a> {
     name: Type,
     error: Type,
-    reply: Option<Ident>,
     custom: Custom<'a>,
     override_entry_points: OverrideEntryPoints,
+    variants: MsgVariants<'a>,
 }
 
 impl<'a> EntryPoints<'a> {
@@ -1068,19 +1068,16 @@ impl<'a> EntryPoints<'a> {
             .unwrap_or_else(|| parse_quote! { #sylvia ::cw_std::StdError });
 
         let generics: Vec<_> = source.generics.params.iter().collect();
-        let reply = MsgVariants::new(source.as_variants(), &generics)
-            .0
-            .into_iter()
-            .find(|variant| variant.msg_type == MsgType::Reply)
-            .map(|variant| variant.function_name.clone());
+
+        let variants = MsgVariants::new(source.as_variants(), &generics);
         let custom = Custom::new(&source.attrs);
 
         Self {
             name,
             error,
-            reply,
             custom,
             override_entry_points,
+            variants,
         }
     }
 
@@ -1088,14 +1085,19 @@ impl<'a> EntryPoints<'a> {
         let Self {
             name,
             error,
-            reply,
             custom,
             override_entry_points,
+            variants,
         } = self;
         let sylvia = crate_module();
 
         let custom_msg = custom.msg_or_default();
         let custom_query = custom.query_or_default();
+        let reply = variants
+            .0
+            .iter()
+            .find(|variant| variant.msg_type == MsgType::Reply)
+            .map(|variant| variant.function_name.clone());
 
         #[cfg(not(tarpaulin_include))]
         {
@@ -1114,6 +1116,26 @@ impl<'a> EntryPoints<'a> {
                     },
                 );
 
+            let migrate_not_overridden = override_entry_points
+                .get_entry_point(MsgType::Migrate)
+                .is_none();
+            let migrate_msg_defined = variants
+                .0
+                .iter()
+                .any(|variant| variant.msg_type == MsgType::Migrate);
+
+            let migrate = if migrate_not_overridden && migrate_msg_defined {
+                OverrideEntryPoint::emit_default_entry_point(
+                    &custom_msg,
+                    &custom_query,
+                    name,
+                    error,
+                    MsgType::Migrate,
+                )
+            } else {
+                quote! {}
+            };
+
             let reply_ep = override_entry_points
                 .get_entry_point(MsgType::Reply)
                 .map(|_| quote! {})
@@ -1137,6 +1159,8 @@ impl<'a> EntryPoints<'a> {
 
                     #(#entry_points)*
 
+                    #migrate
+
                     #reply_ep
                 }
             }
diff --git a/sylvia/tests/messages_generation.rs b/sylvia/tests/messages_generation.rs
index 7cb34e07..2495fc69 100644
--- a/sylvia/tests/messages_generation.rs
+++ b/sylvia/tests/messages_generation.rs
@@ -46,15 +46,17 @@ mod interface {
 }
 
 mod contract {
-    use cosmwasm_std::{Addr, Response, StdResult};
+    use cosmwasm_std::{Addr, Reply, Response, StdResult};
     use sylvia::contract;
-    use sylvia::types::{ExecCtx, InstantiateCtx, MigrateCtx, QueryCtx};
+    use sylvia::types::{ExecCtx, InstantiateCtx, MigrateCtx, QueryCtx, ReplyCtx};
+    use sylvia_derive::entry_points;
 
     use crate::{MyQuery, QueryResult};
 
     pub struct Contract {}
 
     #[cfg(not(tarpaulin_include))]
+    #[entry_points]
     #[contract]
     #[allow(dead_code)]
     #[sv::custom(query=MyQuery)]
@@ -100,6 +102,11 @@ mod contract {
         fn argumented_query(&self, _ctx: QueryCtx<MyQuery>, _user: Addr) -> StdResult<QueryResult> {
             Ok(QueryResult {})
         }
+
+        #[msg(reply)]
+        fn my_reply(&self, _ctx: ReplyCtx<MyQuery>, _reply: Reply) -> StdResult<Response> {
+            Ok(Response::new())
+        }
     }
 }
 
@@ -154,3 +161,16 @@ fn contract_messages_constructible() {
         contract::QueryMsg::ArgumentedQuery { .. } => (),
     }
 }
+
+#[test]
+fn entry_points_generation() {
+    use contract::entry_points;
+
+    let _ = cw_multi_test::ContractWrapper::new(
+        entry_points::execute,
+        entry_points::instantiate,
+        entry_points::query,
+    )
+    .with_migrate(entry_points::migrate)
+    .with_reply(entry_points::reply);
+}