Speedup stackitem json serialization

pkg/io/binaryWriter.go

@@ -11,6 +11,7 @@ import (
 // from a struct with many fields.
 type BinWriter struct {
 	w   io.Writer
+	uv  []byte
 	u64 []byte
 	u32 []byte
 	u16 []byte
@@ -20,11 +21,12 @@ type BinWriter struct {
 
 // NewBinWriterFromIO makes a BinWriter from io.Writer.
 func NewBinWriterFromIO(iow io.Writer) *BinWriter {
-	u64 := make([]byte, 8)
+	uv := make([]byte, 9)
+	u64 := uv[:8]
 	u32 := u64[:4]
 	u16 := u64[:2]
 	u8 := u64[:1]
-	return &BinWriter{w: iow, u64: u64, u32: u32, u16: u16, u8: u8}
+	return &BinWriter{w: iow, uv: uv, u64: u64, u32: u32, u16: u16, u8: u8}
 }
 
 // WriteU64LE writes an uint64 value into the underlying io.Writer in
@@ -106,23 +108,31 @@ func (w *BinWriter) WriteVarUint(val uint64) {
 		return
 	}
 
+	n := PutVarUint(w.uv, val)
+	w.WriteBytes(w.uv[:n])
+}
+
+// PutVarUint puts val in varint form to the pre-allocated buffer.
+func PutVarUint(data []byte, val uint64) int {
+	_ = data[8]
 	if val < 0xfd {
-		w.WriteB(byte(val))
-		return
+		data[0] = byte(val)
+		return 1
 	}
 	if val < 0xFFFF {
-		w.WriteB(byte(0xfd))
-		w.WriteU16LE(uint16(val))
-		return
+		data[0] = byte(0xfd)
+		binary.LittleEndian.PutUint16(data[1:], uint16(val))
+		return 3
 	}
 	if val < 0xFFFFFFFF {
-		w.WriteB(byte(0xfe))
-		w.WriteU32LE(uint32(val))
-		return
+		data[0] = byte(0xfe)
+		binary.LittleEndian.PutUint32(data[1:], uint32(val))
+		return 5
 	}
 
-	w.WriteB(byte(0xff))
-	w.WriteU64LE(val)
+	data[0] = byte(0xff)
+	binary.LittleEndian.PutUint64(data[1:], val)
+	return 9
 }
 
 // WriteBytes writes a variable byte into the underlying io.Writer without prefix.

pkg/vm/stackitem/json.go

@@ -9,8 +9,6 @@ import (
 	gio "io"
 	"math"
 	"math/big"
-
-	"github.com/nspcc-dev/neo-go/pkg/io"
 )
 
 // decoder is a wrapper around json.Decoder helping to mimic C# json decoder behaviour.
@@ -43,87 +41,112 @@ var ErrTooDeep = errors.New("too deep")
 //   Array, Struct -> array
 //   Map -> map with keys as UTF-8 bytes
 func ToJSON(item Item) ([]byte, error) {
-	buf := io.NewBufBinWriter()
-	toJSON(buf, item)
-	if buf.Err != nil {
-		return nil, buf.Err
-	}
-	return buf.Bytes(), nil
+	seen := make(map[Item]sliceNoPointer)
+	return toJSON(nil, seen, item)
+}
+
+// sliceNoPointer represents sub-slice of a known slice.
+// It doesn't contain pointer and uses less memory than `[]byte`.
+type sliceNoPointer struct {
+	start, end int
 }
 
-func toJSON(buf *io.BufBinWriter, item Item) {
-	w := buf.BinWriter
-	if w.Err != nil {
-		return
-	} else if buf.Len() > MaxSize {
-		w.Err = errTooBigSize
+func toJSON(data []byte, seen map[Item]sliceNoPointer, item Item) ([]byte, error) {
+	if len(data) > MaxSize {
+		return nil, errTooBigSize
+	}
+
+	if old, ok := seen[item]; ok {
+		if len(data)+old.end-old.start > MaxSize {
+			return nil, errTooBigSize
+		}
+		return append(data, data[old.start:old.end]...), nil
 	}
+
+	start := len(data)
+	var err error
+
 	switch it := item.(type) {
 	case *Array, *Struct:
-		w.WriteB('[')
-		items := it.Value().([]Item)
+		var items []Item
+		if a, ok := it.(*Array); ok {
+			items = a.value
+		} else {
+			items = it.(*Struct).value
+		}
+
+		data = append(data, '[')
 		for i, v := range items {
-			toJSON(buf, v)
+			data, err = toJSON(data, seen, v)
+			if err != nil {
+				return nil, err
+			}
 			if i < len(items)-1 {
-				w.WriteB(',')
+				data = append(data, ',')
 			}
 		}
-		w.WriteB(']')
+		data = append(data, ']')
+		seen[item] = sliceNoPointer{start, len(data)}
 	case *Map:
-		w.WriteB('{')
+		data = append(data, '{')
 		for i := range it.value {
 			// map key can always be converted to []byte
 			// but are not always a valid UTF-8.
-			writeJSONString(buf.BinWriter, it.value[i].Key)
-			w.WriteBytes([]byte(`:`))
-			toJSON(buf, it.value[i].Value)
+			raw, err := itemToJSONString(it.value[i].Key)
+			if err != nil {
+				return nil, err
+			}
+			data = append(data, raw...)
+			data = append(data, ':')
+			data, err = toJSON(data, seen, it.value[i].Value)
+			if err != nil {
+				return nil, err
+			}
 			if i < len(it.value)-1 {
-				w.WriteB(',')
+				data = append(data, ',')
 			}
 		}
-		w.WriteB('}')
+		data = append(data, '}')
+		seen[item] = sliceNoPointer{start, len(data)}
 	case *BigInteger:
 		if it.value.CmpAbs(big.NewInt(MaxAllowedInteger)) == 1 {
-			w.Err = fmt.Errorf("%w (MaxAllowedInteger)", ErrInvalidValue)
-			return
+			return nil, fmt.Errorf("%w (MaxAllowedInteger)", ErrInvalidValue)
 		}
-		w.WriteBytes([]byte(it.value.String()))
+		data = append(data, it.value.String()...)
 	case *ByteArray, *Buffer:
-		writeJSONString(w, it)
+		raw, err := itemToJSONString(it)
+		if err != nil {
+			return nil, err
+		}
+		data = append(data, raw...)
 	case *Bool:
 		if it.value {
-			w.WriteBytes([]byte("true"))
+			data = append(data, "true"...)
 		} else {
-			w.WriteBytes([]byte("false"))
+			data = append(data, "false"...)
 		}
 	case Null:
-		w.WriteBytes([]byte("null"))
+		data = append(data, "null"...)
 	default:
-		w.Err = fmt.Errorf("%w: %s", ErrUnserializable, it.String())
-		return
+		return nil, fmt.Errorf("%w: %s", ErrUnserializable, it.String())
 	}
-	if w.Err == nil && buf.Len() > MaxSize {
-		w.Err = errTooBigSize
+	if len(data) > MaxSize {
+		return nil, errTooBigSize
 	}
+	return data, nil
 }
 
-// writeJSONString converts it to string and writes it to w as JSON value
+// itemToJSONString converts it to string
 // surrounded in quotes with control characters escaped.
-func writeJSONString(w *io.BinWriter, it Item) {
-	if w.Err != nil {
-		return
-	}
+func itemToJSONString(it Item) ([]byte, error) {
 	s, err := ToString(it)
 	if err != nil {
-		w.Err = err
-		return
+		return nil, err
 	}
 	data, _ := json.Marshal(s) // error never occurs because `ToString` checks for validity
 
 	// ref https://github.com/neo-project/neo-modules/issues/375 and https://github.com/dotnet/runtime/issues/35281
-	data = bytes.Replace(data, []byte{'+'}, []byte("\\u002B"), -1)
-
-	w.WriteBytes(data)
+	return bytes.Replace(data, []byte{'+'}, []byte("\\u002B"), -1), nil
 }
 
 // FromJSON decodes Item from JSON.

pkg/vm/stackitem/json_test.go

@@ -1,6 +1,7 @@
 package stackitem
 
 import (
+	"errors"
 	"math/big"
 	"testing"
 
@@ -105,6 +106,64 @@ func TestFromToJSON(t *testing.T) {
 	})
 }
 
+func testToJSON(t *testing.T, expectedErr error, item Item) {
+	data, err := ToJSON(item)
+	if expectedErr != nil {
+		require.True(t, errors.Is(err, expectedErr), err)
+		return
+	}
+	require.NoError(t, err)
+
+	actual, err := FromJSON(data)
+	require.NoError(t, err)
+	require.Equal(t, item, actual)
+}
+
+func TestToJSONCornerCases(t *testing.T) {
+	// base64 encoding increases size by a factor of ~256/64 = 4
+	const maxSize = MaxSize / 4
+
+	bigByteArray := NewByteArray(make([]byte, maxSize/2))
+	smallByteArray := NewByteArray(make([]byte, maxSize/4))
+	t.Run("Array", func(t *testing.T) {
+		arr := NewArray([]Item{bigByteArray})
+		testToJSON(t, ErrTooBig, NewArray([]Item{arr, arr}))
+
+		arr.value[0] = smallByteArray
+		testToJSON(t, nil, NewArray([]Item{arr, arr}))
+	})
+	t.Run("big ByteArray", func(t *testing.T) {
+		testToJSON(t, ErrTooBig, NewByteArray(make([]byte, maxSize+4)))
+	})
+	t.Run("invalid Map key", func(t *testing.T) {
+		m := NewMap()
+		m.Add(Make([]byte{0xe9}), Make(true))
+		testToJSON(t, ErrInvalidValue, m)
+	})
+}
+
+// getBigArray returns array takes up a lot of storage when serialized.
+func getBigArray(depth int) *Array {
+	arr := NewArray([]Item{})
+	for i := 0; i < depth; i++ {
+		arr = NewArray([]Item{arr, arr})
+	}
+	return arr
+}
+
+func BenchmarkToJSON(b *testing.B) {
+	arr := getBigArray(15)
+
+	b.ResetTimer()
+	b.ReportAllocs()
+	for i := 0; i < b.N; i++ {
+		_, err := ToJSON(arr)
+		if err != nil {
+			b.FailNow()
+		}
+	}
+}
+
 // This test is taken from the C# code
 // https://github.com/neo-project/neo/blob/master/tests/neo.UnitTests/VM/UT_Helper.cs#L30
 func TestToJSONWithTypeCompat(t *testing.T) {