Merge pull request #2112 from nspcc-dev/gas-balance

Improve NEP17 transfer serialization

pkg/core/native/native_gas.go

@@ -59,7 +59,7 @@ func (g *GAS) increaseBalance(_ *interop.Context, _ util.Uint160, si *state.Stor
 	}
 	acc.Balance.Add(&acc.Balance, amount)
 	if acc.Balance.Sign() != 0 {
-		*si = acc.Bytes()
+		*si = acc.Bytes(nil)
 	} else {
 		*si = nil
 	}

pkg/core/state/native_state.go

@@ -7,6 +7,7 @@ import (
 	"math/big"
 
 	"github.com/nspcc-dev/neo-go/pkg/crypto/keys"
+	"github.com/nspcc-dev/neo-go/pkg/encoding/bigint"
 	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
 )
 
@@ -24,17 +25,42 @@ type NEOBalance struct {
 
 // NEP17BalanceFromBytes converts serialized NEP17Balance to structure.
 func NEP17BalanceFromBytes(b []byte) (*NEP17Balance, error) {
-	balance := new(NEP17Balance)
-	err := balanceFromBytes(b, balance)
-	if err != nil {
-		return nil, err
+	if len(b) < 4 {
+		if len(b) == 0 {
+			return new(NEP17Balance), nil
+		}
+		return nil, errors.New("invalid format")
 	}
-	return balance, nil
+	if b[0] != byte(stackitem.StructT) {
+		return nil, errors.New("not a struct")
+	}
+	if b[1] != 1 {
+		return nil, errors.New("invalid item count")
+	}
+	if st := stackitem.Type(b[2]); st != stackitem.IntegerT {
+		return nil, fmt.Errorf("invalid balance: %s", st)
+	}
+	if int(b[3]) != len(b[4:]) {
+		return nil, errors.New("invalid balance format")
+	}
+	return &NEP17Balance{Balance: *bigint.FromBytes(b[4:])}, nil
 }
 
 // Bytes returns serialized NEP17Balance.
-func (s *NEP17Balance) Bytes() []byte {
-	return balanceToBytes(s)
+func (s *NEP17Balance) Bytes(buf []byte) []byte {
+	if cap(buf) < 4+bigint.MaxBytesLen {
+		buf = make([]byte, 4, 4+bigint.MaxBytesLen)
+	} else {
+		buf = buf[:4]
+	}
+	buf[0] = byte(stackitem.StructT)
+	buf[1] = 1
+	buf[2] = byte(stackitem.IntegerT)
+
+	data := bigint.ToPreallocatedBytes(&s.Balance, buf[4:])
+	buf[3] = byte(len(data)) // max is 33, so we are ok here
+	buf = append(buf, data...)
+	return buf
 }
 
 func balanceFromBytes(b []byte, item stackitem.Convertible) error {

pkg/core/state/native_state_test.go

@@ -0,0 +1,95 @@
+package state
+
+import (
+	"testing"
+
+	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNEP17Balance_Bytes(t *testing.T) {
+	var b NEP17Balance
+	b.Balance.SetInt64(0x12345678910)
+
+	data, err := stackitem.SerializeConvertible(&b)
+	require.NoError(t, err)
+	require.Equal(t, data, b.Bytes(nil))
+
+	t.Run("reuse buffer", func(t *testing.T) {
+		buf := make([]byte, 100)
+		ret := b.Bytes(buf[:0])
+		require.Equal(t, ret, buf[:len(ret)])
+	})
+
+	actual, err := NEP17BalanceFromBytes(data)
+	require.NoError(t, err)
+	require.Equal(t, &b, actual)
+}
+
+func TestNEP17BalanceFromBytesInvalid(t *testing.T) {
+	b, err := NEP17BalanceFromBytes(nil) // 0 is ok
+	require.NoError(t, err)
+	require.Equal(t, int64(0), b.Balance.Int64())
+
+	_, err = NEP17BalanceFromBytes([]byte{byte(stackitem.StructT)})
+	require.Error(t, err)
+
+	_, err = NEP17BalanceFromBytes([]byte{byte(stackitem.IntegerT), 4, 0, 1, 2, 3})
+	require.Error(t, err)
+
+	_, err = NEP17BalanceFromBytes([]byte{byte(stackitem.StructT), 0, byte(stackitem.IntegerT), 1, 1})
+	require.Error(t, err)
+
+	_, err = NEP17BalanceFromBytes([]byte{byte(stackitem.StructT), 1, byte(stackitem.ByteArrayT), 1, 1})
+	require.Error(t, err)
+
+	_, err = NEP17BalanceFromBytes([]byte{byte(stackitem.StructT), 1, byte(stackitem.IntegerT), 2, 1})
+	require.Error(t, err)
+}
+
+func BenchmarkNEP17BalanceBytes(b *testing.B) {
+	var bl NEP17Balance
+	bl.Balance.SetInt64(0x12345678910)
+
+	b.Run("stackitem", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_, _ = stackitem.SerializeConvertible(&bl)
+		}
+	})
+	b.Run("bytes", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_ = bl.Bytes(nil)
+		}
+	})
+	b.Run("bytes, prealloc", func(b *testing.B) {
+		bs := bl.Bytes(nil)
+
+		b.ResetTimer()
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_ = bl.Bytes(bs[:0])
+		}
+	})
+}
+
+func BenchmarkNEP17BalanceFromBytes(b *testing.B) {
+	var bl NEP17Balance
+	bl.Balance.SetInt64(0x12345678910)
+
+	buf := bl.Bytes(nil)
+
+	b.Run("stackitem", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_ = stackitem.DeserializeConvertible(buf, new(NEP17Balance))
+		}
+	})
+	b.Run("from bytes", func(b *testing.B) {
+		b.ReportAllocs()
+		for i := 0; i < b.N; i++ {
+			_, _ = NEP17BalanceFromBytes(buf)
+		}
+	})
+}

