[#91] protogen: Support unpacked repeated uint64 fields

Signed-off-by: Evgenii Stratonikov <e.stratonikov@yadro.com>

util/proto/marshal_test.go

@@ -1,12 +1,14 @@
 package proto_test
 
 import (
+	"encoding/binary"
 	"math"
 	"testing"
 
 	"git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/util/proto"
 	"git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/util/proto/test"
 	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/encoding/protowire"
 	goproto "google.golang.org/protobuf/proto"
 )
 
@@ -33,6 +35,8 @@ type stableRepPrimitives struct {
 	FieldD []uint32
 	FieldE []int64
 	FieldF []uint64
+
+	FieldFu []uint64
 }
 
 const (
@@ -185,6 +189,20 @@ func (s *stableRepPrimitives) stableMarshal(buf []byte, wrongField bool) ([]byte
 	}
 	i += proto.RepeatedUInt64Marshal(fieldNum, buf, s.FieldF)
 
+	fieldNum = 7
+	if wrongField {
+		fieldNum++
+	}
+	for j := range s.FieldFu {
+		{
+			prefix := protowire.EncodeTag(
+				protowire.Number(fieldNum),
+				protowire.VarintType)
+			i += binary.PutUvarint(buf[i:], uint64(prefix))
+			i += binary.PutUvarint(buf[i:], s.FieldFu[j])
+		}
+	}
+
 	return buf, nil
 }
 
@@ -196,7 +214,12 @@ func (s *stableRepPrimitives) stableSize() int {
 	f5, _ := proto.RepeatedInt64Size(5, s.FieldE)
 	f6, _ := proto.RepeatedUInt64Size(6, s.FieldF)
 
-	return f1 + f2 + f3 + f4 + f5 + f6
+	var f7 int
+	for i := range s.FieldFu {
+		f7 += protowire.SizeGroup(protowire.Number(7), protowire.SizeVarint(s.FieldFu[i]))
+	}
+
+	return f1 + f2 + f3 + f4 + f5 + f6 + f7
 }
 
 func TestBytesMarshal(t *testing.T) {
@@ -404,6 +427,22 @@ func TestRepeatedUInt64Marshal(t *testing.T) {
 	})
 }
 
