feat: Generate migrate entry point if message defined on contract

CosmWasm · Sep 18, 2023 · 694a40a · 694a40a
1 parent d1c0d6e
commit 694a40a
Showing 1 changed file with 32 additions and 8 deletions.
diff --git 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,34 +1068,36 @@ 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,
         }
     }
 
     pub fn emit(&self) -> TokenStream {
         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
                 }
             }