protogen: Support unpacked repeated uint64 fields #91
4 changed files with 92 additions and 6 deletions
|
@ -1,12 +1,14 @@
|
||||||
package proto_test
|
package proto_test
|
||||||
|
|
||||||
import (
|
import (
|
||||||
|
"encoding/binary"
|
||||||
"math"
|
"math"
|
||||||
"testing"
|
"testing"
|
||||||
|
|
||||||
"git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/util/proto"
|
"git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/util/proto"
|
||||||
"git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/util/proto/test"
|
"git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/util/proto/test"
|
||||||
"github.com/stretchr/testify/require"
|
"github.com/stretchr/testify/require"
|
||||||
|
"google.golang.org/protobuf/encoding/protowire"
|
||||||
goproto "google.golang.org/protobuf/proto"
|
goproto "google.golang.org/protobuf/proto"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
@ -33,6 +35,8 @@ type stableRepPrimitives struct {
|
||||||
FieldD []uint32
|
FieldD []uint32
|
||||||
FieldE []int64
|
FieldE []int64
|
||||||
FieldF []uint64
|
FieldF []uint64
|
||||||
|
|
||||||
|
FieldFu []uint64
|
||||||
}
|
}
|
||||||
|
|
||||||
const (
|
const (
|
||||||
|
@ -185,6 +189,20 @@ func (s *stableRepPrimitives) stableMarshal(buf []byte, wrongField bool) ([]byte
|
||||||
}
|
}
|
||||||
i += proto.RepeatedUInt64Marshal(fieldNum, buf, s.FieldF)
|
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
|
return buf, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@ -196,7 +214,12 @@ func (s *stableRepPrimitives) stableSize() int {
|
||||||
f5, _ := proto.RepeatedInt64Size(5, s.FieldE)
|
f5, _ := proto.RepeatedInt64Size(5, s.FieldE)
|
||||||
f6, _ := proto.RepeatedUInt64Size(6, s.FieldF)
|
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) {
|
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) {
|
func TestFixed64Marshal(t *testing.T) {
|
||||||
t.Run("zero", func(t *testing.T) {
|
t.Run("zero", func(t *testing.T) {
|
||||||
testFixed64Marshal(t, 0, false)
|
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) {
|
func testFixed64Marshal(t *testing.T, n uint64, wrongField bool) {
|
||||||
var (
|
var (
|
||||||
custom = stablePrimitives{FieldI: n}
|
custom = stablePrimitives{FieldI: n}
|
||||||
|
|
BIN
util/proto/test/test.pb.go
generated
BIN
util/proto/test/test.pb.go
generated
Binary file not shown.
|
@ -31,4 +31,5 @@ message RepPrimitives {
|
||||||
repeated uint32 field_d = 4;
|
repeated uint32 field_d = 4;
|
||||||
repeated int64 field_e = 5;
|
repeated int64 field_e = 5;
|
||||||
repeated uint64 field_f = 6;
|
repeated uint64 field_f = 6;
|
||||||
|
repeated uint64 field_fu = 7 [ packed = false ];
|
||||||
}
|
}
|
||||||
|
|
|
@ -8,6 +8,11 @@ import (
|
||||||
"google.golang.org/protobuf/reflect/protoreflect"
|
"google.golang.org/protobuf/reflect/protoreflect"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
var (
|
||||||
|
protowirePackage = protogen.GoImportPath("google.golang.org/protobuf/encoding/protowire")
|
||||||
|
binaryPackage = protogen.GoImportPath("encoding/binary")
|
||||||
|
)
|
||||||
|
|
||||||
func main() {
|
func main() {
|
||||||
protogen.Options{}.Run(func(gen *protogen.Plugin) error {
|
protogen.Options{}.Run(func(gen *protogen.Plugin) error {
|
||||||
for _, f := range gen.Files {
|
for _, f := range gen.Files {
|
||||||
|
@ -62,7 +67,7 @@ func emitMessage(g *protogen.GeneratedFile, msg *protogen.Message) {
|
||||||
g.P("if x == nil { return 0 }")
|
g.P("if x == nil { return 0 }")
|
||||||
if len(fs) != 0 {
|
if len(fs) != 0 {
|
||||||
for _, f := range fs {
|
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")
|
g.P("var n int")
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
|
@ -140,9 +145,19 @@ func emitFieldSize(g *protogen.GeneratedFile, f *protogen.Field) {
|
||||||
}
|
}
|
||||||
|
|
||||||
switch {
|
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("for i := range ", name, "{")
|
||||||
|
if f.Desc.Kind() == protoreflect.MessageKind {
|
||||||
g.P("size += proto.NestedStructureSize(", f.Desc.Number(), ", ", name, "[i])")
|
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("}")
|
g.P("}")
|
||||||
case f.Desc.IsList():
|
case f.Desc.IsList():
|
||||||
if m.RepeatedDouble {
|
if m.RepeatedDouble {
|
||||||
|
@ -177,9 +192,22 @@ func emitFieldMarshal(g *protogen.GeneratedFile, f *protogen.Field) {
|
||||||
prefix = "Repeated" + m.Prefix
|
prefix = "Repeated" + m.Prefix
|
||||||
}
|
}
|
||||||
switch {
|
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("for i := range ", name, "{")
|
||||||
|
if f.Desc.Kind() == protoreflect.MessageKind {
|
||||||
g.P("offset += proto.NestedStructureMarshal(", f.Desc.Number(), ", buf[offset:], ", name, "[i])")
|
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("}")
|
g.P("}")
|
||||||
case f.Desc.IsList():
|
case f.Desc.IsList():
|
||||||
g.P("offset += proto.Repeated", m.Prefix, "Marshal(", f.Desc.Number(), ", buf[offset:], ", name, ")")
|
g.P("offset += proto.Repeated", m.Prefix, "Marshal(", f.Desc.Number(), ", buf[offset:], ", name, ")")
|
||||||
|
|
Loading…
Reference in a new issue