core: add PickN flags to Storage.Find

Allow to pick items by index from serialized struct or array.

pkg/compiler/syscall_test.go

@@ -31,6 +31,8 @@ func TestFindFlags(t *testing.T) {
 	require.EqualValues(t, storage.RemovePrefix, istorage.FindRemovePrefix)
 	require.EqualValues(t, storage.ValuesOnly, istorage.FindValuesOnly)
 	require.EqualValues(t, storage.DeserializeValues, istorage.FindDeserialize)
+	require.EqualValues(t, storage.PickField0, istorage.FindPick0)
+	require.EqualValues(t, storage.PickField1, istorage.FindPick1)
 }
 
 func TestStoragePutGet(t *testing.T) {

pkg/core/interop/storage/find.go

@@ -9,9 +9,11 @@ const (
 	FindRemovePrefix = 1 << 1
 	FindValuesOnly   = 1 << 2
 	FindDeserialize  = 1 << 3
+	FindPick0        = 1 << 4
+	FindPick1        = 1 << 5
 
 	FindAll = FindDefault | FindKeysOnly | FindRemovePrefix | FindValuesOnly |
-		FindDeserialize
+		FindDeserialize | FindPick0 | FindPick1
 )
 
 type Iterator struct {
@@ -52,6 +54,11 @@ func (s *Iterator) Value() stackitem.Item {
 			panic(err)
 		}
 	}
+	if s.opts&FindPick0 != 0 {
+		value = value.Value().([]stackitem.Item)[0]
+	} else if s.opts&FindPick1 != 0 {
+		value = value.Value().([]stackitem.Item)[1]
+	}
 	if s.opts&FindValuesOnly != 0 {
 		return value
 	}

pkg/core/interop_neo.go

@@ -28,13 +28,20 @@ func storageFind(ic *interop.Context) error {
 	if opts&^storage.FindAll != 0 {
 		return fmt.Errorf("%w: unknown flag", errFindInvalidOptions)
 	}
-	if opts&storage.FindKeysOnly != 0 && opts&storage.FindDeserialize != 0 {
+	if opts&storage.FindKeysOnly != 0 &&
+		opts&(storage.FindDeserialize|storage.FindPick0|storage.FindPick1) != 0 {
 		return fmt.Errorf("%w KeysOnly conflicts with other options", errFindInvalidOptions)
 	}
 	if opts&storage.FindValuesOnly != 0 &&
 		opts&(storage.FindKeysOnly|storage.FindRemovePrefix) != 0 {
 		return fmt.Errorf("%w: KeysOnly conflicts with ValuesOnly", errFindInvalidOptions)
 	}
+	if opts&storage.FindPick0 != 0 && opts&storage.FindPick1 != 0 {
+		return fmt.Errorf("%w: Pick0 conflicts with Pick1", errFindInvalidOptions)
+	}
+	if opts&storage.FindDeserialize == 0 && (opts&storage.FindPick0 != 0 || opts&storage.FindPick1 != 0) {
+		return fmt.Errorf("%w: PickN is specified without Deserialize", errFindInvalidOptions)
+	}
 	siMap, err := ic.DAO.GetStorageItemsWithPrefix(stc.ID, prefix)
 	if err != nil {
 		return err

pkg/core/interop_neo_test.go

@@ -1,6 +1,7 @@
 package core
 
 import (
+	"math/big"
 	"testing"
 
 	"github.com/nspcc-dev/neo-go/pkg/config/netmode"
@@ -37,8 +38,20 @@ func TestStorageFind(t *testing.T) {
 	v, contractState, context, chain := createVMAndContractState(t)
 	defer chain.Close()
 
+	arr := []stackitem.Item{
+		stackitem.NewBigInteger(big.NewInt(42)),
+		stackitem.NewByteArray([]byte("second")),
+		stackitem.Null{},
+	}
+	rawArr, err := stackitem.SerializeItem(stackitem.NewArray(arr))
+	require.NoError(t, err)
+	rawArr0, err := stackitem.SerializeItem(stackitem.NewArray(arr[:0]))
+	require.NoError(t, err)
+	rawArr1, err := stackitem.SerializeItem(stackitem.NewArray(arr[:1]))
+	require.NoError(t, err)
+
 	skeys := [][]byte{{0x01, 0x02}, {0x02, 0x01}, {0x01, 0x01},
-		{0x04, 0x00}, {0x05, 0x00}}
+		{0x04, 0x00}, {0x05, 0x00}, {0x06}, {0x07}, {0x08}}
 	items := []*state.StorageItem{
 		{
 			Value: []byte{0x01, 0x02, 0x03, 0x04},
@@ -55,6 +68,15 @@ func TestStorageFind(t *testing.T) {
 		{
 			Value: []byte{0xFF, 0xFF},
 		},
+		{
+			Value: rawArr,
+		},
+		{
+			Value: rawArr0,
+		},
+		{
+			Value: rawArr1,
+		},
 	}
 
 	require.NoError(t, chain.contracts.Management.PutContractState(chain.dao, contractState))
@@ -83,6 +105,10 @@ func TestStorageFind(t *testing.T) {
 			require.True(t, v.Estack().Pop().Bool())
 
 			v.Estack().PushVal(iter)
+			if expected[i] == nil {
+				require.Panics(t, func() { _ = iterator.Value(context) })
+				return
+			}
 			require.NoError(t, iterator.Value(context))
 			require.Equal(t, expected[i], v.Estack().Pop().Item())
 		}
@@ -144,6 +170,19 @@ func TestStorageFind(t *testing.T) {
 			require.Panics(t, func() { _ = iterator.Value(context) })
 		})
 	})
+	t.Run("PickN", func(t *testing.T) {
+		testFind(t, 0x06, istorage.FindPick0|istorage.FindValuesOnly|istorage.FindDeserialize, arr[:1])
+		testFind(t, 0x06, istorage.FindPick1|istorage.FindValuesOnly|istorage.FindDeserialize, arr[1:2])
+		// Array with 0 elements.
+		testFind(t, 0x07, istorage.FindPick0|istorage.FindValuesOnly|istorage.FindDeserialize,
+			[]stackitem.Item{nil})
+		// Array with 1 element.
+		testFind(t, 0x08, istorage.FindPick1|istorage.FindValuesOnly|istorage.FindDeserialize,
+			[]stackitem.Item{nil})
+		// Not an array, but serialized ByteArray.
+		testFind(t, 0x04, istorage.FindPick1|istorage.FindValuesOnly|istorage.FindDeserialize,
+			[]stackitem.Item{nil})
+	})
 
 	t.Run("normal invocation, empty result", func(t *testing.T) {
 		testFind(t, 0x03, istorage.FindDefault, nil)
@@ -154,6 +193,9 @@ func TestStorageFind(t *testing.T) {
 			istorage.FindKeysOnly | istorage.FindValuesOnly,
 			^istorage.FindAll,
 			istorage.FindKeysOnly | istorage.FindDeserialize,
+			istorage.FindPick0,
+			istorage.FindPick0 | istorage.FindPick1 | istorage.FindDeserialize,
+			istorage.FindPick0 | istorage.FindPick1,
 		}
 		for _, opts := range invalid {
 			v.Estack().PushVal(opts)

pkg/interop/storage/storage.go

@@ -29,6 +29,10 @@ const (
 	// DeserializeValues is used for deserializing values on-the-fly.
 	// It can be combined with other options.
 	DeserializeValues FindFlags = 1 << 3
+	// PickField0 is used to get first field in a serialized struct or array.
+	PickField0 FindFlags = 1 << 4
+	// PickField1 is used to get second field in a serialized struct or array.
+	PickField1 FindFlags = 1 << 5
 )
 
 // ConvertContextToReadOnly returns new context from the given one, but with