Merge pull request #992 from nspcc-dev/fix/state

vm,dao: return storage iterator from DAO in Storage.Find interop

pkg/core/blockchain.go

@@ -1147,7 +1147,7 @@ func (bc *Blockchain) GetStorageItem(scripthash util.Uint160, key []byte) *state
 
 // GetStorageItems returns all storage items for a given scripthash.
 func (bc *Blockchain) GetStorageItems(hash util.Uint160) (map[string]*state.StorageItem, error) {
-	siMap, err := bc.dao.GetStorageItems(hash)
+	siMap, err := bc.dao.GetStorageItems(hash, nil)
 	if err != nil {
 		return nil, err
 	}

pkg/core/dao/cacheddao.go

@@ -315,9 +315,55 @@ func (cd *Cached) DeleteStorageItem(scripthash util.Uint160, key []byte) error {
 	return nil
 }
 
+// StorageIteratorFunc is a function returning key-value pair or error.
+type StorageIteratorFunc func() ([]byte, []byte, error)
+
+// GetStorageItemsIterator returns iterator over all storage items.
+// Function returned can be called until first error.
+func (cd *Cached) GetStorageItemsIterator(hash util.Uint160, prefix []byte) (StorageIteratorFunc, error) {
+	items, err := cd.DAO.GetStorageItems(hash, prefix)
+	if err != nil {
+		return nil, err
+	}
+
+	sort.Slice(items, func(i, j int) bool { return bytes.Compare(items[i].Key, items[j].Key) == -1 })
+
+	cache := cd.storage.getItems(hash)
+
+	var getItemFromCache StorageIteratorFunc
+	keyIndex := -1
+	getItemFromCache = func() ([]byte, []byte, error) {
+		keyIndex++
+		for ; keyIndex < len(cd.storage.keys[hash]); keyIndex++ {
+			k := cd.storage.keys[hash][keyIndex]
+			v := cache[k]
+			if v.State != delOp && bytes.HasPrefix([]byte(k), prefix) {
+				val := make([]byte, len(v.StorageItem.Value))
+				copy(val, v.StorageItem.Value)
+				return []byte(k), val, nil
+			}
+		}
+		return nil, nil, errors.New("no more items")
+	}
+
+	var f func() ([]byte, []byte, error)
+	index := -1
+	f = func() ([]byte, []byte, error) {
+		index++
+		for ; index < len(items); index++ {
+			_, ok := cache[string(items[index].Key)]
+			if !ok {
+				return items[index].Key, items[index].Value, nil
+			}
+		}
+		return getItemFromCache()
+	}
+	return f, nil
+}
+
 // GetStorageItems returns all storage items for a given scripthash.
-func (cd *Cached) GetStorageItems(hash util.Uint160) ([]StorageItemWithKey, error) {
+func (cd *Cached) GetStorageItems(hash util.Uint160, prefix []byte) ([]StorageItemWithKey, error) {
-	items, err := cd.DAO.GetStorageItems(hash)
+	items, err := cd.DAO.GetStorageItems(hash, prefix)
 	if err != nil {
 		return nil, err
 	}

pkg/core/dao/dao.go

@@ -35,7 +35,7 @@ type DAO interface {
 	GetNEP5Balances(acc util.Uint160) (*state.NEP5Balances, error)
 	GetNEP5TransferLog(acc util.Uint160, index uint32) (*state.NEP5TransferLog, error)
 	GetStorageItem(scripthash util.Uint160, key []byte) *state.StorageItem
-	GetStorageItems(hash util.Uint160) ([]StorageItemWithKey, error)
+	GetStorageItems(hash util.Uint160, prefix []byte) ([]StorageItemWithKey, error)
 	GetTransaction(hash util.Uint256) (*transaction.Transaction, uint32, error)
 	GetUnspentCoinState(hash util.Uint256) (*state.UnspentCoin, error)
 	GetValidatorState(publicKey *keys.PublicKey) (*state.Validator, error)
@@ -442,7 +442,7 @@ type StorageItemWithKey struct {
 }
 
 // GetStorageItems returns all storage items for a given scripthash.
-func (dao *Simple) GetStorageItems(hash util.Uint160) ([]StorageItemWithKey, error) {
+func (dao *Simple) GetStorageItems(hash util.Uint160, prefix []byte) ([]StorageItemWithKey, error) {
 	var res []StorageItemWithKey
 	var err error
 
@@ -462,7 +462,7 @@ func (dao *Simple) GetStorageItems(hash util.Uint160) ([]StorageItemWithKey, err
 		s.Key = k[21:]
 		res = append(res, s)
 	}
-	dao.Store.Seek(storage.AppendPrefix(storage.STStorage, hash.BytesLE()), saveToMap)
+	dao.Store.Seek(makeStorageItemKey(hash, prefix), saveToMap)
 	if err != nil {
 		return nil, err
 	}

pkg/core/interop_neo.go

@@ -1,7 +1,6 @@
 package core
 
 import (
-	"bytes"
 	"errors"
 	"fmt"
 	"math"
@@ -463,22 +462,12 @@ func (ic *interopContext) storageFind(v *vm.VM) error {
 		return err
 	}
 	pref := v.Estack().Pop().Bytes()
-	siMap, err := ic.dao.GetStorageItems(stc.ScriptHash)
+	next, err := ic.dao.GetStorageItemsIterator(stc.ScriptHash, pref)
 	if err != nil {
 		return err
 	}
-
+	item := newStorageIterator(next)
-	filteredMap := vm.NewMapItem()
+	v.Estack().PushVal(vm.NewInteropItem(item))
-	for i := range siMap {
-		k := siMap[i].Key
-		if bytes.HasPrefix(k, pref) {
-			filteredMap.Add(vm.NewByteArrayItem(siMap[i].Key),
-				vm.NewByteArrayItem(siMap[i].Value))
-		}
-	}
-
-	item := vm.NewMapIterator(filteredMap)
-	v.Estack().PushVal(item)
 
 	return nil
 }
@@ -593,7 +582,7 @@ func (ic *interopContext) contractMigrate(v *vm.VM) error {
 		}
 		if contract.HasStorage() {
 			hash := getContextScriptHash(v, 0)
-			siMap, err := ic.dao.GetStorageItems(hash)
+			siMap, err := ic.dao.GetStorageItems(hash, nil)
 			if err != nil {
 				return err
 			}

pkg/core/interop_system.go

@@ -561,7 +561,7 @@ func (ic *interopContext) contractDestroy(v *vm.VM) error {
 		return err
 	}
 	if cs.HasStorage() {
-		siMap, err := ic.dao.GetStorageItems(hash)
+		siMap, err := ic.dao.GetStorageItems(hash, nil)
 		if err != nil {
 			return err
 		}

pkg/core/storage_find.go

@@ -0,0 +1,46 @@
+package core
+
+import (
+	"github.com/nspcc-dev/neo-go/pkg/core/dao"
+	"github.com/nspcc-dev/neo-go/pkg/vm"
+)
+
+type storageWrapper struct {
+	next       dao.StorageIteratorFunc
+	key, value vm.StackItem
+	finished   bool
+}
+
+var _ vm.Iterator = (*storageWrapper)(nil)
+
+// newStorageIterator returns new storage iterator from the `next()` callback.
+func newStorageIterator(next dao.StorageIteratorFunc) *storageWrapper {
+	return &storageWrapper{
+		next: next,
+	}
+}
+
+// Next implements vm.Iterator interface.
+func (s *storageWrapper) Next() bool {
+	if s.finished {
+		return false
+	}
+	key, value, err := s.next()
+	if err != nil {
+		s.finished = true
+		return false
+	}
+	s.key = vm.NewByteArrayItem(key)
+	s.value = vm.NewByteArrayItem(value)
+	return true
+}
+
+// Value implements vm.Iterator interface.
+func (s *storageWrapper) Value() vm.StackItem {
+	return s.value
+}
+
+// Key implements vm.Iterator interface.
+func (s *storageWrapper) Key() vm.StackItem {
+	return s.key
+}

pkg/vm/interop.go

@@ -205,9 +205,9 @@ func NewMapIterator(m *MapItem) *InteropItem {
 // IteratorConcat handles syscall Neo.Iterator.Concat.
 func IteratorConcat(v *VM) error {
 	iop1 := v.Estack().Pop().Interop()
-	iter1 := iop1.value.(iterator)
+	iter1 := iop1.value.(Iterator)
 	iop2 := v.Estack().Pop().Interop()
-	iter2 := iop2.value.(iterator)
+	iter2 := iop2.value.(Iterator)
 
 	v.Estack().Push(&Element{value: NewInteropItem(
 		&concatIter{
@@ -222,7 +222,7 @@ func IteratorConcat(v *VM) error {
 // IteratorKey handles syscall Neo.Iterator.Key.
 func IteratorKey(v *VM) error {
 	iop := v.estack.Pop().Interop()
-	iter := iop.value.(iterator)
+	iter := iop.value.(Iterator)
 	v.Estack().Push(&Element{value: iter.Key()})
 
 	return nil
@@ -231,7 +231,7 @@ func IteratorKey(v *VM) error {
 // IteratorKeys handles syscall Neo.Iterator.Keys.
 func IteratorKeys(v *VM) error {
 	iop := v.estack.Pop().Interop()
-	iter := iop.value.(iterator)
+	iter := iop.value.(Iterator)
 	v.Estack().Push(&Element{value: NewInteropItem(
 		&keysWrapper{iter},
 	)})
@@ -242,7 +242,7 @@ func IteratorKeys(v *VM) error {
 // IteratorValues handles syscall Neo.Iterator.Values.
 func IteratorValues(v *VM) error {
 	iop := v.estack.Pop().Interop()
-	iter := iop.value.(iterator)
+	iter := iop.value.(Iterator)
 	v.Estack().Push(&Element{value: NewInteropItem(
 		&valuesWrapper{iter},
 	)})

pkg/vm/interop_iterators.go

@@ -18,7 +18,8 @@ type (
 )
 
 type (
-	iterator interface {
+	// Iterator defined public interface for VM's iterator type.
+	Iterator interface {
 		enumerator
 		Key() StackItem
 	}
@@ -29,16 +30,16 @@ type (
 	}
 
 	concatIter struct {
-		current iterator
+		current Iterator
-		second  iterator
+		second  Iterator
 	}
 
 	keysWrapper struct {
-		iter iterator
+		iter Iterator
 	}
 
 	valuesWrapper struct {
-		iter iterator
+		iter Iterator
 	}
 )