From 575ac590e13bb95414089b181ac2a7d805e3eb23 Mon Sep 17 00:00:00 2001
From: Obie Kaku <obiekaku@gmail.com>
Date: Tue, 4 Jun 2024 10:03:19 -0500
Subject: [PATCH 1/2] feature: enable migrate msg to be either a struct or an
 enum

---
 pkg/codegen/instantiate.go | 27 ++++++++++++++++++++++++---
 1 file changed, 24 insertions(+), 3 deletions(-)

diff --git a/pkg/codegen/instantiate.go b/pkg/codegen/instantiate.go
index 86179de..5f8c3e4 100644
--- a/pkg/codegen/instantiate.go
+++ b/pkg/codegen/instantiate.go
@@ -24,7 +24,8 @@ func GenerateMigrateMsg(f *jen.File, schema *schemas.JSONSchema) {
 		return
 	}
 
-	generateStructMsg(f, schema, "MigrateMsg")
+	// MigrateMsg can be either a struct or an enum
+	generateStructOrEnumMsg(f, schema, "MigrateMsg")
 }
 
 func generateStructMsg(f *jen.File, schema *schemas.JSONSchema, allowedTitle string) {
@@ -46,11 +47,31 @@ func generateStructMsg(f *jen.File, schema *schemas.JSONSchema, allowedTitle str
 
 func validateAsStructMsg(schema *schemas.JSONSchema, allowedTitle string) error {
 	if schema.Title != allowedTitle {
-		return fmt.Errorf("title must be InstantiateMsg")
+		return fmt.Errorf("title must be %v", allowedTitle)
 	}
+
+	if schema.Type == nil {
+		return fmt.Errorf("%v type must be defined", allowedTitle)
+	}
+
 	if schema.Type[0] != "object" {
-		return fmt.Errorf("InstantiateMsg type must be object")
+		return fmt.Errorf("%v type must be object", allowedTitle)
 	}
 
 	return nil
 }
+
+// Generate the msg regardless if it is a struct or enum this is mostly needed for migration msgs
+func generateStructOrEnumMsg(f *jen.File, schema *schemas.JSONSchema, allowedTitle string) {
+	if schema == nil {
+		panic(fmt.Errorf("schema of %s is nil", allowedTitle))
+	}
+
+	err := validateAsStructMsg(schema, allowedTitle)
+	if err == nil {
+		generateStructMsg(f, schema, allowedTitle)
+	} else {
+		generateEnumMsg(f, schema, []string{allowedTitle, allowedTitle + "_for_Empty"})
+	}
+
+}

From 48e189595f2985a4e5785fbd1c61119093ce28a4 Mon Sep 17 00:00:00 2001
From: Obie Kaku <obiekaku@gmail.com>
Date: Tue, 4 Jun 2024 10:29:19 -0500
Subject: [PATCH 2/2] chore: ran gofumpt

---
 pkg/codegen/instantiate.go | 1 -
 1 file changed, 1 deletion(-)

diff --git a/pkg/codegen/instantiate.go b/pkg/codegen/instantiate.go
index 5f8c3e4..d08b885 100644
--- a/pkg/codegen/instantiate.go
+++ b/pkg/codegen/instantiate.go
@@ -73,5 +73,4 @@ func generateStructOrEnumMsg(f *jen.File, schema *schemas.JSONSchema, allowedTit
 	} else {
 		generateEnumMsg(f, schema, []string{allowedTitle, allowedTitle + "_for_Empty"})
 	}
-
 }