Fix: wrong reference count after invalid Map[key] = value

src/Neo.VM/Types/Array.cs

@@ -83,10 +83,11 @@ public Array(IReferenceCounter? referenceCounter, IEnumerable<StackItem>? items
             foreach (var item in _array)
             {
                 if (item is CompoundType { ReferenceCounter: null })
-                {
                     throw new InvalidOperationException("Can not set a CompoundType without a ReferenceCounter.");
-                }
+            }
 
+            foreach (var item in _array)
+            {
                 referenceCounter.AddReference(item, this);
             }
         }

src/Neo.VM/Types/Map.cs

@@ -51,14 +51,14 @@ public StackItem this[PrimitiveType key]
                 if (IsReadOnly) throw new InvalidOperationException("The map is readonly, can not set value.");
                 if (ReferenceCounter != null)
                 {
+                    if (value is CompoundType { ReferenceCounter: null })
+                        throw new InvalidOperationException("Can not set value to Map(with ReferenceCounter) without ReferenceCounter.");
+
                     if (dictionary.TryGetValue(key, out StackItem? old_value))
                         ReferenceCounter.RemoveReference(old_value, this);
                     else
                         ReferenceCounter.AddReference(key, this);
-                    if (value is CompoundType { ReferenceCounter: null })
-                    {
-                        throw new InvalidOperationException("Can not set a Map without a ReferenceCounter.");
-                    }
+
                     ReferenceCounter.AddReference(value, this);
                 }
                 dictionary[key] = value;

tests/Neo.VM.Tests/UT_ReferenceCounter.cs

@@ -257,5 +257,36 @@ public void TestInvalidReferenceStackItem()
 
             Assert.ThrowsExactly<InvalidOperationException>(() => arr.Add(arr2));
         }
+
+        [TestMethod]
+        public void TestInvalidMapSet()
+        {
+            {
+                var reference = new ReferenceCounter();
+                var map = new Map(reference);
+                Assert.AreEqual(0, reference.Count); // map is zero-referred
+
+                map["test"] = new Array(reference);
+                Assert.AreEqual(2, reference.Count); // one key + one value
+
+                Assert.ThrowsExactly<InvalidOperationException>(() => map["test"] = new Array());
+                Assert.AreEqual(2, reference.Count); // not changed after invalid set
+
+                map["test"] = new Integer(1); // overwrite the value
+                Assert.AreEqual(2, reference.Count); // not changed after overwrite
+            }
+
+
+            {
+                var reference = new ReferenceCounter();
+                var map = new Map();
+                Assert.ThrowsExactly<InvalidOperationException>(() => new Array(reference, [map]));
+
+                map = new Map(reference);
+                var array = new Array(reference, [map]);
+                Assert.AreEqual(1, reference.Count);
+                Assert.AreEqual(1, array.Count);
+            }
+        }
     }
 }