From 776bd85ded8194bd668bc9e20b6176b62f754b96 Mon Sep 17 00:00:00 2001
From: Evgenii Stratonikov <evgeniy@nspcc.ru>
Date: Wed, 27 May 2020 09:27:43 +0300
Subject: [PATCH] vm,dao: return storage iterator from DAO in Storage.Find
 interop

Reproduce behavior of the reference realization:
- if item was Put in cache after it was encountered during
  Storage.Find, it must appear twice
- checking if item is in cache must be performed in real-time
  during `Iterator.Next()`
---
 pkg/core/dao/cacheddao.go | 46 +++++++++++++++++++++++++++++++++++++++
 pkg/core/interop_neo.go   | 17 +++------------
 pkg/core/storage_find.go  | 44 +++++++++++++++++++++++++++++++++++++
 3 files changed, 93 insertions(+), 14 deletions(-)
 create mode 100644 pkg/core/storage_find.go

diff --git a/pkg/core/dao/cacheddao.go b/pkg/core/dao/cacheddao.go
index 327501157..f35627712 100644
--- a/pkg/core/dao/cacheddao.go
+++ b/pkg/core/dao/cacheddao.go
@@ -315,6 +315,52 @@ 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)
+	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 && bytes.HasPrefix(items[index].Key, prefix) {
+				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) {
 	items, err := cd.DAO.GetStorageItems(hash)
diff --git a/pkg/core/interop_neo.go b/pkg/core/interop_neo.go
index 64cb44614..ccddb26b6 100644
--- a/pkg/core/interop_neo.go
+++ b/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
 	}
-
-	filteredMap := vm.NewMapItem()
-	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)
+	item := newStorageIterator(next)
+	v.Estack().PushVal(vm.NewInteropItem(item))
 
 	return nil
 }
diff --git a/pkg/core/storage_find.go b/pkg/core/storage_find.go
new file mode 100644
index 000000000..14231896d
--- /dev/null
+++ b/pkg/core/storage_find.go
@@ -0,0 +1,44 @@
+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
+}
+
+// newStorageIterator returns new storage iterator from the `next()` callback.
+func newStorageIterator(next dao.StorageIteratorFunc) *storageWrapper {
+	return &storageWrapper{
+		next: next,
+	}
+}
+
+// Next implements 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 iterator interface.
+func (s *storageWrapper) Value() vm.StackItem {
+	return s.value
+}
+
+// Key implements iterator interface.
+func (s *storageWrapper) Key() vm.StackItem {
+	return s.key
+}