Simplify Iterator APIs

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();
  }
  }
 }

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)

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
+ }
+}

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;

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;

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();
  }
 }

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 };
  }
  }
 }

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 };
  }
  }
 }