diff --git a/util/proto/marshal.go b/util/proto/marshal.go index e345492b..b1f55a37 100644 --- a/util/proto/marshal.go +++ b/util/proto/marshal.go @@ -20,13 +20,21 @@ type ( ) func BytesMarshal(field int, buf, v []byte) int { - if len(v) == 0 { - return 0 - } return bytesMarshal(field, buf, v) } -func bytesMarshal(field int, buf, v []byte) int { +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 @@ -38,26 +46,25 @@ func bytesMarshal(field int, buf, v []byte) int { return i } -func BytesSize(field int, v []byte) int { - ln := len(v) - if ln == 0 { +func bytesSize[T ~[]byte | ~string](field int, v T) int { + if len(v) == 0 { return 0 } - return bytesSize(field, v) + return bytesSizeNoCheck(field, v) } -func bytesSize(field int, v []byte) int { +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, []byte(v)) + return bytesMarshal(field, buf, v) } func StringSize(field int, v string) int { - return BytesSize(field, []byte(v)) + return bytesSize(field, v) } func BoolMarshal(field int, buf []byte, v bool) int { @@ -146,7 +153,7 @@ func RepeatedBytesMarshal(field int, buf []byte, v [][]byte) int { var offset int for i := range v { - offset += bytesMarshal(field, buf[offset:], v[i]) + offset += bytesMarshalNoCheck(field, buf[offset:], v[i]) } return offset @@ -154,7 +161,7 @@ func RepeatedBytesMarshal(field int, buf []byte, v [][]byte) int { func RepeatedBytesSize(field int, v [][]byte) (size int) { for i := range v { - size += bytesSize(field, v[i]) + size += bytesSizeNoCheck(field, v[i]) } return size @@ -164,7 +171,7 @@ func RepeatedStringMarshal(field int, buf []byte, v []string) int { var offset int for i := range v { - offset += bytesMarshal(field, buf[offset:], []byte(v[i])) + offset += bytesMarshalNoCheck(field, buf[offset:], v[i]) } return offset @@ -172,7 +179,7 @@ func RepeatedStringMarshal(field int, buf []byte, v []string) int { func RepeatedStringSize(field int, v []string) (size int) { for i := range v { - size += bytesSize(field, []byte(v[i])) + size += bytesSizeNoCheck(field, v[i]) } return size