core: add flags to Storage.Find

It can be iterated over keys, values or both.
Prefix can be stripped.

examples/iterator/iterator.go

@@ -8,7 +8,7 @@ import (
 
 // NotifyKeysAndValues sends notification with `foo` storage keys and values
 func NotifyKeysAndValues() bool {
-	iter := storage.Find(storage.GetContext(), []byte("foo"))
+	iter := storage.Find(storage.GetContext(), []byte("foo"), storage.None)
 	for iterator.Next(iter) {
 		runtime.Notify("found storage key-value pair", iterator.Value(iter))
 	}

examples/storage/storage.go

@@ -32,7 +32,7 @@ func Delete(key []byte) bool {
 
 // Find returns an array of key-value pairs with key that matched the passed value
 func Find(value []byte) []string {
-	iter := storage.Find(ctx, value)
+	iter := storage.Find(ctx, value, storage.None)
 	result := []string{}
 	for iterator.Next(iter) {
 		val := iterator.Value(iter).([]string)

pkg/compiler/syscall_test.go

@@ -4,7 +4,9 @@ import (
 	"math/big"
 	"testing"
 
+	istorage "github.com/nspcc-dev/neo-go/pkg/core/interop/storage"
 	"github.com/nspcc-dev/neo-go/pkg/interop/contract"
+	"github.com/nspcc-dev/neo-go/pkg/interop/storage"
 	"github.com/nspcc-dev/neo-go/pkg/smartcontract/callflag"
 	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
 	"github.com/stretchr/testify/assert"
@@ -23,6 +25,13 @@ func TestCallFlags(t *testing.T) {
 	require.EqualValues(t, contract.NoneFlag, callflag.NoneFlag)
 }
 
+func TestFindFlags(t *testing.T) {
+	require.EqualValues(t, storage.None, istorage.FindDefault)
+	require.EqualValues(t, storage.KeysOnly, istorage.FindKeysOnly)
+	require.EqualValues(t, storage.RemovePrefix, istorage.FindRemovePrefix)
+	require.EqualValues(t, storage.ValuesOnly, istorage.FindValuesOnly)
+}
+
 func TestStoragePutGet(t *testing.T) {
 	src := `
 		package foo

pkg/core/interop/storage/find.go

@@ -0,0 +1,51 @@
+package storage
+
+import "github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
+
+// Storage iterator options.
+const (
+	FindDefault      = 0
+	FindKeysOnly     = 1 << 0
+	FindRemovePrefix = 1 << 1
+	FindValuesOnly   = 1 << 2
+
+	FindAll = FindDefault | FindKeysOnly | FindRemovePrefix | FindValuesOnly
+)
+
+type Iterator struct {
+	m     []stackitem.MapElement
+	opts  int64
+	index int
+}
+
+func NewIterator(m *stackitem.Map, opts int64) *Iterator {
+	return &Iterator{
+		m:     m.Value().([]stackitem.MapElement),
+		opts:  opts,
+		index: -1,
+	}
+}
+
+func (s *Iterator) Next() bool {
+	if s.index < len(s.m) {
+		s.index += 1
+	}
+	return s.index < len(s.m)
+}
+
+func (s *Iterator) Value() stackitem.Item {
+	key := s.m[s.index].Key.Value().([]byte)
+	if s.opts&FindRemovePrefix != 0 {
+		key = key[1:]
+	}
+	if s.opts&FindKeysOnly != 0 {
+		return stackitem.NewByteArray(key)
+	}
+	if s.opts&FindValuesOnly != 0 {
+		return s.m[s.index].Value
+	}
+	return stackitem.NewStruct([]stackitem.Item{
+		stackitem.NewByteArray(key),
+		s.m[s.index].Value,
+	})
+}

pkg/core/interop_neo.go

@@ -7,11 +7,14 @@ import (
 	"sort"
 
 	"github.com/nspcc-dev/neo-go/pkg/core/interop"
-	"github.com/nspcc-dev/neo-go/pkg/vm"
+	"github.com/nspcc-dev/neo-go/pkg/core/interop/storage"
 	"github.com/nspcc-dev/neo-go/pkg/vm/stackitem"
 )
 
-var errGasLimitExceeded = errors.New("gas limit exceeded")
+var (
+	errGasLimitExceeded   = errors.New("gas limit exceeded")
+	errFindInvalidOptions = errors.New("invalid Find options")
+)
 
 // storageFind finds stored key-value pair.
 func storageFind(ic *interop.Context) error {
@@ -21,6 +24,14 @@ func storageFind(ic *interop.Context) error {
 		return fmt.Errorf("%T is not a StorageContext", stcInterface)
 	}
 	prefix := ic.VM.Estack().Pop().Bytes()
+	opts := ic.VM.Estack().Pop().BigInt().Int64()
+	if opts&^storage.FindAll != 0 {
+		return fmt.Errorf("%w: unknown flag", errFindInvalidOptions)
+	}
+	if opts&storage.FindValuesOnly != 0 &&
+		opts&(storage.FindKeysOnly|storage.FindRemovePrefix) != 0 {
+		return fmt.Errorf("%w: KeysOnly conflicts with ValuesOnly", errFindInvalidOptions)
+	}
 	siMap, err := ic.DAO.GetStorageItemsWithPrefix(stc.ID, prefix)
 	if err != nil {
 		return err
@@ -35,8 +46,8 @@ func storageFind(ic *interop.Context) error {
 			filteredMap.Value().([]stackitem.MapElement)[j].Key.Value().([]byte)) == -1
 	})
 
-	item := vm.NewMapIterator(filteredMap)
+	item := storage.NewIterator(filteredMap, opts)
-	ic.VM.Estack().PushVal(item)
+	ic.VM.Estack().PushVal(stackitem.NewInterop(item))
 
 	return nil
 }

pkg/core/interop_neo_test.go

@@ -8,6 +8,7 @@ import (
 	"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/core/interop/iterator"
+	istorage "github.com/nspcc-dev/neo-go/pkg/core/interop/storage"
 	"github.com/nspcc-dev/neo-go/pkg/core/state"
 	"github.com/nspcc-dev/neo-go/pkg/core/storage"
 	"github.com/nspcc-dev/neo-go/pkg/core/transaction"
@@ -58,8 +59,9 @@ func TestStorageFind(t *testing.T) {
 		require.NoError(t, err)
 	}
 
-	t.Run("normal invocation", func(t *testing.T) {
+	testFind := func(t *testing.T, prefix byte, opts int64, expected []stackitem.Item) {
-		v.Estack().PushVal([]byte{0x01})
+		v.Estack().PushVal(opts)
+		v.Estack().PushVal([]byte{prefix})
 		v.Estack().PushVal(stackitem.NewInterop(&StorageContext{ID: id}))
 
 		err := storageFind(context)
@@ -68,38 +70,71 @@ func TestStorageFind(t *testing.T) {
 		var iter *stackitem.Interop
 		require.NotPanics(t, func() { iter = v.Estack().Pop().Interop() })
 
-		for _, id := range []int{2, 0} { // sorted indices with mathing prefix
+		for i := range expected { // sorted indices with mathing prefix
 			v.Estack().PushVal(iter)
 			require.NoError(t, iterator.Next(context))
 			require.True(t, v.Estack().Pop().Bool())
 
 			v.Estack().PushVal(iter)
 			require.NoError(t, iterator.Value(context))
-
+			require.Equal(t, expected[i], v.Estack().Pop().Item())
-			kv, ok := v.Estack().Pop().Value().([]stackitem.Item)
-			require.True(t, ok)
-			require.Len(t, kv, 2)
-			require.Equal(t, skeys[id], kv[0].Value())
-			require.Equal(t, items[id].Value, kv[1].Value())
 		}
 
 		v.Estack().PushVal(iter)
 		require.NoError(t, iterator.Next(context))
 		require.False(t, v.Estack().Pop().Bool())
+	}
+
+	t.Run("normal invocation", func(t *testing.T) {
+		testFind(t, 0x01, istorage.FindDefault, []stackitem.Item{
+			stackitem.NewStruct([]stackitem.Item{
+				stackitem.NewByteArray(skeys[2]),
+				stackitem.NewByteArray(items[2].Value),
+			}),
+			stackitem.NewStruct([]stackitem.Item{
+				stackitem.NewByteArray(skeys[0]),
+				stackitem.NewByteArray(items[0].Value),
+			}),
+		})
+	})
+
+	t.Run("keys only", func(t *testing.T) {
+		testFind(t, 0x01, istorage.FindKeysOnly, []stackitem.Item{
+			stackitem.NewByteArray(skeys[2]),
+			stackitem.NewByteArray(skeys[0]),
+		})
+	})
+	t.Run("remove prefix", func(t *testing.T) {
+		testFind(t, 0x01, istorage.FindKeysOnly|istorage.FindRemovePrefix, []stackitem.Item{
+			stackitem.NewByteArray(skeys[2][1:]),
+			stackitem.NewByteArray(skeys[0][1:]),
+		})
+	})
+	t.Run("values only", func(t *testing.T) {
+		testFind(t, 0x01, istorage.FindValuesOnly, []stackitem.Item{
+			stackitem.NewByteArray(items[2].Value),
+			stackitem.NewByteArray(items[0].Value),
+		})
 	})
 
 	t.Run("normal invocation, empty result", func(t *testing.T) {
-		v.Estack().PushVal([]byte{0x03})
+		testFind(t, 0x03, istorage.FindDefault, nil)
-		v.Estack().PushVal(stackitem.NewInterop(&StorageContext{ID: id}))
-
-		err := storageFind(context)
-		require.NoError(t, err)
-
-		require.NoError(t, iterator.Next(context))
-		require.False(t, v.Estack().Pop().Bool())
 	})
 
+	t.Run("invalid options", func(t *testing.T) {
+		invalid := []int64{
+			istorage.FindKeysOnly | istorage.FindValuesOnly,
+			^istorage.FindAll,
+		}
+		for _, opts := range invalid {
+			v.Estack().PushVal(opts)
+			v.Estack().PushVal([]byte{0x01})
+			v.Estack().PushVal(stackitem.NewInterop(&StorageContext{ID: id}))
+			require.Error(t, storageFind(context))
+		}
+	})
 	t.Run("invalid type for StorageContext", func(t *testing.T) {
+		v.Estack().PushVal(istorage.FindDefault)
 		v.Estack().PushVal([]byte{0x01})
 		v.Estack().PushVal(stackitem.NewInterop(nil))
 
@@ -109,6 +144,7 @@ func TestStorageFind(t *testing.T) {
 	t.Run("invalid id", func(t *testing.T) {
 		invalidID := id + 1
 
+		v.Estack().PushVal(istorage.FindDefault)
 		v.Estack().PushVal([]byte{0x01})
 		v.Estack().PushVal(stackitem.NewInterop(&StorageContext{ID: invalidID}))
 

pkg/interop/iterator/iterator.go

@@ -28,6 +28,7 @@ func Next(it Iterator) bool {
 // successful Next call. This function uses `System.Iterator.Value` syscall.
 // For slices the result is just value.
 // For maps the result can be casted to a slice of 2 elements: key and value.
+// For storage iterators refer to `storage.FindFlags` documentation.
 func Value(it Iterator) interface{} {
 	return nil
 }

pkg/interop/storage/storage.go

@@ -14,6 +14,20 @@ import "github.com/nspcc-dev/neo-go/pkg/interop/iterator"
 // to Neo .net framework's StorageContext class.
 type Context struct{}
 
+// FindFlags represents parameters to `Find` iterator.
+type FindFlags byte
+
+const (
+	// None is default option. Iterator values are key-value pairs.
+	None FindFlags = 0
+	// KeysOnly is used for iterating over keys.
+	KeysOnly FindFlags = 1 << 0
+	// RemovePrefix is used for stripping 1-byte prefix from keys.
+	RemovePrefix FindFlags = 1 << 1
+	// ValuesOnly is used for iterating over values.
+	ValuesOnly FindFlags = 1 << 2
+)
+
 // ConvertContextToReadOnly returns new context from the given one, but with
 // writing capability turned off, so that you could only invoke Get and Find
 // using this new Context. If Context is already read-only this function is a
@@ -58,4 +72,6 @@ func Delete(ctx Context, key interface{}) {}
 // that match the given key (contain it as a prefix). See Put documentation on
 // possible key types and iterator package documentation on how to use the
 // returned value. This function uses `System.Storage.Find` syscall.
-func Find(ctx Context, key interface{}) iterator.Iterator { return iterator.Iterator{} }
+func Find(ctx Context, key interface{}, options FindFlags) iterator.Iterator {
+	return iterator.Iterator{}
+}