Simplify Iterator APIs (#2190)

Co-authored-by: Shargon <shargon@gmail.com>
neo-project · Jan 3, 2021 · 8f98ada · 8f98ada
1 parent 59ea84b
commit 8f98ada
Show file tree

Hide file tree

Showing 22 changed files with 87 additions and 420 deletions.
diff --git a/src/neo/SmartContract/ApplicationEngine.Enumerator.cs b/src/neo/SmartContract/ApplicationEngine.Enumerator.cs
diff --git a/src/neo/SmartContract/ApplicationEngine.Iterator.cs b/src/neo/SmartContract/ApplicationEngine.Iterator.cs
@@ -1,4 +1,3 @@
-using Neo.SmartContract.Enumerators;
 using Neo.SmartContract.Iterators;
 using Neo.VM.Types;
 using System;
@@ -9,35 +8,29 @@ namespace Neo.SmartContract
     partial class ApplicationEngine
     {
         public static readonly InteropDescriptor System_Iterator_Create = Register("System.Iterator.Create", nameof(CreateIterator), 1 << 4, CallFlags.None);
-        public static readonly InteropDescriptor System_Iterator_Key = Register("System.Iterator.Key", nameof(IteratorKey), 1 << 4, CallFlags.None);
-        public static readonly InteropDescriptor System_Iterator_Keys = Register("System.Iterator.Keys", nameof(IteratorKeys), 1 << 4, CallFlags.None);
-        public static readonly InteropDescriptor System_Iterator_Values = Register("System.Iterator.Values", nameof(IteratorValues), 1 << 4, CallFlags.None);
+        public static readonly InteropDescriptor System_Iterator_Next = Register("System.Iterator.Next", nameof(IteratorNext), 1 << 15, CallFlags.None);
+        public static readonly InteropDescriptor System_Iterator_Value = Register("System.Iterator.Value", nameof(IteratorValue), 1 << 4, CallFlags.None);
 
         protected internal IIterator CreateIterator(StackItem item)
         {
             return item switch
             {
                 Array array => new ArrayWrapper(array),
-                Map map => new MapWrapper(map),
+                Map map => new MapWrapper(map, ReferenceCounter),
                 VM.Types.Buffer buffer => new ByteArrayWrapper(buffer),
                 PrimitiveType primitive => new ByteArrayWrapper(primitive),
                 _ => throw new ArgumentException()
             };
         }
 
-        protected internal PrimitiveType IteratorKey(IIterator iterator)
+        protected internal bool IteratorNext(IIterator iterator)
         {
-            return iterator.Key();
+            return iterator.Next();
         }
 
-        protected internal IEnumerator IteratorKeys(IIterator iterator)
+        protected internal StackItem IteratorValue(IIterator iterator)
         {
-            return new IteratorKeysWrapper(iterator);
-        }
-
-        protected internal IEnumerator IteratorValues(IIterator iterator)
-        {
-            return new IteratorValuesWrapper(iterator);
+            return iterator.Value();
         }
     }
 }
diff --git a/src/neo/SmartContract/ApplicationEngine.Storage.cs b/src/neo/SmartContract/ApplicationEngine.Storage.cs
@@ -1,5 +1,4 @@
 using Neo.Ledger;
-using Neo.SmartContract.Enumerators;
 using Neo.SmartContract.Iterators;
 using Neo.SmartContract.Native;
 using System;
@@ -17,7 +16,6 @@ partial class ApplicationEngine
         public static readonly InteropDescriptor System_Storage_AsReadOnly = Register("System.Storage.AsReadOnly", nameof(AsReadOnly), 1 << 4, CallFlags.ReadStates);
         public static readonly InteropDescriptor System_Storage_Get = Register("System.Storage.Get", nameof(Get), 1 << 15, CallFlags.ReadStates);
         public static readonly InteropDescriptor System_Storage_Find = Register("System.Storage.Find", nameof(Find), 1 << 15, CallFlags.ReadStates);
