add: support nested messages (#36)

* add: support nested messages * Update features/fastreflection/descriptors.go Co-authored-by: Aaron Craelius <aaron@regen.network> * chore: document better Co-authored-by: Aaron Craelius <aaron@regen.network>
cosmos · Nov 23, 2021 · 528f500 · 528f500
1 parent f005f0a
commit 528f500
Show file tree

Hide file tree

Showing 5 changed files with 1,226 additions and 1,198 deletions.
diff --git a/features/fastreflection/descriptors.go b/features/fastreflection/descriptors.go
@@ -3,6 +3,7 @@ package fastreflection
 import (
 	"fmt"
 	"github.com/cosmos/cosmos-proto/generator"
+	"google.golang.org/protobuf/reflect/protoreflect"
 
 	"github.com/cosmos/cosmos-proto/features/fastreflection/copied"
 	"google.golang.org/protobuf/compiler/protogen"
@@ -24,7 +25,14 @@ func (g *descGen) generate() {
 	g.P(")")
 	g.P("func init() {")
 	g.P(copied.InitFunctionName(g.file), "()")
-	g.P(messageDescriptorName(g.message), " = ", g.file.GoDescriptorIdent.GoName, ".Messages().ByName(\"", g.message.Desc.Name(), "\")")
+	parentMd, ok := g.message.Desc.Parent().(protoreflect.MessageDescriptor)
+	switch ok {
+	// case nested message 
+	case true:
+		g.P(messageDescriptorName(g.message), " = ", g.file.GoDescriptorIdent.GoName, ".Messages().ByName(\"", parentMd.Name(), "\").Messages().ByName(\"", g.message.Desc.Name(), "\")")
+	default:
+		g.P(messageDescriptorName(g.message), " = ", g.file.GoDescriptorIdent.GoName, ".Messages().ByName(\"", g.message.Desc.Name(), "\")")
+	}
 	for _, field := range g.message.Fields {
 		g.P(fieldDescriptorName(field), " = ", messageDescriptorName(g.message), ".Fields().ByName(\"", field.Desc.Name(), "\")")
 	}

diff --git a/features/fastreflection/proto_message.go b/features/fastreflection/proto_message.go
@@ -21,6 +21,18 @@ const (
 )
 
 func GenProtoMessage(f *protogen.File, g *generator.GeneratedFile, message *protogen.Message) {
+	genMessage(f, g, message)
+	// check for message declarations within a message declaration
+	for _, nested := range message.Messages {
+		// map entries are defines as messages, but we don't want to generate those.
+		if nested.Desc.IsMapEntry() {
+			continue
+		}
+		genMessage(f, g, nested)
+	}
+}
+
+func genMessage(f *protogen.File, g *generator.GeneratedFile, message *protogen.Message) {
 	gen := newGenerator(f, g, message)
 	gen.generateExtraTypes()
 	gen.generateReflectionType()

diff --git a/internal/testprotos/test3/fuzz_test.go b/internal/testprotos/test3/fuzz_test.go
@@ -29,7 +29,7 @@ func TestMarshalUnmarshal(t *testing.T) {
 
 // TestZeroValueOneofIsMarshalled tests that zero values in oneofs are marshalled
 func TestZeroValueOneofIsMarshalled(t *testing.T) {
-	msg1 := &TestAllTypes{OneofField: &TestAllTypes_OneofEnum{OneofEnum: NestedEnum_FOO}}
+	msg1 := &TestAllTypes{OneofField: &TestAllTypes_OneofEnum{OneofEnum: TestAllTypes_FOO}}
 	b, err := proto.Marshal(msg1)
 	require.NoError(t, err)
 

diff --git a/internal/testprotos/test3/test.proto b/internal/testprotos/test3/test.proto
@@ -10,21 +10,22 @@ import "internal/testprotos/test3/test_import.proto";
 
 option go_package = "github.com/cosmos/cosmos-proto/internal/testprotos/test3";
 
-// NestedMessage should eventually become nested when fast-reflection supports it
-message NestedMessage {
-  int32 a = 1;
-  TestAllTypes corecursive = 2;
-}
+message TestAllTypes {
 
-// NestedEnum should eventually become nested when fast-reflection supports it.
-enum NestedEnum {
-  FOO = 0;
-  BAR = 1;
-  BAZ = 2;
-  NEG = -1;  // Intentionally negative.
-}
+  // NestedMessage should eventually become nested when fast-reflection supports it
+  message NestedMessage {
+    int32 a = 1;
+    TestAllTypes corecursive = 2;
+  }
+
+  // NestedEnum should eventually become nested when fast-reflection supports it.
+  enum NestedEnum {
+    FOO = 0;
+    BAR = 1;
+    BAZ = 2;
+    NEG = -1;  // Intentionally negative.
+  }
 
-message TestAllTypes {
 
   int32         singular_int32    = 81;
   int64         singular_int64    = 82;