diff --git a/src/neo/SmartContract/InteropService.NEO.cs b/src/neo/SmartContract/InteropService.NEO.cs index 10798c053d..a133e882fa 100644 --- a/src/neo/SmartContract/InteropService.NEO.cs +++ b/src/neo/SmartContract/InteropService.NEO.cs @@ -259,13 +259,20 @@ private static bool Storage_Find(ApplicationEngine engine) private static bool Enumerator_Create(ApplicationEngine engine) { - if (engine.CurrentContext.EvaluationStack.Pop() is VMArray array) + IEnumerator enumerator; + switch (engine.CurrentContext.EvaluationStack.Pop()) { - IEnumerator enumerator = new ArrayWrapper(array); - engine.CurrentContext.EvaluationStack.Push(StackItem.FromInterface(enumerator)); - return true; + case VMArray array: + enumerator = new ArrayWrapper(array); + break; + case PrimitiveType primitive: + enumerator = new ByteArrayWrapper(primitive); + break; + default: + return false; } - return false; + engine.CurrentContext.EvaluationStack.Push(StackItem.FromInterface(enumerator)); + return true; } private static bool Enumerator_Next(ApplicationEngine engine) @@ -312,6 +319,9 @@ private static bool Iterator_Create(ApplicationEngine engine) case Map map: iterator = new MapWrapper(map); break; + case PrimitiveType primitive: + iterator = new ByteArrayWrapper(primitive); + break; default: return false; } diff --git a/src/neo/SmartContract/Iterators/ByteArrayWrapper.cs b/src/neo/SmartContract/Iterators/ByteArrayWrapper.cs new file mode 100644 index 0000000000..d5baff7da0 --- /dev/null +++ b/src/neo/SmartContract/Iterators/ByteArrayWrapper.cs @@ -0,0 +1,41 @@ +using Neo.VM.Types; +using System; + +namespace Neo.SmartContract.Iterators +{ + internal class ByteArrayWrapper : IIterator + { + private readonly byte[] array; + private int index = -1; + + public ByteArrayWrapper(PrimitiveType value) + { + this.array = value.ToByteArray().ToArray(); + } + + public void Dispose() { } + + public PrimitiveType Key() + { + if (index < 0) + throw new InvalidOperationException(); + return index; + } + + public bool Next() + { + int next = index + 1; + if (next >= array.Length) + return false; + index = next; + return true; + } + + public StackItem Value() + { + if (index < 0) + throw new InvalidOperationException(); + return new Integer(array[index]); + } + } +} diff --git a/tests/neo.UnitTests/SmartContract/Iterators/UT_PrimitiveWrapper.cs b/tests/neo.UnitTests/SmartContract/Iterators/UT_PrimitiveWrapper.cs new file mode 100644 index 0000000000..d8fd144812 --- /dev/null +++ b/tests/neo.UnitTests/SmartContract/Iterators/UT_PrimitiveWrapper.cs @@ -0,0 +1,49 @@ +using FluentAssertions; +using Microsoft.VisualStudio.TestTools.UnitTesting; +using Neo.SmartContract.Iterators; +using Neo.VM; +using Neo.VM.Types; +using System; + +namespace Neo.UnitTests.SmartContract.Iterators +{ + [TestClass] + public class UT_PrimitiveWrapper + { + [TestMethod] + public void TestGeneratorAndDispose() + { + ByteArrayWrapper arrayWrapper = new ByteArrayWrapper(new ByteArray(new byte[0])); + Assert.IsNotNull(arrayWrapper); + Action action = () => arrayWrapper.Dispose(); + action.Should().NotThrow(); + } + + [TestMethod] + public void TestKeyAndValue() + { + ByteArrayWrapper arrayWrapper = new ByteArrayWrapper(new byte[] { 0x01, 0x02 }); + Action action1 = () => arrayWrapper.Key(); + action1.Should().Throw(); + Action action2 = () => arrayWrapper.Value(); + action2.Should().Throw(); + arrayWrapper.Next(); + Assert.AreEqual(0x00, arrayWrapper.Key().GetBigInteger()); + Assert.AreEqual(0x01, arrayWrapper.Value()); + arrayWrapper.Next(); + Assert.AreEqual(0x01, arrayWrapper.Key().GetBigInteger()); + Assert.AreEqual(0x02, arrayWrapper.Value()); + } + + [TestMethod] + public void TestNext() + { + ByteArrayWrapper arrayWrapper = new ByteArrayWrapper(new byte[] { 0x01, 0x02 }); + Assert.AreEqual(true, arrayWrapper.Next()); + Assert.AreEqual(0x01, arrayWrapper.Value()); + Assert.AreEqual(true, arrayWrapper.Next()); + Assert.AreEqual(0x02, arrayWrapper.Value()); + Assert.AreEqual(false, arrayWrapper.Next()); + } + } +} diff --git a/tests/neo.UnitTests/SmartContract/UT_InteropService.NEO.cs b/tests/neo.UnitTests/SmartContract/UT_InteropService.NEO.cs index 351f885f11..e602cb0f74 100644 --- a/tests/neo.UnitTests/SmartContract/UT_InteropService.NEO.cs +++ b/tests/neo.UnitTests/SmartContract/UT_InteropService.NEO.cs @@ -306,7 +306,7 @@ public void TestEnumerator_Create() .Should().Be(new byte[] { 0x01 }.ToHexString()); engine.CurrentContext.EvaluationStack.Push(1); - InteropService.Invoke(engine, InteropService.Neo_Enumerator_Create).Should().BeFalse(); + InteropService.Invoke(engine, InteropService.Neo_Enumerator_Create).Should().BeTrue(); } [TestMethod] @@ -380,6 +380,10 @@ public void TestIterator_Create() ret.GetInterface().Value().GetSpan().ToHexString() .Should().Be(new byte[] { 0x01 }.ToHexString()); + var interop = new InteropInterface(1); + engine.CurrentContext.EvaluationStack.Push(interop); + InteropService.Invoke(engine, InteropService.Neo_Iterator_Create).Should().BeFalse(); + var map = new Map { { new Integer(1), new Integer(2) }, @@ -393,7 +397,7 @@ public void TestIterator_Create() ret.GetInterface().Value().GetBigInteger().Should().Be(2); engine.CurrentContext.EvaluationStack.Push(1); - InteropService.Invoke(engine, InteropService.Neo_Iterator_Create).Should().BeFalse(); + InteropService.Invoke(engine, InteropService.Neo_Iterator_Create).Should().BeTrue(); } [TestMethod]