From c3d989ebdaeada5107a96bf20aa2b4423b95c384 Mon Sep 17 00:00:00 2001
From: Roman Khimov <roman@nspcc.ru>
Date: Tue, 31 May 2022 20:10:20 +0300
Subject: [PATCH] stackitem: reusable serialization context

We serialize items a lot and this allows to avoid a number of allocations.
---
 pkg/core/blockchain.go                  |  1 +
 pkg/core/dao/dao.go                     | 19 +++++++++--
 pkg/core/interop/runtime/engine.go      |  2 +-
 pkg/core/interop/runtime/engine_test.go |  3 +-
 pkg/core/native/native_neo.go           | 12 +++----
 pkg/core/native/neo_types.go            |  4 +--
 pkg/core/native/oracle.go               |  3 +-
 pkg/core/native/std.go                  |  6 ++--
 pkg/core/native/std_test.go             |  9 +++---
 pkg/core/native/util.go                 |  6 +++-
 pkg/core/state/native_state.go          | 17 ++++------
 pkg/core/state/notification_event.go    | 28 ++++++++++++++--
 pkg/vm/stackitem/serialization.go       | 43 +++++++++++++++++++++----
 13 files changed, 111 insertions(+), 42 deletions(-)

diff --git a/pkg/core/blockchain.go b/pkg/core/blockchain.go
index a1616389b..1f881156b 100644
--- a/pkg/core/blockchain.go
+++ b/pkg/core/blockchain.go
@@ -1094,6 +1094,7 @@ func (bc *Blockchain) storeBlock(block *block.Block, txpool *mempool.Pool) error
 		}
 		close(aerdone)
 	}()
