features/pool: reuse map type fields

Makefile

@@ -60,6 +60,7 @@ gen-testproto: get-grpc-testproto gen-wkt-testproto install
 		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/proto3opt/opt.proto \
 		testproto/proto2/scalars.proto \

features/pool/pool.go

@@ -68,6 +68,21 @@ 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() {
+			switch field.Desc.MapValue().Kind() {
+			case protoreflect.MessageKind, protoreflect.GroupKind:
+				valueField := field.Message.Fields[1]
+				p.P(`for _, v := range m.`, fieldName, `{`)
+				if p.ShouldPool(valueField.Message) {
+					p.P(`v.ResetVT()`)
+				} else {
+					p.P(`v.Reset()`)
+				}
+				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 {`)

testproto/pool/pool_test.go

@@ -186,3 +186,22 @@ func Test_Pool_Optional(t *testing.T) {
 	mFromPool := MemoryPoolExtensionFromVTPool()
 	require.True(t, mFromPool.EqualVT(&MemoryPoolExtension{}))
 }
+
+func Test_Pool_Reuse_Map(t *testing.T) {
+	allocs := testing.AllocsPerRun(10, func() {
+		obj := MapReuseTest1FromVTPool()
+		if obj.GetIntToStringMap() == nil {
+			obj.IntToStringMap = make(map[int32]string)
+		}
+		obj.IntToStringMap[1] = "test1"
+		obj.IntToStringMap[2] = "test2"
+		obj.IntToStringMap[3] = "test3"
+		if obj.GetStringToEnumMap() == nil {
+			obj.StringToEnumMap = make(map[string]MapReuseTest1_TEST)
+		}
+		obj.StringToEnumMap["test1"] = MapReuseTest1_test1
+		obj.StringToEnumMap["test2"] = MapReuseTest1_test2
+		obj.ReturnToVTPool()
+	})
+	require.Less(t, int(allocs), 1)
+}