From bd2b1a0521c02826ab18718b7003472c6adc91d5 Mon Sep 17 00:00:00 2001
From: Evgeniy Stratonikov <evgeniy@nspcc.ru>
Date: Tue, 3 Aug 2021 18:50:13 +0300
Subject: [PATCH] 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      |  6 ++++--
 pkg/core/mpt/branch.go    | 11 +++++++++++
 pkg/core/mpt/empty.go     |  3 +++
 pkg/core/mpt/extension.go |  6 ++++++
 pkg/core/mpt/hash.go      |  5 +++++
 pkg/core/mpt/leaf.go      |  5 +++++
 pkg/core/mpt/node.go      |  1 +
 pkg/core/mpt/node_test.go |  1 +
 8 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/pkg/core/mpt/base.go b/pkg/core/mpt/base.go
index f1cd5864d..aa5336880 100644
--- a/pkg/core/mpt/base.go
+++ b/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
 }
diff --git a/pkg/core/mpt/branch.go b/pkg/core/mpt/branch.go
index 5d4c6ad23..0338ff4a7 100644
--- a/pkg/core/mpt/branch.go
+++ b/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++ {
diff --git a/pkg/core/mpt/empty.go b/pkg/core/mpt/empty.go
index 5d5c3f32c..6669ef8c1 100644
--- a/pkg/core/mpt/empty.go
+++ b/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
diff --git a/pkg/core/mpt/extension.go b/pkg/core/mpt/extension.go
index 1b8047e20..2dcbcb66b 100644
--- a/pkg/core/mpt/extension.go
+++ b/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{}{
diff --git a/pkg/core/mpt/hash.go b/pkg/core/mpt/hash.go
index df9ab6017..05ddbe5f3 100644
--- a/pkg/core/mpt/hash.go
+++ b/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 {
diff --git a/pkg/core/mpt/leaf.go b/pkg/core/mpt/leaf.go
index ecb003c23..0f3072b85 100644
--- a/pkg/core/mpt/leaf.go
+++ b/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
diff --git a/pkg/core/mpt/node.go b/pkg/core/mpt/node.go
index 2eea1fb4a..af35286c1 100644
--- a/pkg/core/mpt/node.go
+++ b/pkg/core/mpt/node.go
@@ -33,6 +33,7 @@ type Node interface {
 	io.Serializable
 	json.Marshaler
 	json.Unmarshaler
+	Size() int
 	BaseNodeIface
 }
 
diff --git a/pkg/core/mpt/node_test.go b/pkg/core/mpt/node_test.go
index ca525ac5b..a4425d71c 100644
--- a/pkg/core/mpt/node_test.go
+++ b/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)