fix: ResetVT oneof

planetscale · Nov 15, 2023 · 93e4d89 · 93e4d89
1 parent 5a1a54c
commit 93e4d89
Show file tree

Hide file tree

Showing 4 changed files with 178 additions and 19 deletions.
diff --git a/features/pool/pool.go b/features/pool/pool.go
@@ -45,7 +45,7 @@ func (p *pool) message(message *protogen.Message) {
 
 	p.P(`var vtprotoPool_`, ccTypeName, ` = `, p.Ident("sync", "Pool"), `{`)
 	p.P(`New: func() interface{} {`)
-	p.P(`return &`, message.GoIdent, `{}`)
+	p.P(`return &`, ccTypeName, `{}`)
 	p.P(`},`)
 	p.P(`}`)
 
@@ -68,10 +68,12 @@ 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 && p.ShouldPool(message) && p.ShouldPool(field.Message) {
-			p.P(`if oneof, ok := m.`, field.Oneof.GoName, `.(*`, field.GoIdent, `); ok {`)
-			p.P(`oneof.`, fieldName, `.ReturnToVTPool()`)
-			p.P(`}`)
+		} else if field.Oneof != nil {
+			if p.ShouldPool(field.Message) {
+				p.P(`if oneof, ok := m.`, field.Oneof.GoName, `.(*`, field.GoIdent, `); ok {`)
+				p.P(`oneof.`, fieldName, `.ReturnToVTPool()`)
+				p.P(`}`)
+			}
 		} else {
 			switch field.Desc.Kind() {
 			case protoreflect.MessageKind, protoreflect.GroupKind:

diff --git a/testproto/pool/pool_with_oneof.pb.go b/testproto/pool/pool_with_oneof.pb.go
diff --git a/testproto/pool/pool_with_oneof.proto b/testproto/pool/pool_with_oneof.proto
@@ -31,5 +31,6 @@ message OneofTest {
     Test1 test1 = 1;
     Test2 test2 = 2;
     Test3 test3 = 3;
+    bytes test4 = 4;
   }
 }
diff --git a/testproto/pool/pool_with_oneof_vtproto.pb.go b/testproto/pool/pool_with_oneof_vtproto.pb.go