-        public static readonly InteropDescriptor System_Storage_FindKeys = Register("System.Storage.FindKeys", nameof(FindKeys), 1 << 15, CallFlags.ReadStates);
         public static readonly InteropDescriptor System_Storage_Put = Register("System.Storage.Put", nameof(Put), 0, CallFlags.WriteStates);
         public static readonly InteropDescriptor System_Storage_PutEx = Register("System.Storage.PutEx", nameof(PutEx), 0, CallFlags.WriteStates);
         public static readonly InteropDescriptor System_Storage_Delete = Register("System.Storage.Delete", nameof(Delete), 0, CallFlags.WriteStates);
@@ -62,16 +60,16 @@ protected internal byte[] Get(StorageContext context, byte[] key)
             })?.Value;
         }
 
-        protected internal IIterator Find(StorageContext context, byte[] prefix)
-        {
-            byte[] prefix_key = StorageKey.CreateSearchPrefix(context.Id, prefix);
-            return new StorageIterator(Snapshot.Storages.Find(prefix_key).GetEnumerator());
-        }
-
-        protected internal IEnumerator FindKeys(StorageContext context, byte[] prefix, byte removePrefix)
+        protected internal IIterator Find(StorageContext context, byte[] prefix, FindOptions options)
         {
+            if (options.HasFlag(FindOptions.KeysOnly) && options.HasFlag(FindOptions.ValuesOnly))
+                throw new ArgumentException();
+            if (options.HasFlag(FindOptions.PickField0) && options.HasFlag(FindOptions.PickField1))
+                throw new ArgumentException();
+            if ((options.HasFlag(FindOptions.PickField0) || options.HasFlag(FindOptions.PickField1)) && !options.HasFlag(FindOptions.DeserializeValues))
+                throw new ArgumentException();
             byte[] prefix_key = StorageKey.CreateSearchPrefix(context.Id, prefix);
-            return new StorageKeyEnumerator(Snapshot.Storages.Find(prefix_key).Select(p => p.Key).GetEnumerator(), removePrefix);
+            return new StorageIterator(Snapshot.Storages.Find(prefix_key).GetEnumerator(), options, ReferenceCounter);
         }
 
         protected internal void Put(StorageContext context, byte[] key, byte[] value)

diff --git a/src/neo/SmartContract/Enumerators/IEnumerator.cs b/src/neo/SmartContract/Enumerators/IEnumerator.cs
diff --git a/src/neo/SmartContract/Enumerators/IteratorKeysWrapper.cs b/src/neo/SmartContract/Enumerators/IteratorKeysWrapper.cs
diff --git a/src/neo/SmartContract/Enumerators/IteratorValuesWrapper.cs b/src/neo/SmartContract/Enumerators/IteratorValuesWrapper.cs
diff --git a/src/neo/SmartContract/Enumerators/StorageKeyEnumerator.cs b/src/neo/SmartContract/Enumerators/StorageKeyEnumerator.cs
diff --git a/src/neo/SmartContract/FindOptions.cs b/src/neo/SmartContract/FindOptions.cs
@@ -0,0 +1,17 @@
+using System;
+
+namespace Neo.SmartContract
+{
+    [Flags]
+    public enum FindOptions : byte
+    {
+        None = 0,
+
+        KeysOnly = 1 << 0,
+        RemovePrefix = 1 << 1,
+        ValuesOnly = 1 << 2,
+        DeserializeValues = 1 << 3,
+        PickField0 = 1 << 4,
+        PickField1 = 1 << 5
+    }
+}
diff --git a/src/neo/SmartContract/Iterators/ArrayWrapper.cs b/src/neo/SmartContract/Iterators/ArrayWrapper.cs
@@ -18,13 +18,6 @@ public void Dispose()
         {
         }
 
