frostfs-api-go/util/proto/marshal.go
Evgenii Stratonikov 7a5ee927c8 [#49] util/proto: Do not allocate in StringSize()
It was not catched by the test because most of the time the function is
inlined. However, I've seen it allocating with pprof in one of the
earlier builds.

Signed-off-by: Evgenii Stratonikov <e.stratonikov@yadro.com>
2023-07-26 18:52:59 +03:00

368 lines
7.7 KiB
Go

/*
This package contains help functions for stable marshaller. Their usage is
totally optional. One can implement fast stable marshaller without these
runtime function calls.
*/
package proto
import (
"encoding/binary"
"math"
"math/bits"
)
type (
stableMarshaller interface {
StableMarshal([]byte) []byte
StableSize() int
}
)
func BytesMarshal(field int, buf, v []byte) int {
return bytesMarshal(field, buf, v)
}
func BytesSize(field int, v []byte) int {
return bytesSize(field, v)
}
func bytesMarshal[T ~[]byte | ~string](field int, buf []byte, v T) int {
if len(v) == 0 {
return 0
}
return bytesMarshalNoCheck(field, buf, v)
}
func bytesMarshalNoCheck[T ~[]byte | ~string](field int, buf []byte, v T) int {
prefix := field<<3 | 0x2
// buf length check can prevent panic at PutUvarint, but it will make
// marshaller a bit slower.
i := binary.PutUvarint(buf, uint64(prefix))
i += binary.PutUvarint(buf[i:], uint64(len(v)))
i += copy(buf[i:], v)
return i
}
func bytesSize[T ~[]byte | ~string](field int, v T) int {
if len(v) == 0 {
return 0
}
return bytesSizeNoCheck(field, v)
}
func bytesSizeNoCheck[T ~[]byte | ~string](field int, v T) int {
prefix := field<<3 | 0x2
return VarUIntSize(uint64(prefix)) + VarUIntSize(uint64(len(v))) + len(v)
}
func StringMarshal(field int, buf []byte, v string) int {
return bytesMarshal(field, buf, v)
}
func StringSize(field int, v string) int {
return bytesSize(field, v)
}
func BoolMarshal(field int, buf []byte, v bool) int {
if !v {
return 0
}
prefix := field << 3
// buf length check can prevent panic at PutUvarint, but it will make
// marshaller a bit slower.
i := binary.PutUvarint(buf, uint64(prefix))
buf[i] = 0x1
return i + 1
}
func BoolSize(field int, v bool) int {
if !v {
return 0
}
prefix := field << 3
return VarUIntSize(uint64(prefix)) + 1 // bool is always 1 byte long
}
func UInt64Marshal(field int, buf []byte, v uint64) int {
if v == 0 {
return 0
}
prefix := field << 3
// buf length check can prevent panic at PutUvarint, but it will make
// marshaller a bit slower.
i := binary.PutUvarint(buf, uint64(prefix))
i += binary.PutUvarint(buf[i:], v)
return i
}
func UInt64Size(field int, v uint64) int {
if v == 0 {
return 0
}
prefix := field << 3
return VarUIntSize(uint64(prefix)) + VarUIntSize(v)
}
func Int64Marshal(field int, buf []byte, v int64) int {
return UInt64Marshal(field, buf, uint64(v))
}
func Int64Size(field int, v int64) int {
return UInt64Size(field, uint64(v))
}
func UInt32Marshal(field int, buf []byte, v uint32) int {
return UInt64Marshal(field, buf, uint64(v))
}
func UInt32Size(field int, v uint32) int {
return UInt64Size(field, uint64(v))
}
func Int32Marshal(field int, buf []byte, v int32) int {
return UInt64Marshal(field, buf, uint64(v))
}
func Int32Size(field int, v int32) int {
return UInt64Size(field, uint64(v))
}
func EnumMarshal(field int, buf []byte, v int32) int {
return UInt64Marshal(field, buf, uint64(v))
}
func EnumSize(field int, v int32) int {
return UInt64Size(field, uint64(v))
}
func RepeatedBytesMarshal(field int, buf []byte, v [][]byte) int {
var offset int
for i := range v {
offset += bytesMarshalNoCheck(field, buf[offset:], v[i])
}
return offset
}
func RepeatedBytesSize(field int, v [][]byte) (size int) {
for i := range v {
size += bytesSizeNoCheck(field, v[i])
}
return size
}
func RepeatedStringMarshal(field int, buf []byte, v []string) int {
var offset int
for i := range v {
offset += bytesMarshalNoCheck(field, buf[offset:], v[i])
}
return offset
}
func RepeatedStringSize(field int, v []string) (size int) {
for i := range v {
size += bytesSizeNoCheck(field, v[i])
}
return size
}
func repeatedUIntSize[T ~uint64 | ~int64 | ~uint32 | ~int32](field int, v []T) (size, arraySize int) {
if len(v) == 0 {
return 0, 0
}
for i := range v {
size += VarUIntSize(uint64(v[i]))
}
arraySize = size
size += VarUIntSize(uint64(size))
prefix := field<<3 | 0x2
size += VarUIntSize(uint64(prefix))
return size, arraySize
}
func repeatedUIntMarshal[T ~uint64 | ~int64 | ~uint32 | ~int32](field int, buf []byte, v []T) int {
if len(v) == 0 {
return 0
}
prefix := field<<3 | 0x02
offset := binary.PutUvarint(buf, uint64(prefix))
_, arrSize := repeatedUIntSize(field, v)
offset += binary.PutUvarint(buf[offset:], uint64(arrSize))
for i := range v {
offset += binary.PutUvarint(buf[offset:], uint64(v[i]))
}
return offset
}
func RepeatedUInt64Marshal(field int, buf []byte, v []uint64) int {
return repeatedUIntMarshal(field, buf, v)
}
func RepeatedUInt64Size(field int, v []uint64) (size, arraySize int) {
return repeatedUIntSize(field, v)
}
func RepeatedInt64Marshal(field int, buf []byte, v []int64) int {
return repeatedUIntMarshal(field, buf, v)
}
func RepeatedInt64Size(field int, v []int64) (size, arraySize int) {
return repeatedUIntSize(field, v)
}
func RepeatedUInt32Marshal(field int, buf []byte, v []uint32) int {
return repeatedUIntMarshal(field, buf, v)
}
func RepeatedUInt32Size(field int, v []uint32) (size, arraySize int) {
return repeatedUIntSize(field, v)
}
func RepeatedInt32Marshal(field int, buf []byte, v []int32) int {
return repeatedUIntMarshal(field, buf, v)
}
func RepeatedInt32Size(field int, v []int32) (size, arraySize int) {
return repeatedUIntSize(field, v)
}
// VarUIntSize returns length of varint byte sequence for uint64 value 'x'.
func VarUIntSize(x uint64) int {
return (bits.Len64(x|1) + 6) / 7
}
func NestedStructurePrefix(field int64) (prefix uint64, ln int) {
prefix = uint64(field<<3 | 0x02)
return prefix, VarUIntSize(prefix)
}
func NestedStructureMarshal[T stableMarshaller](field int64, buf []byte, v T) int {
n := v.StableSize()
if n == 0 {
return 0
}
prefix, _ := NestedStructurePrefix(field)
offset := binary.PutUvarint(buf, prefix)
offset += binary.PutUvarint(buf[offset:], uint64(n))
v.StableMarshal(buf[offset:])
return offset + n
}
func NestedStructureSize[T stableMarshaller](field int64, v T) (size int) {
n := v.StableSize()
if n == 0 {
return 0
}
_, ln := NestedStructurePrefix(field)
size = ln + VarUIntSize(uint64(n)) + n
return size
}
func Fixed64Marshal(field int, buf []byte, v uint64) int {
if v == 0 {
return 0
}
prefix := field<<3 | 1
// buf length check can prevent panic at PutUvarint, but it will make
// marshaller a bit slower.
i := binary.PutUvarint(buf, uint64(prefix))
binary.LittleEndian.PutUint64(buf[i:], v)
return i + 8
}
func Fixed64Size(fNum int, v uint64) int {
if v == 0 {
return 0
}
prefix := fNum<<3 | 1
return VarUIntSize(uint64(prefix)) + 8
}
func Float64Marshal(field int, buf []byte, v float64) int {
if v == 0 {
return 0
}
prefix := field<<3 | 1
i := binary.PutUvarint(buf, uint64(prefix))
binary.LittleEndian.PutUint64(buf[i:], math.Float64bits(v))
return i + 8
}
func Float64Size(fNum int, v float64) int {
if v == 0 {
return 0
}
prefix := fNum<<3 | 1
return VarUIntSize(uint64(prefix)) + 8
}
// Fixed32Marshal encodes uint32 value to Protocol Buffers fixed32 field with specified number,
// and writes it to specified buffer. Returns number of bytes written.
//
// Panics if the buffer is undersized.
func Fixed32Marshal(field int, buf []byte, v uint32) int {
if v == 0 {
return 0
}
prefix := field<<3 | 5
// buf length check can prevent panic at PutUvarint, but it will make
// marshaller a bit slower.
i := binary.PutUvarint(buf, uint64(prefix))
binary.LittleEndian.PutUint32(buf[i:], v)
return i + 4
}
// Fixed32Size returns number of bytes required to encode uint32 value to Protocol Buffers fixed32 field
// with specified number.
func Fixed32Size(fNum int, v uint32) int {
if v == 0 {
return 0
}
prefix := fNum<<3 | 5
return VarUIntSize(uint64(prefix)) + 4
}