From 73e4040628bfab5136cb8aecc09bc9f5e4e4afef Mon Sep 17 00:00:00 2001
From: Evgeniy Stratonikov <evgeniy@nspcc.ru>
Date: Fri, 6 Aug 2021 11:32:02 +0300
Subject: [PATCH] mpt: use `BinWriter.Grow()` instead of custom buffer

Also add benchmarks.

Signed-off-by: Evgeniy Stratonikov <evgeniy@nspcc.ru>
---
 pkg/core/mpt/base.go       |  9 ++++-----
 pkg/core/mpt/bench_test.go | 39 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 5 deletions(-)
 create mode 100644 pkg/core/mpt/bench_test.go

diff --git a/pkg/core/mpt/base.go b/pkg/core/mpt/base.go
index aa5336880..c2f7f2f11 100644
--- a/pkg/core/mpt/base.go
+++ b/pkg/core/mpt/base.go
@@ -1,7 +1,6 @@
 package mpt
 
 import (
-	"bytes"
 	"fmt"
 
 	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
@@ -64,10 +63,10 @@ func (b *BaseNode) updateHash(n Node) {
 
 // updateCache updates hash and bytes fields for this BaseNode.
 func (b *BaseNode) updateBytes(n Node) {
-	buf := bytes.NewBuffer(make([]byte, 0, 1+n.Size()))
-	bw := io.NewBinWriterFromIO(buf)
-	encodeNodeWithType(n, bw)
-	b.bytes = buf.Bytes()
+	bw := io.NewBufBinWriter()
+	bw.Grow(1 + n.Size())
+	encodeNodeWithType(n, bw.BinWriter)
+	b.bytes = bw.Bytes()
 	b.bytesValid = true
 }
 
diff --git a/pkg/core/mpt/bench_test.go b/pkg/core/mpt/bench_test.go
new file mode 100644
index 000000000..7c3c00d51
--- /dev/null
+++ b/pkg/core/mpt/bench_test.go
@@ -0,0 +1,39 @@
+package mpt
+
+import (
+	"testing"
+
+	"github.com/nspcc-dev/neo-go/internal/random"
+)
+
+func benchmarkBytes(b *testing.B, n Node) {
+	inv := n.(interface{ invalidateCache() })
+	b.ReportAllocs()
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		inv.invalidateCache()
+		_ = n.Bytes()
+	}
+}
+
+func BenchmarkBytes(b *testing.B) {
+	b.Run("extension", func(b *testing.B) {
+		n := NewExtensionNode(random.Bytes(10), NewLeafNode(random.Bytes(10)))
+		benchmarkBytes(b, n)
+	})
+	b.Run("leaf", func(b *testing.B) {
+		n := NewLeafNode(make([]byte, 15))
+		benchmarkBytes(b, n)
+	})
+	b.Run("hash", func(b *testing.B) {
+		n := NewHashNode(random.Uint256())
+		benchmarkBytes(b, n)
+	})
+	b.Run("branch", func(b *testing.B) {
+		n := NewBranchNode()
+		n.Children[0] = NewLeafNode(random.Bytes(10))
+		n.Children[4] = NewLeafNode(random.Bytes(10))
+		n.Children[7] = NewLeafNode(random.Bytes(10))
+		n.Children[8] = NewLeafNode(random.Bytes(10))
+	})
+}