[#40] types: Generate StableMarshaler methods for refs

* 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

Makefile

@@ -1,6 +1,8 @@
 #!/usr/bin/make -f
 SHELL = bash
 
+BIN = bin
+
 VERSION ?= $(shell git describe --tags --match "v*" --abbrev=8 --dirty --always)
 
 .PHONY: dep fmts fmt imports protoc test lint version help
@@ -47,6 +49,8 @@ protoc:
 		echo "⇒ Processing $$f "; \
 		protoc \
 			--proto_path=.:./vendor:/usr/local/include \
+			--plugin=protoc-gen-go-frostfs=$(BIN)/protogen \
+			--go-frostfs_out=. --go-frostfs_opt=paths=source_relative \
 			--go_out=. --go_opt=paths=source_relative \
 			--go-grpc_opt=require_unimplemented_servers=false \
 			--go-grpc_out=. --go-grpc_opt=paths=source_relative $$f; \

refs/marshal.go

@@ -30,25 +30,11 @@ const (
 )
 
 func (o *OwnerID) StableMarshal(buf []byte) []byte {
-	if o == nil {
-		return []byte{}
-	}
-
-	if buf == nil {
-		buf = make([]byte, o.StableSize())
-	}
-
-	proto.BytesMarshal(ownerIDValField, buf, o.val)
-
-	return buf
+	return message.StableMarshal[*refs.OwnerID](o, buf)
 }
 
 func (o *OwnerID) StableSize() int {
-	if o == nil {
-		return 0
-	}
-
-	return proto.BytesSize(ownerIDValField, o.val)
+	return message.StableSize[*refs.OwnerID](o)
 }
 
 func (o *OwnerID) Unmarshal(data []byte) error {
@@ -56,25 +42,11 @@ func (o *OwnerID) Unmarshal(data []byte) error {
 }
 
 func (c *ContainerID) StableMarshal(buf []byte) []byte {
-	if c == nil {
-		return []byte{}
-	}
-
-	if buf == nil {
-		buf = make([]byte, c.StableSize())
-	}
-
-	proto.BytesMarshal(containerIDValField, buf, c.val)
-
-	return buf
+	return message.StableMarshal[*refs.ContainerID](c, buf)
 }
 
 func (c *ContainerID) StableSize() int {
-	if c == nil {
-		return 0
-	}
-
-	return proto.BytesSize(containerIDValField, c.val)
+	return message.StableSize[*refs.ContainerID](c)
 }
 
 func (c *ContainerID) Unmarshal(data []byte) error {
@@ -82,22 +54,13 @@ func (c *ContainerID) Unmarshal(data []byte) error {
 }
 
 func (o *ObjectID) StableMarshal(buf []byte) []byte {
-	if o == nil {
-		return []byte{}
-	}
-
-	if buf == nil {
-		buf = make([]byte, o.StableSize())
-	}
-
-	proto.BytesMarshal(objectIDValField, buf, o.val)
-
-	return buf
+	return message.StableMarshal[*refs.ObjectID](o, buf)
 }
 
 // ObjectIDNestedListSize returns byte length of nested
 // repeated ObjectID field with fNum number.
 func ObjectIDNestedListSize(fNum int64, ids []ObjectID) (sz int) {
+	// TODO (aarifullin): remove this method when all marshalers are refactored
 	for i := range ids {
 		sz += proto.NestedStructureSize(fNum, &ids[i])
 	}
@@ -106,16 +69,13 @@ func ObjectIDNestedListSize(fNum int64, ids []ObjectID) (sz int) {
 }
 
 func (o *ObjectID) StableSize() int {
-	if o == nil {
-		return 0
-	}
-
-	return proto.BytesSize(objectIDValField, o.val)
+	return message.StableSize[*refs.ObjectID](o)
 }
 
 // ObjectIDNestedListMarshal writes protobuf repeated ObjectID field
 // with fNum number to buf.
 func ObjectIDNestedListMarshal(fNum int64, buf []byte, ids []ObjectID) (off int) {
+	// TODO (aarifullin): remove this method when all marshalers are refactored
 	prefix, _ := proto.NestedStructurePrefix(fNum)
 	for i := range ids {
 		off += binary.PutUvarint(buf[off:], prefix)
@@ -133,31 +93,11 @@ func (o *ObjectID) Unmarshal(data []byte) error {
 }
 
 func (a *Address) StableMarshal(buf []byte) []byte {
-	if a == nil {
-		return []byte{}
-	}
-
-	if buf == nil {
-		buf = make([]byte, a.StableSize())
-	}
-
-	var offset int
-
-	offset += proto.NestedStructureMarshal(addressContainerField, buf[offset:], a.cid)
-	proto.NestedStructureMarshal(addressObjectField, buf[offset:], a.oid)
-
-	return buf
+	return message.StableMarshal[*refs.Address](a, buf)
 }
 
 func (a *Address) StableSize() (size int) {
-	if a == nil {
-		return 0
-	}
-
-	size += proto.NestedStructureSize(addressContainerField, a.cid)
-	size += proto.NestedStructureSize(addressObjectField, a.oid)
-
-	return size
+	return message.StableSize[*refs.Address](a)
 }
 
 func (a *Address) Unmarshal(data []byte) error {
@@ -165,31 +105,11 @@ func (a *Address) Unmarshal(data []byte) error {
 }
 
 func (c *Checksum) StableMarshal(buf []byte) []byte {
-	if c == nil {
-		return []byte{}
-	}
-
-	if buf == nil {
-		buf = make([]byte, c.StableSize())
-	}
-
-	var offset int
-
-	offset += proto.EnumMarshal(checksumTypeField, buf[offset:], int32(c.typ))
-	proto.BytesMarshal(checksumValueField, buf[offset:], c.sum)
-
-	return buf
+	return message.StableMarshal[*refs.Checksum](c, buf)
 }
 
 func (c *Checksum) StableSize() (size int) {
-	if c == nil {
-		return 0
-	}
-
-	size += proto.EnumSize(checksumTypeField, int32(c.typ))
-	size += proto.BytesSize(checksumValueField, c.sum)
-
-	return size
+	return message.StableSize[*refs.Checksum](c)
 }
 
 func (c *Checksum) Unmarshal(data []byte) error {
@@ -197,33 +117,11 @@ func (c *Checksum) Unmarshal(data []byte) error {
 }
 
 func (s *Signature) StableMarshal(buf []byte) []byte {
-	if s == nil {
-		return []byte{}
-	}
-
-	if buf == nil {
-		buf = make([]byte, s.StableSize())
-	}
-
-	var offset int
-
-	offset += proto.BytesMarshal(signatureKeyField, buf[offset:], s.key)
-	offset += proto.BytesMarshal(signatureValueField, buf[offset:], s.sign)
-	proto.EnumMarshal(signatureSchemeField, buf[offset:], int32(s.scheme))
-
-	return buf
+	return message.StableMarshal[*refs.Signature](s, buf)
 }
 
 func (s *Signature) StableSize() (size int) {
-	if s == nil {
-		return 0
-	}
-
-	size += proto.BytesSize(signatureKeyField, s.key)
-	size += proto.BytesSize(signatureValueField, s.sign)
-	size += proto.EnumSize(signatureSchemeField, int32(s.scheme))
-
-	return size
+	return message.StableSize[*refs.Signature](s)
 }
 
 func (s *Signature) Unmarshal(data []byte) error {
@@ -231,31 +129,11 @@ func (s *Signature) Unmarshal(data []byte) error {
 }
 
 func (v *Version) StableMarshal(buf []byte) []byte {
-	if v == nil {
-		return []byte{}
-	}
-
-	if buf == nil {
-		buf = make([]byte, v.StableSize())
-	}
-
-	var offset int
-
-	offset += proto.UInt32Marshal(versionMajorField, buf[offset:], v.major)
-	proto.UInt32Marshal(versionMinorField, buf[offset:], v.minor)
-
-	return buf
+	return message.StableMarshal[*refs.Version](v, buf)
 }
 
 func (v *Version) StableSize() (size int) {
-	if v == nil {
-		return 0
-	}
-
-	size += proto.UInt32Size(versionMajorField, v.major)
-	size += proto.UInt32Size(versionMinorField, v.minor)
-
-	return size
+	return message.StableSize[*refs.Version](v)
 }
 
 func (v *Version) Unmarshal(data []byte) error {

rpc/message/encoding.go

@@ -2,6 +2,7 @@ package message
 
 import (
 	"git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/rpc/grpc"
+	protoutil "git.frostfs.info/TrueCloudLab/frostfs-api-go/v2/util/proto"
 	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/proto"
 )
@@ -46,3 +47,13 @@ func UnmarshalJSON(m Message, data []byte, gm GRPCConvertedMessage) error {
 
 	return m.FromGRPCMessage(gm)
 }
+
+// StableMarshal encodes message 'm' using the invocation from related gRPC message
+func StableMarshal[M protoutil.StableMarshaller](m Message, buf []byte) []byte {
+	return m.ToGRPCMessage().(M).StableMarshal(buf)
+}
+
+// StableSize encodes message 'm' using the invocation from related gRPC message
+func StableSize[M protoutil.StableMarshaller](m Message) int {
+	return m.ToGRPCMessage().(M).StableSize()
+}

util/proto/marshal.go

@@ -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
 	}

util/protogen/main.go

@@ -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)
+		}
 	}
 }