TrueCloudLab/frostfs-api-go
22
0
Fork 15
Code Issues 10 Pull requests 1 Projects Releases Activity Actions

[#40] types: Generate StableMarshaler methods for refs

Browse source 
* Add plugin option for protogen in Makefile
* Fix the generator for the plugin in util/protogen
* Fix marshaler.go for refs

Signed-off-by: Airat Arifullin a.arifullin@yadro.com
This commit is contained in:
Airat Arifullin 2023-06-13 01:33:53 +03:00
parent 8266b31092
commit f0642d7f13
8 changed files with 343 additions and 191 deletions
Download patch file Download diff file Expand all files Collapse all files

6
util/proto/marshal.go
View file

@ -14,7 +14,7 @@ import (
)
type (
stableMarshaller interface {
StableMarshaller interface {
StableMarshal([]byte) []byte
StableSize() int
}
@ -302,7 +302,7 @@ func NestedStructurePrefix(field int64) (prefix uint64, ln int) {
return prefix, VarUIntSize(prefix)
}
func NestedStructureMarshal(field int64, buf []byte, v stableMarshaller) int {
func NestedStructureMarshal(field int64, buf []byte, v StableMarshaller) int {
if v == nil || reflect.ValueOf(v).IsNil() {
return 0
}
@ -317,7 +317,7 @@ func NestedStructureMarshal(field int64, buf []byte, v stableMarshaller) int {
return offset + n
}
func NestedStructureSize(field int64, v stableMarshaller) (size int) {
func NestedStructureSize(field int64, v StableMarshaller) (size int) {
if v == nil || reflect.ValueOf(v).IsNil() {
return 0
}

4
util/proto/test/test.pb.go generated
View file

@ -1,7 +1,7 @@
// Code generated by protoc-gen-go. DO NOT EDIT.
// versions:
// protoc-gen-go v1.28.1
// protoc v3.21.9
// protoc-gen-go v1.30.0
// protoc v3.12.4
// source: util/proto/test/test.proto
package test

110
util/protogen/main.go
View file

@ -8,14 +8,18 @@ import (
"google.golang.org/protobuf/reflect/protoreflect"
)
type MethodMask int
const (
Signature MethodMask = 1 << iota
)
func main() {
protogen.Options{}.Run(func(gen *protogen.Plugin) error {
for _, f := range gen.Files {
//if !f.Generate {
// continue
//}
imp := string(f.GoImportPath)
if strings.HasSuffix(imp, "/tree") || strings.HasSuffix(imp, "/control") {
impPath := f.GoImportPath.String()
if strings.HasSuffix(impPath, "/tree") || strings.HasSuffix(impPath, "/control") ||
strings.Contains(impPath, "/refs/") {
generateFile(gen, f)
}
}
@ -25,6 +29,11 @@ func main() {
// generateFile generates a *.pb.go file enforcing field-order serialization.
func generateFile(gen *protogen.Plugin, file *protogen.File) *protogen.GeneratedFile {
emitMask := MethodMask(0)
if imp := string(file.GoImportPath); strings.HasSuffix(imp, "/tree") || strings.HasSuffix(imp, "/control") {
emitMask |= Signature
}
filename := file.GeneratedFilenamePrefix + "_frostfs.pb.go"
g := gen.NewGeneratedFile(filename, file.GoImportPath)
g.P("// Code generated by protoc-gen-go-frostfs. DO NOT EDIT.")
@ -33,28 +42,14 @@ func generateFile(gen *protogen.Plugin, file *protogen.File) *protogen.Generated
g.P()
g.P(`import "git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/util/proto"`)
//for _, e := range file.Enums {
// g.P("type " + e.GoIdent.GoName + " int32")
// g.P("const (")
// for _, ev := range e.Values {
// g.P(ev.GoIdent.GoName, " = ", ev.Desc.Number())
// }
// g.P(")")
//}
for _, msg := range file.Messages {
emitMessage(g, msg)
emitMessage(g, msg, emitMask)
}
return g
}
func emitMessage(g *protogen.GeneratedFile, msg *protogen.Message) {
for _, inner := range msg.Messages {
emitMessage(g, inner)
}
func genStableSizeFunc(g *protogen.GeneratedFile, msg *protogen.Message) {
fs := sortFields(msg.Fields)
// StableSize implementation.
g.P("// StableSize returns the size of x in protobuf format.")
g.P("//")
g.P("// Structures with the same field values have the same binary size.")
@ -72,8 +67,10 @@ func emitMessage(g *protogen.GeneratedFile, msg *protogen.Message) {
}
g.P("return size")
g.P("}\n")
}
// StableMarshal implementation.
func genStableMarshalFunc(g *protogen.GeneratedFile, msg *protogen.Message) {
fs := sortFields(msg.Fields)
g.P("// StableMarshal marshals x in protobuf binary format with stable field order.")
g.P("//")
g.P("// If buffer length is less than x.StableSize(), new buffer is allocated.")
@ -93,32 +90,53 @@ func emitMessage(g *protogen.GeneratedFile, msg *protogen.Message) {
}
g.P("return buf")
g.P("}\n")
}
func genReadSignedDataFunc(g *protogen.GeneratedFile, msg *protogen.Message) {
g.P("// ReadSignedData fills buf with signed data of x.")
g.P("// If buffer length is less than x.SignedDataSize(), new buffer is allocated.")
g.P("//")
g.P("// Returns any error encountered which did not allow writing the data completely.")
g.P("// Otherwise, returns the buffer in which the data is written.")
g.P("//")
g.P("func (x *", msg.GoIdent.GoName, ") ReadSignedData(buf []byte) ([]byte, error) {")
g.P("return x.GetBody().StableMarshal(buf), nil")
g.P("}\n")
}
func genSignedDataSizeFunc(g *protogen.GeneratedFile, msg *protogen.Message) {
g.P("// SignedDataSize returns size of the request signed data in bytes.")
g.P("//")
g.P("// Structures with the same field values have the same signed data size.")
g.P("// Structures with the same field values have the same signed data.")
g.P("func (x *", msg.GoIdent.GoName, ") SignedDataSize() int {")
g.P("return x.GetBody().StableSize()")
g.P("}\n")
}
func genSetSignatureFunc(g *protogen.GeneratedFile, msg *protogen.Message) {
g.P("func (x *", msg.GoIdent.GoName, ") SetSignature(sig *Signature) {")
g.P("x.Signature = sig")
g.P("}\n")
}
func emitMessage(g *protogen.GeneratedFile, msg *protogen.Message, mask MethodMask) {
for _, inner := range msg.Messages {
emitMessage(g, inner, mask)
}
genStableSizeFunc(g, msg)
genStableMarshalFunc(g, msg)
if strings.HasSuffix(msg.GoIdent.GoName, "Request") || strings.HasSuffix(msg.GoIdent.GoName, "Response") {
// SignedDataSize implementation (only for requests and responses).
g.P("// ReadSignedData fills buf with signed data of x.")
g.P("// If buffer length is less than x.SignedDataSize(), new buffer is allocated.")
g.P("//")
g.P("// Returns any error encountered which did not allow writing the data completely.")
g.P("// Otherwise, returns the buffer in which the data is written.")
g.P("//")
g.P("// Structures with the same field values have the same signed data.")
g.P("func (x *", msg.GoIdent.GoName, ") SignedDataSize() int {")
g.P("return x.GetBody().StableSize()")
g.P("}\n")
// ReadSignedData implementation (only for requests and responses).
g.P("// SignedDataSize returns size of the request signed data in bytes.")
g.P("//")
g.P("// Structures with the same field values have the same signed data size.")
g.P("func (x *", msg.GoIdent.GoName, ") ReadSignedData(buf []byte) ([]byte, error) {")
g.P("return x.GetBody().StableMarshal(buf), nil")
g.P("}\n")
// Signature setters and getters.
g.P("func (x *", msg.GoIdent.GoName, ") SetSignature(sig *Signature) {")
g.P("x.Signature = sig")
g.P("}\n")
if mask&Signature != 0 {
genReadSignedDataFunc(g, msg)
genSignedDataSizeFunc(g, msg)
genSetSignatureFunc(g, msg)
}
}
}