+	_ = cache.GetItemCtx() // Prime serialization context cache (it'll be reused by upper layer DAOs).
 	aer, err := bc.runPersist(bc.contracts.GetPersistScript(), block, cache, trigger.OnPersist)
 	if err != nil {
 		// Release goroutines, don't care about errors, we already have one.
diff --git a/pkg/core/dao/dao.go b/pkg/core/dao/dao.go
index daf2e60c2..6995612ce 100644
--- a/pkg/core/dao/dao.go
+++ b/pkg/core/dao/dao.go
@@ -17,6 +17,7 @@ import (
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/trigger"
 	"github.com/nspcc-dev/neo-go/pkg/util"
 	"github.com/nspcc-dev/neo-go/pkg/util/slice"
+	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
 )
 
 // HasTransaction errors.
@@ -44,6 +45,7 @@ type Simple struct {
 	nativeCachePS *Simple
 
 	private bool
+	serCtx  *stackitem.SerializationContext
 	keyBuf  []byte
 	dataBuf *io.BufBinWriter
 }
@@ -96,6 +98,7 @@ func (dao *Simple) GetPrivate() *Simple {
 		Version: dao.Version,
 		keyBuf:  dao.keyBuf,
 		dataBuf: dao.dataBuf,
+		serCtx:  dao.serCtx,
 	} // Inherit everything...
 	d.Store = storage.NewPrivateMemCachedStore(dao.Store) // except storage, wrap another layer.
 	d.private = true
@@ -710,10 +713,10 @@ func (dao *Simple) StoreAsBlock(block *block.Block, aer1 *state.AppExecResult, a
 	buf.WriteB(storage.ExecBlock)
 	block.EncodeTrimmed(buf.BinWriter)
 	if aer1 != nil {
-		aer1.EncodeBinary(buf.BinWriter)
+		aer1.EncodeBinaryWithContext(buf.BinWriter, dao.GetItemCtx())
 	}
 	if aer2 != nil {
-		aer2.EncodeBinary(buf.BinWriter)
+		aer2.EncodeBinaryWithContext(buf.BinWriter, dao.GetItemCtx())
 	}
 	if buf.Err != nil {
 		return buf.Err
@@ -790,7 +793,7 @@ func (dao *Simple) StoreAsTransaction(tx *transaction.Transaction, index uint32,
 	buf.WriteU32LE(index)
 	tx.EncodeBinary(buf.BinWriter)
 	if aer != nil {
-		aer.EncodeBinary(buf.BinWriter)
+		aer.EncodeBinaryWithContext(buf.BinWriter, dao.GetItemCtx())
 	}
 	if buf.Err != nil {
 		return buf.Err
@@ -835,6 +838,16 @@ func (dao *Simple) getDataBuf() *io.BufBinWriter {
 	return io.NewBufBinWriter()
 }
 
+func (dao *Simple) GetItemCtx() *stackitem.SerializationContext {
+	if dao.private {
+		if dao.serCtx == nil {
+			dao.serCtx = stackitem.NewSerializationContext()
+		}
+		return dao.serCtx
+	}
+	return stackitem.NewSerializationContext()
+}
+
 // Persist flushes all the changes made into the (supposedly) persistent
 // underlying store. It doesn't block accesses to DAO from other threads.
 func (dao *Simple) Persist() (int, error) {
diff --git a/pkg/core/interop/runtime/engine.go b/pkg/core/interop/runtime/engine.go
index f34358479..749d704dd 100644
--- a/pkg/core/interop/runtime/engine.go
+++ b/pkg/core/interop/runtime/engine.go
@@ -61,7 +61,7 @@ func Notify(ic *interop.Context) error {
 	// But it has to be serializable, otherwise we either have some broken
 	// (recursive) structure inside or an interop item that can't be used
 	// outside of the interop subsystem anyway.
-	bytes, err := stackitem.Serialize(elem.Item())
+	bytes, err := ic.DAO.GetItemCtx().Serialize(elem.Item(), false)
 	if err != nil {
 		return fmt.Errorf("bad notification: %w", err)
 	}
diff --git a/pkg/core/interop/runtime/engine_test.go b/pkg/core/interop/runtime/engine_test.go
index 41b1c2dc9..fffe8900f 100644
--- a/pkg/core/interop/runtime/engine_test.go
+++ b/pkg/core/interop/runtime/engine_test.go
@@ -8,6 +8,7 @@ import (
 
 	"github.com/nspcc-dev/neo-go/internal/random"
 	"github.com/nspcc-dev/neo-go/pkg/core/block"
+	"github.com/nspcc-dev/neo-go/pkg/core/dao"
 	"github.com/nspcc-dev/neo-go/pkg/core/interop"
 	"github.com/nspcc-dev/neo-go/pkg/crypto/hash"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/callflag"
@@ -133,7 +134,7 @@ func TestLog(t *testing.T) {
 func TestNotify(t *testing.T) {
 	h := random.Uint160()
 	newIC := func(name string, args interface{}) *interop.Context {
-		ic := &interop.Context{VM: vm.New()}
+		ic := &interop.Context{VM: vm.New(), DAO: &dao.Simple{}}
 		ic.VM.LoadScriptWithHash([]byte{1}, h, callflag.NoneFlag)
 		ic.VM.Estack().PushVal(args)
 		ic.VM.Estack().PushVal(name)
diff --git a/pkg/core/native/native_neo.go b/pkg/core/native/native_neo.go
index 734667014..7e601d3c8 100644
--- a/pkg/core/native/native_neo.go
+++ b/pkg/core/native/native_neo.go
@@ -281,7 +281,7 @@ func (n *NEO) Initialize(ic *interop.Context) error {
 		return err
 	}
 
-	ic.DAO.PutStorageItem(n.ID, prefixCommittee, cvs.Bytes())
+	ic.DAO.PutStorageItem(n.ID, prefixCommittee, cvs.Bytes(ic.DAO.GetItemCtx()))
 
 	h, err := getStandbyValidatorsHash(ic)
 	if err != nil {
@@ -355,7 +355,7 @@ func (n *NEO) updateCache(cache *NeoCache, cvs keysWithVotes, blockHeight uint32
 func (n *NEO) updateCommittee(cache *NeoCache, ic *interop.Context) error {
 	if !cache.votesChanged {
 		// We need to put in storage anyway, as it affects dumps
-		ic.DAO.PutStorageItem(n.ID, prefixCommittee, cache.committee.Bytes())
+		ic.DAO.PutStorageItem(n.ID, prefixCommittee, cache.committee.Bytes(ic.DAO.GetItemCtx()))
 		return nil
 	}
 
@@ -367,7 +367,7 @@ func (n *NEO) updateCommittee(cache *NeoCache, ic *interop.Context) error {
 		return err
 	}
 	cache.votesChanged = false
-	ic.DAO.PutStorageItem(n.ID, prefixCommittee, cvs.Bytes())
+	ic.DAO.PutStorageItem(n.ID, prefixCommittee, cvs.Bytes(ic.DAO.GetItemCtx()))
 	return nil
 }
 
@@ -495,7 +495,7 @@ func (n *NEO) increaseBalance(ic *interop.Context, h util.Uint160, si *state.Sto
 		postF = func() { n.GAS.mint(ic, h, newGas, true) }
 	}
 	if amount.Sign() == 0 {
-		*si = acc.Bytes()
+		*si = acc.Bytes(ic.DAO.GetItemCtx())
 		return postF, nil
 	}
 	if err := n.ModifyAccountVotes(acc, ic.DAO, amount, false); err != nil {
@@ -508,7 +508,7 @@ func (n *NEO) increaseBalance(ic *interop.Context, h util.Uint160, si *state.Sto
 	}
 	acc.Balance.Add(&acc.Balance, amount)
 	if acc.Balance.Sign() != 0 {
-		*si = acc.Bytes()
+		*si = acc.Bytes(ic.DAO.GetItemCtx())
 	} else {
 		*si = nil
 	}
@@ -872,7 +872,7 @@ func (n *NEO) VoteInternal(ic *interop.Context, h util.Uint160, pub *keys.Public
 	if err := n.ModifyAccountVotes(acc, ic.DAO, &acc.Balance, true); err != nil {
 		return err
 	}
-	ic.DAO.PutStorageItem(n.ID, key, acc.Bytes())
+	ic.DAO.PutStorageItem(n.ID, key, acc.Bytes(ic.DAO.GetItemCtx()))
 
 	ic.AddNotification(n.Hash, "Vote", stackitem.NewArray([]stackitem.Item{
 		stackitem.NewByteArray(h.BytesBE()),
diff --git a/pkg/core/native/neo_types.go b/pkg/core/native/neo_types.go
index 7c27e45bc..530f7b935 100644
--- a/pkg/core/native/neo_types.go
+++ b/pkg/core/native/neo_types.go
@@ -82,8 +82,8 @@ func (k *keysWithVotes) fromStackItem(item stackitem.Item) error {
 }
 
 // Bytes serializes keys with votes slice.
-func (k keysWithVotes) Bytes() []byte {
-	buf, err := stackitem.Serialize(k.toStackItem())
+func (k keysWithVotes) Bytes(sc *stackitem.SerializationContext) []byte {
+	buf, err := sc.Serialize(k.toStackItem(), false)
 	if err != nil {
 		panic(err)
 	}
diff --git a/pkg/core/native/oracle.go b/pkg/core/native/oracle.go
index 800423656..0f819104a 100644
--- a/pkg/core/native/oracle.go
+++ b/pkg/core/native/oracle.go
@@ -364,13 +364,14 @@ func (o *Oracle) RequestInternal(ic *interop.Context, url string, filter *string
 		return err
 	}
 
-	data, err := stackitem.Serialize(userData)
+	data, err := ic.DAO.GetItemCtx().Serialize(userData, false)
 	if err != nil {
 		return err
 	}
 	if len(data) > maxUserDataLength {
 		return ErrBigArgument
 	}
+	data = slice.Copy(data) // Serialization context will be used in PutRequestInternal again.
 
 	var filterNotif stackitem.Item
 	if filter != nil {
diff --git a/pkg/core/native/std.go b/pkg/core/native/std.go
index 00ff22562..733fd7338 100644
--- a/pkg/core/native/std.go
+++ b/pkg/core/native/std.go
@@ -160,8 +160,8 @@ func newStd() *Std {
 	return s
 }
 
-func (s *Std) serialize(_ *interop.Context, args []stackitem.Item) stackitem.Item {
-	data, err := stackitem.Serialize(args[0])
+func (s *Std) serialize(ic *interop.Context, args []stackitem.Item) stackitem.Item {
+	data, err := ic.DAO.GetItemCtx().Serialize(args[0], false)
 	if err != nil {
 		panic(err)
 	}
@@ -169,7 +169,7 @@ func (s *Std) serialize(_ *interop.Context, args []stackitem.Item) stackitem.Ite
 		panic(errors.New("too big item"))
 	}
 
-	return stackitem.NewByteArray(data)
+	return stackitem.NewByteArray(slice.Copy(data)) // Serialization context can be reused.
 }
 
 func (s *Std) deserialize(_ *interop.Context, args []stackitem.Item) stackitem.Item {
diff --git a/pkg/core/native/std_test.go b/pkg/core/native/std_test.go
index da56a4f90..2ee87f31a 100644
--- a/pkg/core/native/std_test.go
+++ b/pkg/core/native/std_test.go
@@ -9,6 +9,7 @@ import (
 	"testing"
 
 	"github.com/mr-tron/base58"
+	"github.com/nspcc-dev/neo-go/pkg/core/dao"
 	"github.com/nspcc-dev/neo-go/pkg/core/interop"
 	base58neogo "github.com/nspcc-dev/neo-go/pkg/encoding/base58"
 	"github.com/nspcc-dev/neo-go/pkg/vm"
@@ -19,7 +20,7 @@ import (
 
 func TestStdLibItoaAtoi(t *testing.T) {
 	s := newStd()
-	ic := &interop.Context{VM: vm.New()}
+	ic := &interop.Context{VM: vm.New(), DAO: &dao.Simple{}}
 	var actual stackitem.Item
 
 	t.Run("itoa-atoi", func(t *testing.T) {
@@ -251,7 +252,7 @@ func TestStdLibEncodeDecode(t *testing.T) {
 
 func TestStdLibSerialize(t *testing.T) {
 	s := newStd()
-	ic := &interop.Context{VM: vm.New()}
+	ic := &interop.Context{VM: vm.New(), DAO: &dao.Simple{}}
 
 	t.Run("recursive", func(t *testing.T) {
 		arr := stackitem.NewArray(nil)
@@ -294,7 +295,7 @@ func TestStdLibSerialize(t *testing.T) {
 
 func TestStdLibSerializeDeserialize(t *testing.T) {
 	s := newStd()
-	ic := &interop.Context{VM: vm.New()}
+	ic := &interop.Context{VM: vm.New(), DAO: &dao.Simple{}}
 	var actual stackitem.Item
 
 	checkSerializeDeserialize := func(t *testing.T, value interface{}, expected stackitem.Item) {
@@ -381,7 +382,7 @@ func TestStdLibSerializeDeserialize(t *testing.T) {
 
 func TestMemoryCompare(t *testing.T) {
 	s := newStd()
-	ic := &interop.Context{VM: vm.New()}
+	ic := &interop.Context{VM: vm.New(), DAO: &dao.Simple{}}
 
 	check := func(t *testing.T, result int64, s1, s2 string) {
 		actual := s.memoryCompare(ic, []stackitem.Item{stackitem.Make(s1), stackitem.Make(s2)})
diff --git a/pkg/core/native/util.go b/pkg/core/native/util.go
index 38fe03eee..a9d85ba9e 100644
--- a/pkg/core/native/util.go
+++ b/pkg/core/native/util.go
@@ -24,7 +24,11 @@ func getConvertibleFromDAO(id int32, d *dao.Simple, key []byte, conv stackitem.C
 }
 
 func putConvertibleToDAO(id int32, d *dao.Simple, key []byte, conv stackitem.Convertible) error {
-	data, err := stackitem.SerializeConvertible(conv)
+	item, err := conv.ToStackItem()
+	if err != nil {
+		return err
+	}
+	data, err := d.GetItemCtx().Serialize(item, false)
 	if err != nil {
 		return err
 	}
diff --git a/pkg/core/state/native_state.go b/pkg/core/state/native_state.go
index 0f255b01c..48eef1a0f 100644
--- a/pkg/core/state/native_state.go
+++ b/pkg/core/state/native_state.go
@@ -70,14 +70,6 @@ func balanceFromBytes(b []byte, item stackitem.Convertible) error {
 	return stackitem.DeserializeConvertible(b, item)
 }
 
-func balanceToBytes(item stackitem.Convertible) []byte {
-	data, err := stackitem.SerializeConvertible(item)
-	if err != nil {
-		panic(err)
-	}
-	return data
-}
-
 // ToStackItem implements stackitem.Convertible. It never returns an error.
 func (s *NEP17Balance) ToStackItem() (stackitem.Item, error) {
 	return stackitem.NewStruct([]stackitem.Item{stackitem.NewBigInteger(&s.Balance)}), nil
@@ -111,8 +103,13 @@ func NEOBalanceFromBytes(b []byte) (*NEOBalance, error) {
 }
 
 // Bytes returns a serialized NEOBalance.
-func (s *NEOBalance) Bytes() []byte {
-	return balanceToBytes(s)
+func (s *NEOBalance) Bytes(sc *stackitem.SerializationContext) []byte {
+	item, _ := s.ToStackItem() // Never returns an error.
+	data, err := sc.Serialize(item, false)
+	if err != nil {
+		panic(err)
+	}
+	return data
 }
 
 // ToStackItem implements stackitem.Convertible interface. It never returns an error.
diff --git a/pkg/core/state/notification_event.go b/pkg/core/state/notification_event.go
index 709d7d165..d87bb40e0 100644
--- a/pkg/core/state/notification_event.go
+++ b/pkg/core/state/notification_event.go
@@ -29,9 +29,20 @@ type AppExecResult struct {
 
 // EncodeBinary implements the Serializable interface.
 func (ne *NotificationEvent) EncodeBinary(w *io.BinWriter) {
+	ne.EncodeBinaryWithContext(w, stackitem.NewSerializationContext())
+}
+
+// EncodeBinaryWithContext is the same as EncodeBinary, but allows to efficiently reuse
+// stack item serialization context.
+func (ne *NotificationEvent) EncodeBinaryWithContext(w *io.BinWriter, sc *stackitem.SerializationContext) {
 	ne.ScriptHash.EncodeBinary(w)
 	w.WriteString(ne.Name)
-	stackitem.EncodeBinary(ne.Item, w)
+	b, err := sc.Serialize(ne.Item, false)
+	if err != nil {
+		w.Err = err
+		return
+	}
+	w.WriteBytes(b)
 }
 
 // DecodeBinary implements the Serializable interface.
@@ -52,6 +63,12 @@ func (ne *NotificationEvent) DecodeBinary(r *io.BinReader) {
 
 // EncodeBinary implements the Serializable interface.
 func (aer *AppExecResult) EncodeBinary(w *io.BinWriter) {
+	aer.EncodeBinaryWithContext(w, stackitem.NewSerializationContext())
+}
+
+// EncodeBinaryWithContext is the same as EncodeBinary, but allows to efficiently reuse
+// stack item serialization context.
+func (aer *AppExecResult) EncodeBinaryWithContext(w *io.BinWriter, sc *stackitem.SerializationContext) {
 	w.WriteBytes(aer.Container[:])
 	w.WriteB(byte(aer.Trigger))
 	w.WriteB(byte(aer.VMState))
@@ -59,11 +76,16 @@ func (aer *AppExecResult) EncodeBinary(w *io.BinWriter) {
 	// Stack items are expected to be marshaled one by one.
 	w.WriteVarUint(uint64(len(aer.Stack)))
 	for _, it := range aer.Stack {
-		stackitem.EncodeBinaryProtected(it, w)
+		b, err := sc.Serialize(it, true)
+		if err != nil {
+			w.Err = err
+			return
+		}
+		w.WriteBytes(b)
 	}
 	w.WriteVarUint(uint64(len(aer.Events)))
 	for i := range aer.Events {
-		aer.Events[i].EncodeBinary(w)
+		aer.Events[i].EncodeBinaryWithContext(w, sc)
 	}
 	w.WriteVarBytes([]byte(aer.FaultException))
 }
diff --git a/pkg/vm/stackitem/serialization.go b/pkg/vm/stackitem/serialization.go
index fcc0c3c2b..0ddccfe61 100644
--- a/pkg/vm/stackitem/serialization.go
+++ b/pkg/vm/stackitem/serialization.go
@@ -27,8 +27,8 @@ var ErrRecursive = errors.New("recursive item")
 // be serialized (like Interop item or Pointer).
 var ErrUnserializable = errors.New("unserializable")
 
-// serContext is an internal serialization context.
-type serContext struct {
+// SerializationContext is a serialization context.
+type SerializationContext struct {
 	uv           [9]byte
 	data         []byte
 	allowInvalid bool
@@ -44,7 +44,7 @@ type deserContext struct {
 
 // Serialize encodes the given Item into a byte slice.
 func Serialize(item Item) ([]byte, error) {
-	sc := serContext{
+	sc := SerializationContext{
 		allowInvalid: false,
 		seen:         make(map[Item]sliceNoPointer, typicalNumOfItems),
 	}
@@ -73,7 +73,7 @@ func EncodeBinary(item Item, w *io.BinWriter) {
 // (like recursive array) is encountered, it just writes the special InvalidT
 // type of an element to the w.
 func EncodeBinaryProtected(item Item, w *io.BinWriter) {
-	sc := serContext{
+	sc := SerializationContext{
 		allowInvalid: true,
 		seen:         make(map[Item]sliceNoPointer, typicalNumOfItems),
 	}
@@ -85,7 +85,7 @@ func EncodeBinaryProtected(item Item, w *io.BinWriter) {
 	w.WriteBytes(sc.data)
 }
 
-func (w *serContext) writeArray(item Item, arr []Item, start int) error {
+func (w *SerializationContext) writeArray(item Item, arr []Item, start int) error {
 	w.seen[item] = sliceNoPointer{}
 	w.appendVarUint(uint64(len(arr)))
 	for i := range arr {
@@ -97,7 +97,36 @@ func (w *serContext) writeArray(item Item, arr []Item, start int) error {
 	return nil
 }
 
-func (w *serContext) serialize(item Item) error {
+// NewSerializationContext returns reusable stack item serialization context.
+func NewSerializationContext() *SerializationContext {
+	return &SerializationContext{
+		seen: make(map[Item]sliceNoPointer, typicalNumOfItems),
+	}
+}
+
+// Serialize returns flat slice of bytes with the given item. The process can be protected
+// from bad elements if appropriate flag is given (otherwise an error is returned on
+// encountering any of them). The buffer returned is only valid until the call to Serialize.
+func (w *SerializationContext) Serialize(item Item, protected bool) ([]byte, error) {
+	w.allowInvalid = protected
+	if w.data != nil {
+		w.data = w.data[:0]
+	}
+	for k := range w.seen {
+		delete(w.seen, k)
+	}
+	err := w.serialize(item)
+	if err != nil && protected {
+		if w.data == nil {
+			w.data = make([]byte, 0, 1)
+		}
+		w.data = append(w.data[:0], byte(InvalidT))
+		err = nil
+	}
+	return w.data, err
+}
+
+func (w *SerializationContext) serialize(item Item) error {
 	if v, ok := w.seen[item]; ok {
 		if v.start == v.end {
 			return ErrRecursive
@@ -180,7 +209,7 @@ func (w *serContext) serialize(item Item) error {
 	return nil
 }
 
-func (w *serContext) appendVarUint(val uint64) {
+func (w *SerializationContext) appendVarUint(val uint64) {
 	n := io.PutVarUint(w.uv[:], val)
 	w.data = append(w.data, w.uv[:n]...)
 }