[#91] protogen: Support unpacked repeated uint64 fields

Signed-off-by: Evgenii Stratonikov <e.stratonikov@yadro.com>
This commit is contained in:
Evgenii Stratonikov 2024-07-15 11:16:29 +03:00
parent 3639563d80
commit f517e39491
4 changed files with 113 additions and 17 deletions

View file

@ -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, ")")