+func TestRepeatedUInt64MarshalUnpacked(t *testing.T) {
+	t.Run("not empty", func(t *testing.T) {
+		data := []uint64{0, 1, 2, 3, 4, 5}
+		testRepeatedUInt64MarshalUnpacked(t, data, false)
+		testRepeatedUInt64MarshalUnpacked(t, data, true)
+	})
+
+	t.Run("empty", func(t *testing.T) {
+		testRepeatedUInt64MarshalUnpacked(t, []uint64{}, false)
+	})
+
+	t.Run("nil", func(t *testing.T) {
+		testRepeatedUInt64MarshalUnpacked(t, nil, false)
+	})
+}
+
 func TestFixed64Marshal(t *testing.T) {
 	t.Run("zero", func(t *testing.T) {
 		testFixed64Marshal(t, 0, false)
@@ -738,6 +777,24 @@ func testRepeatedUInt64Marshal(t *testing.T, n []uint64, wrongField bool) {
 	}
 }
 
+func testRepeatedUInt64MarshalUnpacked(t *testing.T, n []uint64, wrongField bool) {
+	var (
+		custom    = stableRepPrimitives{FieldFu: n}
+		transport = test.RepPrimitives{FieldFu: n}
+	)
+
+	result := testRepMarshal(t, custom, &transport, wrongField)
+
+	if !wrongField {
+		require.Len(t, result.FieldFu, len(n))
+		if len(n) > 0 {
+			require.Equal(t, n, result.FieldFu)
+		}
+	} else {
+		require.Len(t, result.FieldFu, 0)
+	}
+}
+
 func testFixed64Marshal(t *testing.T, n uint64, wrongField bool) {
 	var (
 		custom    = stablePrimitives{FieldI: n}

util/proto/test/test.pb.go

@@ -1,7 +1,7 @@
 // Code generated by protoc-gen-go. DO NOT EDIT.
 // versions:
 // 	protoc-gen-go v1.33.0
-// 	protoc        v4.25.3
+// 	protoc        v5.27.2
 // source: util/proto/test/test.proto
 
 package test
@@ -201,12 +201,13 @@ type RepPrimitives struct {
 	sizeCache     protoimpl.SizeCache
 	unknownFields protoimpl.UnknownFields
 
-	FieldA [][]byte `protobuf:"bytes,1,rep,name=field_a,json=fieldA,proto3" json:"field_a,omitempty"`
+	FieldA  [][]byte `protobuf:"bytes,1,rep,name=field_a,json=fieldA,proto3" json:"field_a,omitempty"`
-	FieldB []string `protobuf:"bytes,2,rep,name=field_b,json=fieldB,proto3" json:"field_b,omitempty"`
+	FieldB  []string `protobuf:"bytes,2,rep,name=field_b,json=fieldB,proto3" json:"field_b,omitempty"`
-	FieldC []int32  `protobuf:"varint,3,rep,packed,name=field_c,json=fieldC,proto3" json:"field_c,omitempty"`
+	FieldC  []int32  `protobuf:"varint,3,rep,packed,name=field_c,json=fieldC,proto3" json:"field_c,omitempty"`
-	FieldD []uint32 `protobuf:"varint,4,rep,packed,name=field_d,json=fieldD,proto3" json:"field_d,omitempty"`
+	FieldD  []uint32 `protobuf:"varint,4,rep,packed,name=field_d,json=fieldD,proto3" json:"field_d,omitempty"`
-	FieldE []int64  `protobuf:"varint,5,rep,packed,name=field_e,json=fieldE,proto3" json:"field_e,omitempty"`
+	FieldE  []int64  `protobuf:"varint,5,rep,packed,name=field_e,json=fieldE,proto3" json:"field_e,omitempty"`
-	FieldF []uint64 `protobuf:"varint,6,rep,packed,name=field_f,json=fieldF,proto3" json:"field_f,omitempty"`
+	FieldF  []uint64 `protobuf:"varint,6,rep,packed,name=field_f,json=fieldF,proto3" json:"field_f,omitempty"`
+	FieldFu []uint64 `protobuf:"varint,7,rep,name=field_fu,json=fieldFu,proto3" json:"field_fu,omitempty"`
 }
 
 func (x *RepPrimitives) Reset() {
@@ -283,6 +284,13 @@ func (x *RepPrimitives) GetFieldF() []uint64 {
 	return nil
 }
 
+func (x *RepPrimitives) GetFieldFu() []uint64 {
+	if x != nil {
+		return x.FieldFu
+	}
+	return nil
+}
+
 var File_util_proto_test_test_proto protoreflect.FileDescriptor
 
 var file_util_proto_test_test_proto_rawDesc = []byte{
@@ -312,7 +320,7 @@ var file_util_proto_test_test_proto_rawDesc = []byte{
 	0x45, 0x6e, 0x75, 0x6d, 0x12, 0x0b, 0x0a, 0x07, 0x55, 0x4e, 0x4b, 0x4e, 0x4f, 0x57, 0x4e, 0x10,
 	0x00, 0x12, 0x0c, 0x0a, 0x08, 0x50, 0x4f, 0x53, 0x49, 0x54, 0x49, 0x56, 0x45, 0x10, 0x01, 0x12,
 	0x15, 0x0a, 0x08, 0x4e, 0x45, 0x47, 0x41, 0x54, 0x49, 0x56, 0x45, 0x10, 0xff, 0xff, 0xff, 0xff,
-	0xff, 0xff, 0xff, 0xff, 0xff, 0x01, 0x22, 0xa5, 0x01, 0x0a, 0x0d, 0x52, 0x65, 0x70, 0x50, 0x72,
+	0xff, 0xff, 0xff, 0xff, 0xff, 0x01, 0x22, 0xc4, 0x01, 0x0a, 0x0d, 0x52, 0x65, 0x70, 0x50, 0x72,
 	0x69, 0x6d, 0x69, 0x74, 0x69, 0x76, 0x65, 0x73, 0x12, 0x17, 0x0a, 0x07, 0x66, 0x69, 0x65, 0x6c,
 	0x64, 0x5f, 0x61, 0x18, 0x01, 0x20, 0x03, 0x28, 0x0c, 0x52, 0x06, 0x66, 0x69, 0x65, 0x6c, 0x64,
 	0x41, 0x12, 0x17, 0x0a, 0x07, 0x66, 0x69, 0x65, 0x6c, 0x64, 0x5f, 0x62, 0x18, 0x02, 0x20, 0x03,
@@ -322,9 +330,11 @@ var file_util_proto_test_test_proto_rawDesc = []byte{
 	0x20, 0x03, 0x28, 0x0d, 0x52, 0x06, 0x66, 0x69, 0x65, 0x6c, 0x64, 0x44, 0x12, 0x17, 0x0a, 0x07,
 	0x66, 0x69, 0x65, 0x6c, 0x64, 0x5f, 0x65, 0x18, 0x05, 0x20, 0x03, 0x28, 0x03, 0x52, 0x06, 0x66,
 	0x69, 0x65, 0x6c, 0x64, 0x45, 0x12, 0x17, 0x0a, 0x07, 0x66, 0x69, 0x65, 0x6c, 0x64, 0x5f, 0x66,
-	0x18, 0x06, 0x20, 0x03, 0x28, 0x04, 0x52, 0x06, 0x66, 0x69, 0x65, 0x6c, 0x64, 0x46, 0x42, 0x11,
+	0x18, 0x06, 0x20, 0x03, 0x28, 0x04, 0x52, 0x06, 0x66, 0x69, 0x65, 0x6c, 0x64, 0x46, 0x12, 0x1d,
-	0x5a, 0x0f, 0x75, 0x74, 0x69, 0x6c, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2f, 0x74, 0x65, 0x73,
+	0x0a, 0x08, 0x66, 0x69, 0x65, 0x6c, 0x64, 0x5f, 0x66, 0x75, 0x18, 0x07, 0x20, 0x03, 0x28, 0x04,
-	0x74, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
+	0x42, 0x02, 0x10, 0x00, 0x52, 0x07, 0x66, 0x69, 0x65, 0x6c, 0x64, 0x46, 0x75, 0x42, 0x11, 0x5a,
+	0x0f, 0x75, 0x74, 0x69, 0x6c, 0x2f, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x2f, 0x74, 0x65, 0x73, 0x74,
+	0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
 }
 
 var (

util/proto/test/test.proto

@@ -31,4 +31,5 @@ message RepPrimitives {
   repeated uint32 field_d = 4;
   repeated int64 field_e = 5;
   repeated uint64 field_f = 6;
+  repeated uint64 field_fu = 7 [ packed = false ];
 }

util/protogen/main.go

@@ -8,6 +8,11 @@ import (
 	"google.golang.org/protobuf/reflect/protoreflect"
 )
 
+var (
+	protowirePackage = protogen.GoImportPath("google.golang.org/protobuf/encoding/protowire")
+	binaryPackage    = protogen.GoImportPath("encoding/binary")
+)
+
 func main() {
 	protogen.Options{}.Run(func(gen *protogen.Plugin) error {
 		for _, f := range gen.Files {
@@ -62,7 +67,7 @@ func emitMessage(g *protogen.GeneratedFile, msg *protogen.Message) {
 	g.P("if x == nil { return 0 }")
 	if len(fs) != 0 {
 		for _, f := range fs {
-			if f.Desc.IsList() && marshalers[f.Desc.Kind()].RepeatedDouble {
+			if f.Desc.IsList() && marshalers[f.Desc.Kind()].RepeatedDouble && !(f.Desc.Kind() == protoreflect.Uint64Kind && !f.Desc.IsPacked()) {
 				g.P("var n int")
 				break
 			}
@@ -140,9 +145,19 @@ func emitFieldSize(g *protogen.GeneratedFile, f *protogen.Field) {
 	}
 
 	switch {
-	case f.Desc.IsList() && f.Desc.Kind() == protoreflect.MessageKind:
+	case f.Desc.IsList() && (f.Desc.Kind() == protoreflect.MessageKind || f.Desc.Kind() == protoreflect.Uint64Kind && !f.Desc.IsPacked()):
 		g.P("for i := range ", name, "{")
-		g.P("size += proto.NestedStructureSize(", f.Desc.Number(), ", ", name, "[i])")
+		if f.Desc.Kind() == protoreflect.MessageKind {
+			g.P("size += proto.NestedStructureSize(", f.Desc.Number(), ", ", name, "[i])")
+		} else {
+			if f.Desc.Kind() != protoreflect.Uint64Kind {
+				panic("only uint64 unpacked primitive is supported")
+			}
+
+			g.P("size += ", protowirePackage.Ident("SizeGroup"), "(",
+				protowirePackage.Ident("Number"), "(", f.Desc.Number(), "), ",
+				protowirePackage.Ident("SizeVarint"), "(", name, "[i]))")
+		}
 		g.P("}")
 	case f.Desc.IsList():
 		if m.RepeatedDouble {
@@ -177,9 +192,22 @@ func emitFieldMarshal(g *protogen.GeneratedFile, f *protogen.Field) {
 		prefix = "Repeated" + m.Prefix
 	}
 	switch {
-	case f.Desc.IsList() && f.Desc.Kind() == protoreflect.MessageKind:
+	case f.Desc.IsList() && (f.Desc.Kind() == protoreflect.MessageKind || f.Desc.Kind() == protoreflect.Uint64Kind && !f.Desc.IsPacked()):
 		g.P("for i := range ", name, "{")
-		g.P("offset += proto.NestedStructureMarshal(", f.Desc.Number(), ", buf[offset:], ", name, "[i])")
+		if f.Desc.Kind() == protoreflect.MessageKind {
+			g.P("offset += proto.NestedStructureMarshal(", f.Desc.Number(), ", buf[offset:], ", name, "[i])")
+		} else {
+			if f.Desc.Kind() != protoreflect.Uint64Kind {
+				panic("only uint64 unpacked primitive is supported")
+			}
+			g.P("{")
+			g.P("prefix := ", protowirePackage.Ident("EncodeTag"), "(",
+				protowirePackage.Ident("Number"), "(", f.Desc.Number(), "), ",
+				protowirePackage.Ident("VarintType"), ")")
+			g.P("offset += ", binaryPackage.Ident("PutUvarint"), "(buf[offset:], uint64(prefix))")
+			g.P("offset += ", binaryPackage.Ident("PutUvarint"), "(buf[offset:], ", name, "[i])")
+			g.P("}")
+		}
 		g.P("}")
 	case f.Desc.IsList():
 		g.P("offset += proto.Repeated", m.Prefix, "Marshal(", f.Desc.Number(), ", buf[offset:], ", name, ")")