Merge pull request #111 from evgfedotov/fix/optional-message

fix: check oneof on syntetic
planetscale · Nov 21, 2023 · a8e005f · a8e005f
2 parents a19b16b + 0b07b23
commit a8e005f
Show file tree

Hide file tree

Showing 5 changed files with 447 additions and 21 deletions.
diff --git a/features/pool/pool.go b/features/pool/pool.go
@@ -68,7 +68,7 @@ func (p *pool) message(message *protogen.Message) {
 			}
 			p.P(fmt.Sprintf("f%d", len(saved)), ` := m.`, fieldName, `[:0]`)
 			saved = append(saved, field)
-		} else if field.Oneof != nil {
+		} else if field.Oneof != nil && !field.Oneof.Desc.IsSynthetic() {
 			if p.ShouldPool(field.Message) {
 				p.P(`if oneof, ok := m.`, field.Oneof.GoName, `.(*`, field.GoIdent, `); ok {`)
 				p.P(`oneof.`, fieldName, `.ReturnToVTPool()`)

diff --git a/testproto/pool/pool.pb.go b/testproto/pool/pool.pb.go
diff --git a/testproto/pool/pool.proto b/testproto/pool/pool.proto
@@ -3,8 +3,13 @@ option go_package = "testproto/pool";
 
 import "github.com/planetscale/vtprotobuf/vtproto/ext.proto";
 
+message OptionalMessage{
+  option (vtproto.mempool) = true;
+}
+
 message MemoryPoolExtension {
   option (vtproto.mempool) = true;
   string foo1 = 1;
   uint64 foo2 = 2;
+  optional OptionalMessage foo3 = 3;
 }
diff --git a/testproto/pool/pool_test.go b/testproto/pool/pool_test.go
@@ -6,6 +6,8 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"google.golang.org/protobuf/proto"
 )
 
 func Test_Pool_slice_data_override(t *testing.T) {
@@ -163,3 +165,24 @@ func Test_Pool_Oneof(t *testing.T) {
 	require.NoError(t, t8.UnmarshalVT(t4Bytes))
 	require.Equal(t, &t4, &t8)
 }
+
+func Test_Pool_Optional(t *testing.T) {
+	m := &MemoryPoolExtension{
+		Foo1: "foo1",
+		Foo2: 123,
+		Foo3: &OptionalMessage{},
+	}
+
+	mBytes, err := m.MarshalVT()
+	require.NoError(t, err)
+
+	mUnmarshal := &MemoryPoolExtension{}
+	err = proto.Unmarshal(mBytes, mUnmarshal)
+	require.NoError(t, err)
+
+	require.True(t, m.EqualVT(mUnmarshal))
+
+	m.ReturnToVTPool()
+	mFromPool := MemoryPoolExtensionFromVTPool()
+	require.True(t, mFromPool.EqualVT(&MemoryPoolExtension{}))
+}