pkg/core/state/nep17.go

@@ -1,6 +1,7 @@
 package state
 
 import (
+	"bytes"
 	"math/big"
 
 	"github.com/nspcc-dev/neo-go/pkg/encoding/bigint"
@@ -78,19 +79,20 @@ func (bs *NEP17TransferInfo) EncodeBinary(w *io.BinWriter) {
 
 // Append appends single transfer to a log.
 func (lg *NEP17TransferLog) Append(tr *NEP17Transfer) error {
-	w := io.NewBufBinWriter()
 	// The first entry, set up counter.
 	if len(lg.Raw) == 0 {
-		w.WriteB(1)
+		lg.Raw = append(lg.Raw, 0)
 	}
-	tr.EncodeBinary(w.BinWriter)
+
+	b := bytes.NewBuffer(lg.Raw)
+	w := io.NewBinWriterFromIO(b)
+
+	tr.EncodeBinary(w)
 	if w.Err != nil {
 		return w.Err
 	}
-	if len(lg.Raw) != 0 {
-		lg.Raw[0]++
-	}
-	lg.Raw = append(lg.Raw, w.Bytes()...)
+	lg.Raw = b.Bytes()
+	lg.Raw[0]++
 	return nil
 }
 

pkg/core/state/nep17_test.go

@@ -38,6 +38,26 @@ func TestNEP17TransferLog_Append(t *testing.T) {
 	require.True(t, cont)
 }
 
+func BenchmarkNEP17TransferLog_Append(b *testing.B) {
+	r := rand.New(rand.NewSource(time.Now().UnixNano()))
+	ts := make([]*NEP17Transfer, NEP17TransferBatchSize)
+	for i := range ts {
+		ts[i] = randomTransfer(r)
+	}
+
+	lg := new(NEP17TransferLog)
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		for _, tr := range ts {
+			err := lg.Append(tr)
+			if err != nil {
+				b.FailNow()
+			}
+		}
+	}
+}
+
 func TestNEP17Transfer_DecodeBinary(t *testing.T) {
 	expected := &NEP17Transfer{
 		Asset:     123,