Add CreateArray and CreateMap to Helper (neo-project#2145)

ShawnYun · Jan 8, 2021 · 6c95505 · 6c95505
1 parent 8ad5a44
commit 6c95505
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 0 deletions.
diff --git a/src/neo/VM/Helper.cs b/src/neo/VM/Helper.cs
@@ -15,6 +15,30 @@ namespace Neo.VM
 {
     public static class Helper
     {
+        public static ScriptBuilder CreateArray<T>(this ScriptBuilder sb, IReadOnlyList<T> list = null)
+        {
+            if (list is null || list.Count == 0)
+                return sb.Emit(OpCode.NEWARRAY0);
+            for (int i = list.Count - 1; i >= 0; i--)
+                sb.EmitPush(list[i]);
+            sb.EmitPush(list.Count);
+            return sb.Emit(OpCode.PACK);
+        }
+
+        public static ScriptBuilder CreateMap<TKey, TValue>(this ScriptBuilder sb, IReadOnlyDictionary<TKey, TValue> map = null)
+        {
+            sb.Emit(OpCode.NEWMAP);
+            if (map != null)
+                foreach (var p in map)
+                {
+                    sb.Emit(OpCode.DUP);
+                    sb.EmitPush(p.Key);
+                    sb.EmitPush(p.Value);
+                    sb.Emit(OpCode.SETITEM);
+                }
+            return sb;
+        }
+
         public static ScriptBuilder Emit(this ScriptBuilder sb, params OpCode[] ops)
         {
             foreach (OpCode op in ops)

diff --git a/tests/neo.UnitTests/VM/UT_Helper.cs b/tests/neo.UnitTests/VM/UT_Helper.cs
@@ -68,6 +68,48 @@ public void TestEmitAppCall1()
             CollectionAssert.AreEqual(tempArray, sb.ToArray());
         }
 
+        [TestMethod]
+        public void TestEmitArray()
+        {
+            var expected = new BigInteger[] { 1, 2, 3 };
+            var sb = new ScriptBuilder();
+            sb.CreateArray(expected);
+
+            using var engine = ApplicationEngine.Create(TriggerType.Application, null, null);
+            engine.LoadScript(sb.ToArray());
+            Assert.AreEqual(VMState.HALT, engine.Execute());
+
+            CollectionAssert.AreEqual(expected, engine.ResultStack.Pop<VM.Types.Array>().Select(u => u.GetInteger()).ToArray());
+
+            expected = new BigInteger[] { };
+            sb = new ScriptBuilder();
+            sb.CreateArray(expected);
+
+            using var engine2 = ApplicationEngine.Create(TriggerType.Application, null, null);
+            engine2.LoadScript(sb.ToArray());
+            Assert.AreEqual(VMState.HALT, engine2.Execute());
+
+            Assert.AreEqual(0, engine2.ResultStack.Pop<VM.Types.Array>().Count);
+        }
+
+        [TestMethod]
+        public void TestEmitMap()
+        {
+            var expected = new Dictionary<BigInteger, BigInteger>() { { 1, 2 }, { 3, 4 } };
+            var sb = new ScriptBuilder();
+            sb.CreateMap(expected);
+
+            using var engine = ApplicationEngine.Create(TriggerType.Application, null, null);
+            engine.LoadScript(sb.ToArray());
+            Assert.AreEqual(VMState.HALT, engine.Execute());
+
+            var map = engine.ResultStack.Pop<VM.Types.Map>();
+            var dic = map.ToDictionary(u => u.Key, u => u.Value);
+
+            CollectionAssert.AreEqual(expected.Keys, dic.Keys.Select(u => u.GetInteger()).ToArray());
+            CollectionAssert.AreEqual(expected.Values, dic.Values.Select(u => u.GetInteger()).ToArray());
+        }
+
         [TestMethod]
         public void TestEmitAppCall2()
         {