feat(pool): adds pooling support for well-known type wrappers and maps

Makefile

@@ -41,6 +41,7 @@ gen-wkt: bin/protoc-gen-go-vtproto
 		--plugin protoc-gen-go-vtproto="${GOBIN}/protoc-gen-go-vtproto" \
 		--go-vtproto_out=. \
 		--go-vtproto_opt=module=google.golang.org/protobuf,wrap=true \
+		--go-vtproto_opt=pool=google.golang.org/protobuf/types/known/*.* \
 		$(PROTOBUF_ROOT)/src/google/protobuf/any.proto \
         $(PROTOBUF_ROOT)/src/google/protobuf/duration.proto \
         $(PROTOBUF_ROOT)/src/google/protobuf/empty.proto \
@@ -55,12 +56,15 @@ gen-testproto: get-grpc-testproto gen-wkt-testproto install
 		--proto_path=include \
 		--go_out=. --plugin protoc-gen-go="${GOBIN}/protoc-gen-go" \
 		--go-vtproto_out=allow-empty=true:. --plugin protoc-gen-go-vtproto="${GOBIN}/protoc-gen-go-vtproto" \
+		--go-vtproto_opt=pool=google.golang.org/protobuf/types/known/*.* \
 		-I$(PROTOBUF_ROOT)/src \
 		testproto/ignore_unknown_fields/opt.proto \
 		testproto/empty/empty.proto \
 		testproto/pool/pool.proto \
 		testproto/pool/pool_with_slice_reuse.proto \
+		testproto/pool/pool_with_map_reuse.proto \
 		testproto/pool/pool_with_oneof.proto \
+		testproto/pool/wkt_pool.proto \
 		testproto/proto3opt/opt.proto \
 		testproto/proto2/scalars.proto \
 		testproto/unsafe/unsafe.proto \

features/pool/pool.go

@@ -39,9 +39,9 @@ func (p *pool) message(message *protogen.Message) {
 	}
 
 	p.once = true
-	ccTypeName := message.GoIdent
+	ccTypeName := p.PoolMessageIdent(message)
 
-	p.P(`var vtprotoPool_`, ccTypeName, ` = `, p.Ident("sync", "Pool"), `{`)
+	p.P(`var vtprotoPool_`, ccTypeName.GoName, ` = `, p.Ident("sync", "Pool"), `{`)
 	p.P(`New: func() interface{} {`)
 	p.P(`return &`, ccTypeName, `{}`)
 	p.P(`},`)
@@ -58,7 +58,7 @@ func (p *pool) message(message *protogen.Message) {
 			case protoreflect.MessageKind, protoreflect.GroupKind:
 				p.P(`for _, mm := range m.`, fieldName, `{`)
 				if p.ShouldPool(field.Message) {
-					p.P(`mm.ResetVT()`)
+					p.PResetVTElement("mm", field.Message)
 				} else {
 					p.P(`mm.Reset()`)
 				}
@@ -68,17 +68,30 @@ 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.Desc.IsMap() {
+			// Pooled message values go back to their pool; other values are dropped
+			// by clear. The map itself is kept so its buckets are reused.
+			if valueField := field.Message.Fields[1]; valueField.Message != nil && p.ShouldPool(valueField.Message) {
+				p.P(`for _, v := range m.`, fieldName, `{`)
+				p.P(`if v != nil {`)
+				p.PReturnToVTPool("v", valueField.Message)
+				p.P(`}`)
+				p.P(`}`)
+			}
+			p.P(`clear(m.`, fieldName, `)`)
+			p.P(fmt.Sprintf("f%d", len(saved)), ` := m.`, fieldName)
+			saved = append(saved, field)
 		} 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()`)
+				p.PReturnToVTPool("oneof."+fieldName, field.Message)
 				p.P(`}`)
 			}
 		} else {
 			switch field.Desc.Kind() {
 			case protoreflect.MessageKind, protoreflect.GroupKind:
-				if !field.Desc.IsMap() && p.ShouldPool(field.Message) {
-					p.P(`m.`, fieldName, `.ReturnToVTPool()`)
+				if p.ShouldPool(field.Message) {
+					p.PReturnToVTPool("m."+fieldName, field.Message)
 				}
 			case protoreflect.BytesKind:
 				p.P(fmt.Sprintf("f%d", len(saved)), ` := m.`, fieldName, `[:0]`)
@@ -87,7 +100,11 @@ func (p *pool) message(message *protogen.Message) {
 		}
 	}
 
-	p.P(`m.Reset()`)
+	if p.IsLocalWrapper(message) {
+		p.P(`(*`, p.QualifiedGoIdent(message.GoIdent), `)(m).Reset()`)
+	} else {
+		p.P(`m.Reset()`)
+	}
 	for i, field := range saved {
 		p.P(`m.`, field.GoName, ` = `, fmt.Sprintf("f%d", i))
 	}
@@ -97,11 +114,11 @@ func (p *pool) message(message *protogen.Message) {
 	p.P(`func (m *`, ccTypeName, `) ReturnToVTPool() {`)
 	p.P(`if m != nil {`)
 	p.P(`m.ResetVT()`)
-	p.P(`vtprotoPool_`, ccTypeName, `.Put(m)`)
+	p.P(`vtprotoPool_`, ccTypeName.GoName, `.Put(m)`)
 	p.P(`}`)
 	p.P(`}`)
 
-	p.P(`func `, ccTypeName, `FromVTPool() *`, ccTypeName, `{`)
-	p.P(`return vtprotoPool_`, ccTypeName, `.Get().(*`, ccTypeName, `)`)
+	p.P(`func `, ccTypeName.GoName, `FromVTPool() *`, ccTypeName, `{`)
+	p.P(`return vtprotoPool_`, ccTypeName.GoName, `.Get().(*`, ccTypeName, `)`)
 	p.P(`}`)
 }

features/unmarshal/unmarshal.go

@@ -157,7 +157,7 @@ func (p *unmarshal) declareMapField(varName string, nullable bool, field *protog
 	}
 }
 
-func (p *unmarshal) mapField(varName string, field *protogen.Field, unique bool) {
+func (p *unmarshal) mapField(varName string, field *protogen.Field, unique, pool bool) {
 	switch field.Desc.Kind() {
 	case protoreflect.DoubleKind:
 		p.P(`var `, varName, `temp uint64`)
@@ -226,7 +226,11 @@ func (p *unmarshal) mapField(varName string, field *protogen.Field, unique bool)
 		p.P(`return `, p.Ident("io", `ErrUnexpectedEOF`))
 		p.P(`}`)
 		buf := `dAtA[iNdEx:postmsgIndex]`
-		p.P(varName, ` = &`, p.noStarOrSliceType(field), `{}`)
+		if pool {
+			p.PAssignFromVTPool(varName, field.Message)
+		} else {
+			p.P(varName, ` = &`, p.noStarOrSliceType(field), `{}`)
+		}
 		p.decodeMessage(varName, buf, field.Message)
 		p.P(`iNdEx = postmsgIndex`)
 	case protoreflect.BytesKind:
@@ -499,12 +503,16 @@ func (p *unmarshal) fieldItem(field *protogen.Field, fieldname string, message *
 			p.decodeMessage("oneof."+field.GoName, buf, field.Message)
 			p.P(`} else {`)
 			if p.ShouldPool(message) && p.ShouldPool(field.Message) {
-				p.P(`v := `, msgname, `FromVTPool()`)
+				p.PDeclFromVTPool("v", field.Message)
 			} else {
 				p.P(`v := &`, msgname, `{}`)
 			}
 			p.decodeMessage("v", buf, field.Message)
-			p.P(`m.`, fieldname, ` = &`, field.GoIdent, "{", field.GoName, `: v}`)
+			if p.IsWellKnownType(field.Message) && p.IsLocalWrapper(field.Message) {
+				p.P(`m.`, fieldname, ` = &`, field.GoIdent, "{", field.GoName, `: (*`, p.QualifiedGoIdent(field.Message.GoIdent), `)(v)}`)
+			} else {
+				p.P(`m.`, fieldname, ` = &`, field.GoIdent, "{", field.GoName, `: v}`)
+			}
 			p.P(`}`)
 		} else if field.Desc.IsMap() {
 			unique := proto.GetExtension(field.Desc.Options(), vtproto.E_Options).(*vtproto.Opts).GetUnique()
@@ -527,9 +535,9 @@ func (p *unmarshal) fieldItem(field *protogen.Field, fieldname string, message *
 			p.P(`fieldNum := int32(wire >> 3)`)
 
 			p.P(`if fieldNum == 1 {`)
-			p.mapField("mapkey", field.Message.Fields[0], unique)
+			p.mapField("mapkey", field.Message.Fields[0], unique, false)
 			p.P(`} else if fieldNum == 2 {`)
-			p.mapField("mapvalue", field.Message.Fields[1], unique)
+			p.mapField("mapvalue", field.Message.Fields[1], unique, p.ShouldPool(message) && p.ShouldPool(field.Message.Fields[1].Message))
 			p.P(`} else {`)
 			p.P(`iNdEx = entryPreIndex`)
 			p.P(`skippy, err := `, p.Helper("Skip"), `(dAtA[iNdEx:])`)
@@ -565,7 +573,7 @@ func (p *unmarshal) fieldItem(field *protogen.Field, fieldname string, message *
 		} else {
 			p.P(`if m.`, fieldname, ` == nil {`)
 			if p.ShouldPool(message) && p.ShouldPool(field.Message) {
-				p.P(`m.`, fieldname, ` = `, field.Message.GoIdent, `FromVTPool()`)
+				p.PAssignFromVTPool("m."+fieldname, field.Message)
 			} else {
 				p.P(`m.`, fieldname, ` = &`, field.Message.GoIdent, `{}`)
 			}

generator/generatedfile.go

@@ -16,8 +16,9 @@ import (
 
 type GeneratedFile struct {
 	*protogen.GeneratedFile
-	Config        *Config
-	LocalPackages map[protoreflect.FullName]bool
+	Config             *Config
+	LocalPackages      map[protoreflect.FullName]bool
+	OutputGoImportPath protogen.GoImportPath
 }
 
 func (p *GeneratedFile) Ident(path, ident string) string {
@@ -52,16 +53,79 @@ func (b *GeneratedFile) ShouldIgnoreUnknownFields(message *protogen.Message) boo
 }
 
 func (b *GeneratedFile) Alloc(vname string, message *protogen.Message, isQualifiedIdent bool) {
+	if b.ShouldPool(message) {
+		b.PDeclFromVTPool(vname, message)
+		return
+	}
+
 	ident := message.GoIdent.GoName
 	if isQualifiedIdent {
 		ident = b.QualifiedGoIdent(message.GoIdent)
 	}
+	b.P(vname, " := new(", ident, `)`)
+}
 
-	if b.ShouldPool(message) {
-		b.P(vname, " := ", ident, `FromVTPool()`)
-	} else {
-		b.P(vname, " := new(", ident, `)`)
+// PoolMessageIdent is the type used for pool method receivers in the file being generated.
+func (p *GeneratedFile) PoolMessageIdent(message *protogen.Message) protogen.GoIdent {
+	if p.IsLocalWrapper(message) {
+		return protogen.GoIdent{GoName: message.GoIdent.GoName}
+	}
+	return message.GoIdent
+}
+
+// PDeclFromVTPool emits `vname := <pooled message>`. Caller must ensure ShouldPool(message).
+func (p *GeneratedFile) PDeclFromVTPool(vname string, message *protogen.Message) {
+	p.P(append([]any{vname, " := "}, p.fromVTPool(message)...)...)
+}
+
+// PAssignFromVTPool emits `target = <pooled message>`. Caller must ensure ShouldPool(message).
+func (p *GeneratedFile) PAssignFromVTPool(target string, message *protogen.Message) {
+	expr := p.fromVTPool(message)
+	if p.IsLocalWrapper(message) {
+		// The wrapper shares its layout with the underlying message, which is the type targets hold.
+		expr = p.convertExpr(message, expr)
+	}
+	p.P(append([]any{target, " = "}, expr...)...)
+}
+
+// convertExpr wraps expr in a pointer conversion to the message's own Go type.
+func (p *GeneratedFile) convertExpr(message *protogen.Message, expr []any) []any {
+	out := []any{"(*", p.QualifiedGoIdent(message.GoIdent), ")("}
+	out = append(out, expr...)
+	return append(out, ")")
+}
+
+func (p *GeneratedFile) fromVTPool(message *protogen.Message) []any {
+	if p.IsWellKnownType(message) && !p.IsLocalWrapper(message) {
+		wkt := p.WellKnownTypeMap(message)
+		fromPool := protogen.GoIdent{GoName: wkt.GoName + "FromVTPool", GoImportPath: wkt.GoImportPath}
+		return []any{"(*", p.QualifiedGoIdent(message.GoIdent), ")((*", p.QualifiedGoIdent(wkt), ")(", p.QualifiedGoIdent(fromPool), "()))"}
+	}
+	if p.IsLocalWrapper(message) || message.GoIdent.GoImportPath == p.OutputGoImportPath {
+		return []any{p.PoolMessageIdent(message).GoName, "FromVTPool()"}
+	}
+	fromPool := protogen.GoIdent{GoName: message.GoIdent.GoName + "FromVTPool", GoImportPath: message.GoIdent.GoImportPath}
+	return []any{p.QualifiedGoIdent(fromPool), "()"}
+}
+
+// PReturnToVTPool emits a statement that returns a nested message to its pool.
+func (p *GeneratedFile) PReturnToVTPool(expr string, message *protogen.Message) {
+	if p.IsWellKnownType(message) {
+		wkt := p.WellKnownTypeMap(message)
+		p.P(`(*`, p.QualifiedGoIdent(wkt), `)(`, expr, `).ReturnToVTPool()`)
+		return
+	}
+	p.P(expr, `.ReturnToVTPool()`)
+}
+
+// PResetVTElement emits a reset for a repeated element.
+func (p *GeneratedFile) PResetVTElement(expr string, message *protogen.Message) {
+	if p.IsWellKnownType(message) {
+		wkt := p.WellKnownTypeMap(message)
+		p.P(`(*`, p.QualifiedGoIdent(wkt), `)(`, expr, `).ResetVT()`)
+		return
 	}
+	p.P(expr, `.ResetVT()`)
 }
 
 func (p *GeneratedFile) FieldGoType(field *protogen.Field) (goType string, pointer bool) {
@@ -206,3 +270,8 @@ func (p *GeneratedFile) WellKnownTypeMap(message *protogen.Message) protogen.GoI
 func (p *GeneratedFile) Wrapper() bool {
 	return p.Config.Wrap
 }
+
+// IsLocalWrapper is true when wrap=true and the message belongs to a locally generated package.
+func (p *GeneratedFile) IsLocalWrapper(message *protogen.Message) bool {
+	return p.Wrapper() && p.IsLocalMessage(message)
+}

generator/generator.go

@@ -114,9 +114,10 @@ func (gen *Generator) Generate() {
 
 func (gen *Generator) generateFile(gf *protogen.GeneratedFile, file *protogen.File) {
 	p := &GeneratedFile{
-		GeneratedFile: gf,
-		Config:        gen.cfg,
-		LocalPackages: gen.local,
+		GeneratedFile:      gf,
+		Config:             gen.cfg,
+		LocalPackages:      gen.local,
+		OutputGoImportPath: file.GoImportPath,
 	}
 
 	if p.Config.BuildTag != "" {