-        public PrimitiveType Key()
-        {
-            if (index < 0)
-                throw new InvalidOperationException();
-            return index;
-        }
-
         public bool Next()
         {
             int next = index + 1;

diff --git a/src/neo/SmartContract/Iterators/ByteArrayWrapper.cs b/src/neo/SmartContract/Iterators/ByteArrayWrapper.cs
@@ -20,13 +20,6 @@ public ByteArrayWrapper(VM.Types.Buffer value)
 
         public void Dispose() { }
 
-        public PrimitiveType Key()
-        {
-            if (index < 0)
-                throw new InvalidOperationException();
-            return index;
-        }
-
         public bool Next()
         {
             int next = index + 1;

diff --git a/src/neo/SmartContract/Iterators/IIterator.cs b/src/neo/SmartContract/Iterators/IIterator.cs
@@ -1,10 +1,11 @@
-using Neo.SmartContract.Enumerators;
 using Neo.VM.Types;
+using System;
 
 namespace Neo.SmartContract.Iterators
 {
-    public interface IIterator : IEnumerator
+    public interface IIterator : IDisposable
     {
-        PrimitiveType Key();
+        bool Next();
+        StackItem Value();
     }
 }
diff --git a/src/neo/SmartContract/Iterators/MapWrapper.cs b/src/neo/SmartContract/Iterators/MapWrapper.cs
@@ -1,3 +1,4 @@
+using Neo.VM;
 using Neo.VM.Types;
 using System.Collections.Generic;
 
@@ -6,30 +7,27 @@ namespace Neo.SmartContract.Iterators
     internal class MapWrapper : IIterator
     {
         private readonly IEnumerator<KeyValuePair<PrimitiveType, StackItem>> enumerator;
+        private readonly ReferenceCounter referenceCounter;
 
-        public MapWrapper(IEnumerable<KeyValuePair<PrimitiveType, StackItem>> map)
+        public MapWrapper(IEnumerable<KeyValuePair<PrimitiveType, StackItem>> map, ReferenceCounter referenceCounter)
         {
             this.enumerator = map.GetEnumerator();
+            this.referenceCounter = referenceCounter;
         }
 
         public void Dispose()
         {
             enumerator.Dispose();
         }
 
-        public PrimitiveType Key()
-        {
-            return enumerator.Current.Key;
-        }
-
         public bool Next()
         {
             return enumerator.MoveNext();
         }
 
         public StackItem Value()
         {
-            return enumerator.Current.Value;
+            return new Struct(referenceCounter) { enumerator.Current.Key, enumerator.Current.Value };
         }
     }
 }
diff --git a/src/neo/SmartContract/Iterators/StorageIterator.cs b/src/neo/SmartContract/Iterators/StorageIterator.cs
@@ -1,4 +1,5 @@
 using Neo.Ledger;
+using Neo.VM;
 using Neo.VM.Types;
 using System.Collections.Generic;
 
@@ -7,30 +8,48 @@ namespace Neo.SmartContract.Iterators
     internal class StorageIterator : IIterator
     {
         private readonly IEnumerator<(StorageKey Key, StorageItem Value)> enumerator;
+        private readonly FindOptions options;
+        private readonly ReferenceCounter referenceCounter;
 
-        public StorageIterator(IEnumerator<(StorageKey, StorageItem)> enumerator)
+        public StorageIterator(IEnumerator<(StorageKey, StorageItem)> enumerator, FindOptions options, ReferenceCounter referenceCounter)
         {
             this.enumerator = enumerator;
+            this.options = options;
+            this.referenceCounter = referenceCounter;
         }
 
         public void Dispose()
         {
             enumerator.Dispose();
         }
 
-        public PrimitiveType Key()
-        {
-            return enumerator.Current.Key.Key;
-        }
-
         public bool Next()
         {
             return enumerator.MoveNext();
         }
 
         public StackItem Value()
         {
-            return enumerator.Current.Value.Value;
+            byte[] key = enumerator.Current.Key.Key;
+            byte[] value = enumerator.Current.Value.Value;
+
+            if (options.HasFlag(FindOptions.RemovePrefix))
+                key = key[1..];
+
+            StackItem item = options.HasFlag(FindOptions.DeserializeValues)
+                ? BinarySerializer.Deserialize(value, 1024, (uint)value.Length, referenceCounter)
+                : value;
+
+            if (options.HasFlag(FindOptions.PickField0))
+                item = ((Array)item)[0];
+            else if (options.HasFlag(FindOptions.PickField1))
+                item = ((Array)item)[1];
+
+            if (options.HasFlag(FindOptions.KeysOnly))
+                return key;
+            if (options.HasFlag(FindOptions.ValuesOnly))
+                return item;
+            return new Struct(referenceCounter) { key, item };
         }
     }
 }