Optimize encodeUint (#2882)

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]) {

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()
+			}
+		})
+	}
+}