refactor(schema)!: collect all types in ModuleSchema (cosmos#21036)

schema/decoding/decoding_test.go

@@ -366,8 +366,9 @@ func (e exampleBankModule) subBalance(acct, denom string, amount uint64) error {
 	return nil
 }
 
-var exampleBankSchema = schema.ModuleSchema{
-	ObjectTypes: []schema.ObjectType{
+func init() {
+	var err error
+	exampleBankSchema, err = schema.NewModuleSchema([]schema.ObjectType{
 		{
 			Name: "balances",
 			KeyFields: []schema.Field{
@@ -387,9 +388,14 @@ var exampleBankSchema = schema.ModuleSchema{
 				},
 			},
 		},
-	},
+	})
+	if err != nil {
+		panic(err)
+	}
 }
 
+var exampleBankSchema schema.ModuleSchema
+
 func (e exampleBankModule) ModuleCodec() (schema.ModuleCodec, error) {
 	return schema.ModuleCodec{
 		Schema: exampleBankSchema,
@@ -426,17 +432,25 @@ type oneValueModule struct {
 	store *testStore
 }
 
-var oneValueModSchema = schema.ModuleSchema{
-	ObjectTypes: []schema.ObjectType{
-		{
-			Name: "item",
-			ValueFields: []schema.Field{
-				{Name: "value", Kind: schema.StringKind},
+func init() {
+	var err error
+	oneValueModSchema, err = schema.NewModuleSchema(
+		[]schema.ObjectType{
+			{
+				Name: "item",
+				ValueFields: []schema.Field{
+					{Name: "value", Kind: schema.StringKind},
+				},
 			},
 		},
-	},
+	)
+	if err != nil {
+		panic(err)
+	}
 }
 
+var oneValueModSchema schema.ModuleSchema
+
 func (i oneValueModule) ModuleCodec() (schema.ModuleCodec, error) {
 	return schema.ModuleCodec{
 		Schema: oneValueModSchema,

schema/decoding/resolver_test.go

@@ -10,16 +10,24 @@ import (
 type modA struct{}
 
 func (m modA) ModuleCodec() (schema.ModuleCodec, error) {
+	modSchema, err := schema.NewModuleSchema([]schema.ObjectType{{Name: "A", KeyFields: []schema.Field{{Name: "field1", Kind: schema.StringKind}}}})
+	if err != nil {
+		return schema.ModuleCodec{}, err
+	}
 	return schema.ModuleCodec{
-		Schema: schema.ModuleSchema{ObjectTypes: []schema.ObjectType{{Name: "A"}}},
+		Schema: modSchema,
 	}, nil
 }
 
 type modB struct{}
 
 func (m modB) ModuleCodec() (schema.ModuleCodec, error) {
+	modSchema, err := schema.NewModuleSchema([]schema.ObjectType{{Name: "B", KeyFields: []schema.Field{{Name: "field2", Kind: schema.StringKind}}}})
+	if err != nil {
+		return schema.ModuleCodec{}, err
+	}
 	return schema.ModuleCodec{
-		Schema: schema.ModuleSchema{ObjectTypes: []schema.ObjectType{{Name: "B"}}},
+		Schema: modSchema,
 	}, nil
 }
 
@@ -36,7 +44,13 @@ var testResolver = ModuleSetDecoderResolver(moduleSet)
 func TestModuleSetDecoderResolver_IterateAll(t *testing.T) {
 	objectTypes := map[string]bool{}
 	err := testResolver.IterateAll(func(moduleName string, cdc schema.ModuleCodec) error {
-		objectTypes[cdc.Schema.ObjectTypes[0].Name] = true
+		cdc.Schema.Types(func(t schema.Type) bool {
+			objTyp, ok := t.(schema.ObjectType)
+			if ok {
+				objectTypes[objTyp.Name] = true
+			}
+			return true
+		})
 		return nil
 	})
 	if err != nil {
@@ -66,8 +80,9 @@ func TestModuleSetDecoderResolver_LookupDecoder(t *testing.T) {
 		t.Fatalf("expected to find decoder for modA")
 	}
 
-	if decoder.Schema.ObjectTypes[0].Name != "A" {
-		t.Fatalf("expected object type A, got %s", decoder.Schema.ObjectTypes[0].Name)
+	_, ok := decoder.Schema.LookupType("A")
+	if !ok {
+		t.Fatalf("expected object type A")
 	}
 
 	decoder, found, err = testResolver.LookupDecoder("modB")
@@ -79,8 +94,9 @@ func TestModuleSetDecoderResolver_LookupDecoder(t *testing.T) {
 		t.Fatalf("expected to find decoder for modB")
 	}
 
-	if decoder.Schema.ObjectTypes[0].Name != "B" {
-		t.Fatalf("expected object type B, got %s", decoder.Schema.ObjectTypes[0].Name)
+	_, ok = decoder.Schema.LookupType("B")
+	if !ok {
+		t.Fatalf("expected object type B")
 	}
 
 	decoder, found, err = testResolver.LookupDecoder("modC")

schema/enum.go

@@ -18,6 +18,8 @@ type EnumDefinition struct {
 	Values []string
 }
 
+func (EnumDefinition) isType() {}
+
 // Validate validates the enum definition.
 func (e EnumDefinition) Validate() error {
 	if !ValidateName(e.Name) {
@@ -50,31 +52,3 @@ func (e EnumDefinition) ValidateValue(value string) error {
 	}
 	return fmt.Errorf("value %q is not a valid enum value for %s", value, e.Name)
 }
-
-// checkEnumCompatibility checks that the enum values are consistent across object types and fields.
-func checkEnumCompatibility(enumValueMap map[string]map[string]bool, field Field) error {
-	if field.Kind != EnumKind {
-		return nil
-	}
-
-	enum := field.EnumDefinition
-
-	if existing, ok := enumValueMap[enum.Name]; ok {
-		if len(existing) != len(enum.Values) {
-			return fmt.Errorf("enum %q has different number of values in different object types", enum.Name)
-		}
-
-		for _, value := range enum.Values {
-			if !existing[value] {
-				return fmt.Errorf("enum %q has different values in different object types", enum.Name)
-			}
-		}
-	} else {
-		valueMap := map[string]bool{}
-		for _, value := range enum.Values {
-			valueMap[value] = true
-		}
-		enumValueMap[enum.Name] = valueMap
-	}
-	return nil
-}

schema/module_schema.go

@@ -1,18 +1,82 @@
 package schema
 
-import "fmt"
+import (
+	"fmt"
+	"sort"
+)
 
 // ModuleSchema represents the logical schema of a module for purposes of indexing and querying.
 type ModuleSchema struct {
-	// ObjectTypes describe the types of objects that are part of the module's schema.
-	ObjectTypes []ObjectType
+	types map[string]Type
+}
+
+// NewModuleSchema constructs a new ModuleSchema and validates it. Any module schema returned without an error
+// is guaranteed to be valid.
+func NewModuleSchema(objectTypes []ObjectType) (ModuleSchema, error) {
+	types := map[string]Type{}
+
+	for _, objectType := range objectTypes {
+		types[objectType.Name] = objectType
+	}
+
+	res := ModuleSchema{types: types}
+
+	// validate adds all enum types to the type map
+	err := res.Validate()
+	if err != nil {
+		return ModuleSchema{}, err
+	}
+
+	return res, nil
+}
+
+func addEnumType(types map[string]Type, field Field) error {
+	enumDef := field.EnumDefinition
+	if enumDef.Name == "" {
+		return nil
+	}
+
+	existing, ok := types[enumDef.Name]
+	if !ok {
+		types[enumDef.Name] = enumDef
+		return nil
+	}
+
+	existingEnum, ok := existing.(EnumDefinition)
+	if !ok {
+		return fmt.Errorf("enum %q already exists as a different non-enum type", enumDef.Name)
+	}
+
+	if len(existingEnum.Values) != len(enumDef.Values) {
+		return fmt.Errorf("enum %q has different number of values in different fields", enumDef.Name)
+	}
+
+	existingValues := map[string]bool{}
+	for _, value := range existingEnum.Values {
+		existingValues[value] = true
+	}
+
+	for _, value := range enumDef.Values {
+		_, ok := existingValues[value]
+		if !ok {
+			return fmt.Errorf("enum %q has different values in different fields", enumDef.Name)
+		}
+	}
+
+	return nil
 }
 
 // Validate validates the module schema.
 func (s ModuleSchema) Validate() error {
-	enumValueMap := map[string]map[string]bool{}
-	for _, objType := range s.ObjectTypes {
-		if err := objType.validate(enumValueMap); err != nil {
+	for _, typ := range s.types {
+		objTyp, ok := typ.(ObjectType)
+		if !ok {
+			continue
+		}
+
+		// all enum types get added to the type map when we call ObjectType.validate
+		err := objTyp.validate(s.types)
+		if err != nil {
 			return err
 		}
 	}
@@ -22,10 +86,36 @@ func (s ModuleSchema) Validate() error {
 
 // ValidateObjectUpdate validates that the update conforms to the module schema.
 func (s ModuleSchema) ValidateObjectUpdate(update ObjectUpdate) error {
-	for _, objType := range s.ObjectTypes {
-		if objType.Name == update.TypeName {
-			return objType.ValidateObjectUpdate(update)
+	typ, ok := s.types[update.TypeName]
+	if !ok {
+		return fmt.Errorf("object type %q not found in module schema", update.TypeName)
+	}
+
+	objTyp, ok := typ.(ObjectType)
+	if !ok {
+		return fmt.Errorf("type %q is not an object type", update.TypeName)
+	}
+
+	return objTyp.ValidateObjectUpdate(update)
+}
+
+// LookupType looks up a type by name in the module schema.
+func (s ModuleSchema) LookupType(name string) (Type, bool) {
+	typ, ok := s.types[name]
+	return typ, ok
+}
+
+// Types calls the provided function for each type in the module schema and stops if the function returns false.
+// The types are iterated over in sorted order by name. This function is compatible with go 1.23 iterators.
+func (s ModuleSchema) Types(f func(Type) bool) {
+	keys := make([]string, 0, len(s.types))
+	for k := range s.types {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+	for _, k := range keys {
+		if !f(s.types[k]) {
+			break
 		}
 	}
-	return fmt.Errorf("object type %q not found in module schema", update.TypeName)
 }