From f7def4dde17d150c79e1c089bcf858d461830076 Mon Sep 17 00:00:00 2001
From: Stephen Buttolph <stephen@avalabs.org>
Date: Tue, 2 Apr 2024 13:47:56 -0400
Subject: [PATCH] Improve performance of marshalling small keys (#2895)

---
 x/merkledb/codec.go      |  8 ++--
 x/merkledb/codec_test.go | 94 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+), 4 deletions(-)

diff --git a/x/merkledb/codec.go b/x/merkledb/codec.go
index af144317bc4..ba13a97a1b5 100644
--- a/x/merkledb/codec.go
+++ b/x/merkledb/codec.go
@@ -362,10 +362,10 @@ func decodeID(src *bytes.Reader) (ids.ID, error) {
 }
 
 func encodeKey(key Key) []byte {
-	estimatedLen := binary.MaxVarintLen64 + len(key.Bytes())
-	dst := bytes.NewBuffer(make([]byte, 0, estimatedLen))
-	encodeKeyToBuffer(dst, key)
-	return dst.Bytes()
+	keyLen := uintSize(uint64(key.length)) + len(key.Bytes())
+	buf := bytes.NewBuffer(make([]byte, 0, keyLen))
+	encodeKeyToBuffer(buf, key)
+	return buf.Bytes()
 }
 
 func encodeKeyToBuffer(dst *bytes.Buffer, key Key) {
diff --git a/x/merkledb/codec_test.go b/x/merkledb/codec_test.go
index 19aa4d55ed4..811b1d6b60e 100644
--- a/x/merkledb/codec_test.go
+++ b/x/merkledb/codec_test.go
@@ -374,6 +374,81 @@ var (
 			},
 		},
 	}
+	encodeKeyTests = []struct {
+		name          string
+		key           Key
+		expectedBytes []byte
+	}{
+		{
+			name: "empty",
+			key:  ToKey([]byte{}),
+			expectedBytes: []byte{
+				0x00, // length
+			},
+		},
+		{
+			name: "1 byte",
+			key:  ToKey([]byte{0}),
+			expectedBytes: []byte{
+				0x08, // length
+				0x00, // key
+			},
+		},
+		{
+			name: "2 bytes",
+			key:  ToKey([]byte{0, 1}),
+			expectedBytes: []byte{
+				0x10,       // length
+				0x00, 0x01, // key
+			},
+		},
+		{
+			name: "4 bytes",
+			key:  ToKey([]byte{0, 1, 2, 3}),
+			expectedBytes: []byte{
+				0x20,                   // length
+				0x00, 0x01, 0x02, 0x03, // key
+			},
+		},
+		{
+			name: "8 bytes",
+			key:  ToKey([]byte{0, 1, 2, 3, 4, 5, 6, 7}),
+			expectedBytes: []byte{
+				0x40,                                           // length
+				0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, // key
+			},
+		},
+		{
+			name: "32 bytes",
+			key:  ToKey(make([]byte, 32)),
+			expectedBytes: append(
+				[]byte{
+					0x80, 0x02, // length
+				},
+				make([]byte, 32)..., // key
+			),
+		},
+		{
+			name: "64 bytes",
+			key:  ToKey(make([]byte, 64)),
+			expectedBytes: append(
+				[]byte{
+					0x80, 0x04, // length
+				},
+				make([]byte, 64)..., // key
+			),
+		},
+		{
+			name: "1024 bytes",
+			key:  ToKey(make([]byte, 1024)),
+			expectedBytes: append(
+				[]byte{
+					0x80, 0x40, // length
+				},
+				make([]byte, 1024)..., // key
+			),
+		},
+	}
 )
 
 func FuzzCodecBool(f *testing.F) {
@@ -609,6 +684,15 @@ func TestEncodeDBNode(t *testing.T) {
 	}
 }
 
+func TestEncodeKey(t *testing.T) {
+	for _, test := range encodeKeyTests {
+		t.Run(test.name, func(t *testing.T) {
+			bytes := encodeKey(test.key)
+			require.Equal(t, test.expectedBytes, bytes)
+		})
+	}
+}
+
 func TestCodecDecodeKeyLengthOverflowRegression(t *testing.T) {
 	_, err := decodeKey(binary.AppendUvarint(nil, math.MaxInt))
 	require.ErrorIs(t, err, io.ErrUnexpectedEOF)
@@ -654,6 +738,16 @@ func Benchmark_EncodeDBNode(b *testing.B) {
 	}
 }
 
+func Benchmark_EncodeKey(b *testing.B) {
+	for _, benchmark := range encodeKeyTests {
+		b.Run(benchmark.name, func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				encodeKey(benchmark.key)
+			}
+		})
+	}
+}
+
 func Benchmark_EncodeUint(b *testing.B) {
 	var dst bytes.Buffer
 	dst.Grow(binary.MaxVarintLen64)