Update uintsize implementation (#2590)

x/merkledb/codec.go

@@ -9,6 +9,7 @@ import (
 	"errors"
 	"io"
 	"math"
+	"math/bits"
 	"sync"
 
 	"golang.org/x/exp/maps"
@@ -101,14 +102,10 @@ func (c *codecImpl) childSize(index byte, childEntry *child) int {
 
 // based on the current implementation of codecImpl.encodeUint which uses binary.PutUvarint
 func (*codecImpl) uintSize(value uint64) int {
-	// binary.PutUvarint repeatedly divides by 128 until the value is under 128,
-	// so count the number of times that will occur
-	i := 0
-	for value >= 0x80 {
-		value >>= 7
-		i++
-	}
-	return i + 1
+	if value == 0 {
+		return 1
+	}
+	return (bits.Len64(value) + 6) / 7
 }
 
 func (c *codecImpl) keySize(p Key) int {

x/merkledb/codec_test.go

@@ -247,9 +247,22 @@ func TestCodecDecodeKeyLengthOverflowRegression(t *testing.T) {
 
 func TestUintSize(t *testing.T) {
 	c := codec.(*codecImpl)
-	for i := uint64(0); i < math.MaxInt16; i++ {
-		expectedSize := c.uintSize(i)
-		actualSize := binary.PutUvarint(make([]byte, binary.MaxVarintLen64), i)
-		require.Equal(t, expectedSize, actualSize, i)
+
+	// Test lower bound
+	expectedSize := c.uintSize(0)
+	actualSize := binary.PutUvarint(make([]byte, binary.MaxVarintLen64), 0)
+	require.Equal(t, expectedSize, actualSize)
+
+	// Test upper bound
+	expectedSize = c.uintSize(math.MaxUint64)
+	actualSize = binary.PutUvarint(make([]byte, binary.MaxVarintLen64), math.MaxUint64)
+	require.Equal(t, expectedSize, actualSize)
+
+	// Test powers of 2
+	for power := 0; power < 64; power++ {
+		n := uint64(1) << uint(power)
+		expectedSize := c.uintSize(n)
+		actualSize := binary.PutUvarint(make([]byte, binary.MaxVarintLen64), n)
+		require.Equal(t, expectedSize, actualSize, power)
 	}
 }