mpt: add Size method to trie nodes

Knowing serialized size of the node is useful for
preallocating byte-slice in advance.

Signed-off-by: Evgeniy Stratonikov <evgeniy@nspcc.ru>

pkg/core/mpt/base.go

@@ -1,6 +1,7 @@
 package mpt
 
 import (
+	"bytes"
 	"fmt"
 
 	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
@@ -63,8 +64,9 @@ func (b *BaseNode) updateHash(n Node) {
 
 // updateCache updates hash and bytes fields for this BaseNode.
 func (b *BaseNode) updateBytes(n Node) {
-	buf := io.NewBufBinWriter()
-	encodeNodeWithType(n, buf.BinWriter)
+	buf := bytes.NewBuffer(make([]byte, 0, 1+n.Size()))
+	bw := io.NewBinWriterFromIO(buf)
+	encodeNodeWithType(n, bw)
 	b.bytes = buf.Bytes()
 	b.bytesValid = true
 }

pkg/core/mpt/branch.go

@@ -45,6 +45,17 @@ func (b *BranchNode) Bytes() []byte {
 	return b.getBytes(b)
 }
 
+// Size implements Node interface.
+func (b *BranchNode) Size() int {
+	sz := childrenCount
+	for i := range b.Children {
+		if !isEmpty(b.Children[i]) {
+			sz += util.Uint256Size
+		}
+	}
+	return sz
+}
+
 // EncodeBinary implements io.Serializable.
 func (b *BranchNode) EncodeBinary(w *io.BinWriter) {
 	for i := 0; i < childrenCount; i++ {

pkg/core/mpt/empty.go

@@ -19,6 +19,9 @@ func (e EmptyNode) DecodeBinary(*io.BinReader) {
 func (e EmptyNode) EncodeBinary(*io.BinWriter) {
 }
 
+// Size implements Node interface.
+func (EmptyNode) Size() int { return 0 }
+
 // MarshalJSON implements Node interface.
 func (e EmptyNode) MarshalJSON() ([]byte, error) {
 	return []byte(`{}`), nil

pkg/core/mpt/extension.go

@@ -72,6 +72,12 @@ func (e ExtensionNode) EncodeBinary(w *io.BinWriter) {
 	encodeBinaryAsChild(e.next, w)
 }
 
+// Size implements Node interface.
+func (e *ExtensionNode) Size() int {
+	return io.GetVarSize(len(e.key)) + len(e.key) +
+		1 + util.Uint256Size // e.next is never empty
+}
+
 // MarshalJSON implements json.Marshaler.
 func (e *ExtensionNode) MarshalJSON() ([]byte, error) {
 	m := map[string]interface{}{

pkg/core/mpt/hash.go

@@ -27,6 +27,11 @@ func NewHashNode(h util.Uint256) *HashNode {
 // Type implements Node interface.
 func (h *HashNode) Type() NodeType { return HashT }
 
+// Size implements Node interface.
+func (h *HashNode) Size() int {
+	return util.Uint256Size
+}
+
 // Hash implements Node interface.
 func (h *HashNode) Hash() util.Uint256 {
 	if !h.hashValid {

pkg/core/mpt/leaf.go

@@ -56,6 +56,11 @@ func (n LeafNode) EncodeBinary(w *io.BinWriter) {
 	w.WriteVarBytes(n.value)
 }
 
+// Size implements Node interface.
+func (n *LeafNode) Size() int {
+	return io.GetVarSize(len(n.value)) + len(n.value)
+}
+
 // MarshalJSON implements json.Marshaler.
 func (n *LeafNode) MarshalJSON() ([]byte, error) {
 	return []byte(`{"value":"` + hex.EncodeToString(n.value) + `"}`), nil

pkg/core/mpt/node.go

@@ -33,6 +33,7 @@ type Node interface {
 	io.Serializable
 	json.Marshaler
 	json.Unmarshaler
+	Size() int
 	BaseNodeIface
 }
 

pkg/core/mpt/node_test.go

@@ -27,6 +27,7 @@ func getTestFuncEncode(ok bool, expected, actual Node) func(t *testing.T) {
 			require.NoError(t, err)
 			require.Equal(t, expected.Type(), actual.Type())
 			require.Equal(t, expected.Hash(), actual.Hash())
+			require.Equal(t, 1+expected.Size(), len(expected.Bytes()))
 		})
 		t.Run("JSON", func(t *testing.T) {
 			bs, err := json.Marshal(expected)