Extend ExecutionContext with states

src/neo-vm/ExecutionContext.cs

@@ -1,11 +1,15 @@
-using System.Diagnostics;
+using System;
+using System.Collections.Generic;
+using System.Diagnostics;
 using System.Runtime.CompilerServices;
 
 namespace Neo.VM
 {
     [DebuggerDisplay("RVCount={RVCount}, InstructionPointer={InstructionPointer}")]
-    public class ExecutionContext
+    public sealed class ExecutionContext
     {
+        private readonly Dictionary<Type, object> states = new Dictionary<Type, object>();
+
         /// <summary>
         /// Number of items to be returned
         /// </summary>
@@ -82,6 +86,28 @@ internal ExecutionContext Clone()
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         private Instruction GetInstruction(int ip) => Script.GetInstruction(ip);
 
+        public T GetState<T>()
+        {
+            return (T)states[typeof(T)];
+        }
+
+        public bool TryGetState<T>(out T value)
+        {
+            if (states.TryGetValue(typeof(T), out var val))
+            {
+                value = (T)val;
+                return true;
+            }
+
+            value = default;
+            return false;
+        }
+
+        public void SetState<T>(T state)
+        {
+            states[typeof(T)] = state;
+        }
+
         internal bool MoveNext()
         {
             InstructionPointer += CurrentInstruction.Size;

src/neo-vm/ExecutionEngine.cs

@@ -62,13 +62,6 @@ public class ExecutionEngine : IDisposable
         public ExecutionContext EntryContext => InvocationStack.Count > 0 ? InvocationStack.Peek(InvocationStack.Count - 1) : null;
         public VMState State { get; internal protected set; } = VMState.BREAK;
 
-        #region Events
-
-        public event EventHandler<ExecutionContext> ContextLoaded;
-        public event EventHandler<ExecutionContext> ContextUnloaded;
-
-        #endregion
-
         #region Limits
 
         /// <summary>
@@ -170,6 +163,10 @@ private static int GetItemCount(IEnumerable<StackItem> items)
 
         #endregion
 
+        protected virtual void ContextUnloaded(ExecutionContext context)
+        {
+        }
+
         public virtual void Dispose()
         {
             InvocationStack.Clear();
@@ -298,11 +295,11 @@ private bool ExecuteInstruction()
                                 context_pop.AltStack.CopyTo(CurrentContext.AltStack);
                             }
                             CheckStackSize(false, 0);
-                            ContextUnloaded?.Invoke(this, context_pop);
                             if (InvocationStack.Count == 0)
                             {
                                 State = VMState.HALT;
                             }
+                            ContextUnloaded(context_pop);
                             return true;
                         }
                     case OpCode.SYSCALL:
@@ -1075,12 +1072,11 @@ private bool ExecuteInstruction()
             return true;
         }
 
-        private void LoadContext(ExecutionContext context)
+        protected virtual void LoadContext(ExecutionContext context)
         {
             if (InvocationStack.Count >= MaxInvocationStackSize)
                 throw new InvalidOperationException();
             InvocationStack.Push(context);
-            ContextLoaded?.Invoke(this, context);
         }
 
         public ExecutionContext LoadScript(byte[] script, int rvcount = -1)

tests/neo-vm.Tests/UtExecutionContext.cs

@@ -0,0 +1,21 @@
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Neo.VM;
+
+namespace Neo.Test
+{
+    [TestClass]
+    public class UtExecutionContext
+    {
+        [TestMethod]
+        public void StateTest()
+        {
+            var context = new ExecutionContext(null, null, 0);
+
+            Assert.IsFalse(context.TryGetState<int>(out var i));
+            context.SetState(5);
+            Assert.AreEqual(5, context.GetState<int>());
+            Assert.IsTrue(context.TryGetState(out i));
+            Assert.AreEqual(5, i);
+        }
+    }
+}