From 849de02bc3968ec69fda328d6a7c673bbbbd0928 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <e.stratonikov@yadro.com>
Date: Wed, 26 Jul 2023 17:52:54 +0300
Subject: [PATCH] [#49] util/proto: Calculate repeated field size without
 allocations

Signed-off-by: Evgenii Stratonikov <e.stratonikov@yadro.com>
---
 util/proto/marshal.go | 100 ++++++++++++------------------------------
 1 file changed, 27 insertions(+), 73 deletions(-)

diff --git a/util/proto/marshal.go b/util/proto/marshal.go
index ff0d30b..7bc8162 100644
--- a/util/proto/marshal.go
+++ b/util/proto/marshal.go
@@ -178,30 +178,13 @@ func RepeatedStringSize(field int, v []string) (size int) {
 	return size
 }
 
-func RepeatedUInt64Marshal(field int, buf []byte, v []uint64) int {
-	if len(v) == 0 {
-		return 0
-	}
-
-	prefix := field<<3 | 0x02
-	offset := binary.PutUvarint(buf, uint64(prefix))
-
-	_, arrSize := RepeatedUInt64Size(field, v)
-	offset += binary.PutUvarint(buf[offset:], uint64(arrSize))
-	for i := range v {
-		offset += binary.PutUvarint(buf[offset:], v[i])
-	}
-
-	return offset
-}
-
-func RepeatedUInt64Size(field int, v []uint64) (size, arraySize int) {
+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(v[i])
+		size += VarUIntSize(uint64(v[i]))
 	}
 	arraySize = size
 
@@ -213,82 +196,53 @@ func RepeatedUInt64Size(field int, v []uint64) (size, arraySize int) {
 	return size, arraySize
 }
 
-func RepeatedInt64Marshal(field int, buf []byte, v []int64) int {
+func repeatedUIntMarshal[T ~uint64 | ~int64 | ~uint32 | ~int32](field int, buf []byte, v []T) int {
 	if len(v) == 0 {
 		return 0
 	}
 
-	convert := make([]uint64, len(v))
+	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 {
-		convert[i] = uint64(v[i])
+		offset += binary.PutUvarint(buf[offset:], uint64(v[i]))
 	}
 
-	return RepeatedUInt64Marshal(field, buf, convert)
+	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) {
-	if len(v) == 0 {
-		return 0, 0
-	}
-
-	convert := make([]uint64, len(v))
-	for i := range v {
-		convert[i] = uint64(v[i])
-	}
-
-	return RepeatedUInt64Size(field, convert)
+	return repeatedUIntSize(field, v)
 }
 
 func RepeatedUInt32Marshal(field int, buf []byte, v []uint32) int {
-	if len(v) == 0 {
-		return 0
-	}
-
-	convert := make([]uint64, len(v))
-	for i := range v {
-		convert[i] = uint64(v[i])
-	}
-
-	return RepeatedUInt64Marshal(field, buf, convert)
+	return repeatedUIntMarshal(field, buf, v)
 }
 
 func RepeatedUInt32Size(field int, v []uint32) (size, arraySize int) {
-	if len(v) == 0 {
-		return 0, 0
-	}
-
-	convert := make([]uint64, len(v))
-	for i := range v {
-		convert[i] = uint64(v[i])
-	}
-
-	return RepeatedUInt64Size(field, convert)
+	return repeatedUIntSize(field, v)
 }
 
 func RepeatedInt32Marshal(field int, buf []byte, v []int32) int {
-	if len(v) == 0 {
-		return 0
-	}
-
-	convert := make([]uint64, len(v))
-	for i := range v {
-		convert[i] = uint64(v[i])
-	}
-
-	return RepeatedUInt64Marshal(field, buf, convert)
+	return repeatedUIntMarshal(field, buf, v)
 }
 
 func RepeatedInt32Size(field int, v []int32) (size, arraySize int) {
-	if len(v) == 0 {
-		return 0, 0
-	}
-
-	convert := make([]uint64, len(v))
-	for i := range v {
-		convert[i] = uint64(v[i])
-	}
-
-	return RepeatedUInt64Size(field, convert)
+	return repeatedUIntSize(field, v)
 }
 
 // VarUIntSize returns length of varint byte sequence for uint64 value 'x'.