From 7a5ee927c8a272ff97f26cd8b7aa65a4c9d6d161 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <e.stratonikov@yadro.com>
Date: Wed, 26 Jul 2023 18:51:37 +0300
Subject: [PATCH] [#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>
---
 util/proto/marshal.go | 37 ++++++++++++++++++++++---------------
 1 file changed, 22 insertions(+), 15 deletions(-)

diff --git a/util/proto/marshal.go b/util/proto/marshal.go
index e345492..b1f55a3 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