merkledb -- dynamic root

scripts/mocks.mockgen.txt

@@ -43,5 +43,4 @@ github.com/ava-labs/avalanchego/vms/registry=VMGetter=vms/registry/mock_vm_gette
 github.com/ava-labs/avalanchego/vms/registry=VMRegisterer=vms/registry/mock_vm_registerer.go
 github.com/ava-labs/avalanchego/vms/registry=VMRegistry=vms/registry/mock_vm_registry.go
 github.com/ava-labs/avalanchego/vms=Factory,Manager=vms/mock_manager.go
-github.com/ava-labs/avalanchego/x/merkledb=MerkleDB=x/merkledb/mock_db.go
 github.com/ava-labs/avalanchego/x/sync=Client=x/sync/mock_client.go

x/merkledb/codec.go

@@ -66,11 +66,13 @@ type encoder interface {
 	// Returns the bytes that will be hashed to generate [n]'s ID.
 	// Assumes [n] is non-nil.
 	encodeHashValues(n *node) []byte
+	encodeKey(key Key) []byte
 }
 
 type decoder interface {
 	// Assumes [n] is non-nil.
 	decodeDBNode(bytes []byte, n *dbNode) error
+	decodeKey(bytes []byte) (Key, error)
 }
 
 func newCodec() encoderDecoder {
@@ -98,7 +100,6 @@ func (c *codecImpl) encodeDBNode(n *dbNode) []byte {
 		estimatedLen = estimatedValueLen + minVarIntLen + estimatedNodeChildLen*numChildren
 		buf          = bytes.NewBuffer(make([]byte, 0, estimatedLen))
 	)
-
 	c.encodeMaybeByteSlice(buf, n.value)
 	c.encodeUint(buf, uint64(numChildren))
 	// Note we insert children in order of increasing index
@@ -108,7 +109,7 @@ func (c *codecImpl) encodeDBNode(n *dbNode) []byte {
 	for _, index := range keys {
 		entry := n.children[index]
 		c.encodeUint(buf, uint64(index))
-		c.encodeKey(buf, entry.compressedKey)
+		c.encodeKeyToBuffer(buf, entry.compressedKey)
 		_, _ = buf.Write(entry.id[:])
 		c.encodeBool(buf, entry.hasValue)
 	}
@@ -134,7 +135,7 @@ func (c *codecImpl) encodeHashValues(n *node) []byte {
 		_, _ = buf.Write(entry.id[:])
 	}
 	c.encodeMaybeByteSlice(buf, n.valueDigest)
-	c.encodeKey(buf, n.key)
+	c.encodeKeyToBuffer(buf, n.key)
 
 	return buf.Bytes()
 }
@@ -172,7 +173,7 @@ func (c *codecImpl) decodeDBNode(b []byte, n *dbNode) error {
 		}
 		previousChild = index
 
-		compressedKey, err := c.decodeKey(src)
+		compressedKey, err := c.decodeKeyFromReader(src)
 		if err != nil {
 			return err
 		}
@@ -330,12 +331,31 @@ func (*codecImpl) decodeID(src *bytes.Reader) (ids.ID, error) {
 	return id, err
 }
 
-func (c *codecImpl) encodeKey(dst *bytes.Buffer, key Key) {
+func (c *codecImpl) encodeKey(key Key) []byte {
+	estimatedLen := binary.MaxVarintLen64 + len(key.Bytes())
+	dst := bytes.NewBuffer(make([]byte, 0, estimatedLen))
+	c.encodeKeyToBuffer(dst, key)
+	return dst.Bytes()
+}
+
+func (c *codecImpl) encodeKeyToBuffer(dst *bytes.Buffer, key Key) {
 	c.encodeUint(dst, uint64(key.length))
 	_, _ = dst.Write(key.Bytes())
 }
 
-func (c *codecImpl) decodeKey(src *bytes.Reader) (Key, error) {
+func (c *codecImpl) decodeKey(b []byte) (Key, error) {
+	src := bytes.NewReader(b)
+	key, err := c.decodeKeyFromReader(src)
+	if err != nil {
+		return Key{}, err
+	}
+	if src.Len() != 0 {
+		return Key{}, errExtraSpace
+	}
+	return key, err
+}
+
+func (c *codecImpl) decodeKeyFromReader(src *bytes.Reader) (Key, error) {
 	if minKeyLen > src.Len() {
 		return Key{}, io.ErrUnexpectedEOF
 	}

x/merkledb/codec_test.go

@@ -81,21 +81,14 @@ func FuzzCodecKey(f *testing.F) {
 		) {
 			require := require.New(t)
 			codec := codec.(*codecImpl)
-			reader := bytes.NewReader(b)
-			startLen := reader.Len()
-			got, err := codec.decodeKey(reader)
+			got, err := codec.decodeKey(b)
 			if err != nil {
 				t.SkipNow()
 			}
-			endLen := reader.Len()
-			numRead := startLen - endLen
 
 			// Encoding [got] should be the same as [b].
-			var buf bytes.Buffer
-			codec.encodeKey(&buf, got)
-			bufBytes := buf.Bytes()
-			require.Len(bufBytes, numRead)
-			require.Equal(b[:numRead], bufBytes)
+			gotBytes := codec.encodeKey(got)
+			require.Equal(b, gotBytes)
 		},
 	)
 }
@@ -248,7 +241,6 @@ func FuzzEncodeHashValues(f *testing.F) {
 
 func TestCodecDecodeKeyLengthOverflowRegression(t *testing.T) {
 	codec := codec.(*codecImpl)
-	bytes := bytes.NewReader(binary.AppendUvarint(nil, math.MaxInt))
-	_, err := codec.decodeKey(bytes)
+	_, err := codec.decodeKey(binary.AppendUvarint(nil, math.MaxInt))
 	require.ErrorIs(t, err, io.ErrUnexpectedEOF)
 }