From de5eb918a3209fecf9ef0110a2fb9f0304aafaf8 Mon Sep 17 00:00:00 2001
From: Stephen Buttolph <stephen@avalabs.org>
Date: Thu, 28 Mar 2024 19:06:44 -0400
Subject: [PATCH] Optimize encodeUint

---
 x/merkledb/codec.go      | 22 +++++-----------------
 x/merkledb/codec_test.go | 17 +++++++++++++++++
 2 files changed, 22 insertions(+), 17 deletions(-)

diff --git a/x/merkledb/codec.go b/x/merkledb/codec.go
index a5d4a922b0d..3161fd16c86 100644
--- a/x/merkledb/codec.go
+++ b/x/merkledb/codec.go
@@ -11,7 +11,6 @@ import (
 	"math"
 	"math/bits"
 	"slices"
-	"sync"
 
 	"golang.org/x/exp/maps"
 
@@ -75,22 +74,12 @@ type decoder interface {
 }
 
 func newCodec() encoderDecoder {
-	return &codecImpl{
-		varIntPool: sync.Pool{
-			New: func() interface{} {
-				return make([]byte, binary.MaxVarintLen64)
-			},
-		},
-	}
+	return &codecImpl{}
 }
 
 // Note that bytes.Buffer.Write always returns nil, so we
 // can ignore its return values in [codecImpl] methods.
-type codecImpl struct {
-	// Invariant: Every byte slice returned by [varIntPool] has
-	// length [binary.MaxVarintLen64].
-	varIntPool sync.Pool
-}
+type codecImpl struct{}
 
 func (c *codecImpl) childSize(index byte, childEntry *child) int {
 	// * index
@@ -285,11 +274,10 @@ func (*codecImpl) decodeUint(src *bytes.Reader) (uint64, error) {
 	return val64, nil
 }
 
-func (c *codecImpl) encodeUint(dst *bytes.Buffer, value uint64) {
-	buf := c.varIntPool.Get().([]byte)
-	size := binary.PutUvarint(buf, value)
+func (*codecImpl) encodeUint(dst *bytes.Buffer, value uint64) {
+	var buf [binary.MaxVarintLen64]byte
+	size := binary.PutUvarint(buf[:], value)
 	_, _ = dst.Write(buf[:size])
-	c.varIntPool.Put(buf)
 }
 
 func (c *codecImpl) encodeMaybeByteSlice(dst *bytes.Buffer, maybeValue maybe.Maybe[[]byte]) {
diff --git a/x/merkledb/codec_test.go b/x/merkledb/codec_test.go
index 455b75e1bed..7d8e6028cfe 100644
--- a/x/merkledb/codec_test.go
+++ b/x/merkledb/codec_test.go
@@ -9,6 +9,7 @@ import (
 	"io"
 	"math"
 	"math/rand"
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -266,3 +267,19 @@ func TestUintSize(t *testing.T) {
 		require.Equal(t, expectedSize, actualSize, power)
 	}
 }
+
+func Benchmark_EncodeUint(b *testing.B) {
+	c := codec.(*codecImpl)
+
+	var dst bytes.Buffer
+	dst.Grow(binary.MaxVarintLen64)
+
+	for _, v := range []uint64{0, 1, 2, 32, 1024, 32768} {
+		b.Run(strconv.FormatUint(v, 10), func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				c.encodeUint(&dst, v)
+				dst.Reset()
+			}
+		})
+	}
+}