Storage payback

src/neo/IO/Caching/DataCache.cs

@@ -260,7 +260,7 @@ public TValue GetOrAdd(TKey key, Func<TValue> factory)
             }
         }
 
-        public TValue TryGet(TKey key)
+        public virtual TValue TryGet(TKey key)
         {
             lock (dictionary)
             {

src/neo/SmartContract/ApplicationEngine.cs

@@ -1,6 +1,7 @@
 using Neo.Ledger;
 using Neo.Network.P2P.Payloads;
 using Neo.Persistence;
+using Neo.SmartContract.Native;
 using Neo.VM;
 using Neo.VM.Types;
 using System;
@@ -19,6 +20,8 @@ public partial class ApplicationEngine : ExecutionEngine
         private readonly bool testMode;
         private readonly List<NotifyEventArgs> notifications = new List<NotifyEventArgs>();
         private readonly List<IDisposable> disposables = new List<IDisposable>();
+        private readonly List<byte[]> updatedKeys = new List<byte[]>();
+        private long maxConsumedGas = 0;
 
         public TriggerType Trigger { get; }
         public IVerifiable ScriptContainer { get; }
@@ -45,12 +48,34 @@ internal T AddDisposable<T>(T disposable) where T : IDisposable
             return disposable;
         }
 
+        internal bool TryAddUpdatedKey(byte[] key)
+        {
+            bool keyAdded = false;
+            if (!updatedKeys.Contains(key))
+            {
+                updatedKeys.Add(key);
+                keyAdded = true;
+            }
+            return keyAdded;
+        }
+
         private bool AddGas(long gas)
         {
+            if (gas < 0 && GasConsumed > maxConsumedGas)
+                maxConsumedGas = GasConsumed;
             GasConsumed = checked(GasConsumed + gas);
             return testMode || GasConsumed <= gas_amount;
         }
 
+        /// <summary>
+        /// Recalculate the property GasConsumed to use the gas required to run the whole transaction.
+        /// </summary>
+        private void RecalculateConsumedGas()
+        {
+            if (maxConsumedGas > GasConsumed)
+                GasConsumed = maxConsumedGas;
+        }
+
         protected override void LoadContext(ExecutionContext context)
         {
             // Set default execution context state
@@ -70,7 +95,7 @@ public override void Dispose()
 
         protected override bool OnSysCall(uint method)
         {
-            if (!AddGas(InteropService.GetPrice(method, CurrentContext.EvaluationStack)))
+            if (!AddGas(InteropService.GetPrice(method, this)))
                 return false;
             return InteropService.Invoke(this, method);
         }
@@ -103,6 +128,14 @@ private static Block CreateDummyBlock(StoreView snapshot)
             };
         }
 
+        public override VMState Execute()
+        {
+            var resultingState = base.Execute();
+            RecalculateConsumedGas();
+            return resultingState;
+        }
+
+
         public static ApplicationEngine Run(byte[] script, StoreView snapshot,
             IVerifiable container = null, Block persistingBlock = null, bool testMode = false, long extraGAS = default)
         {

src/neo/SmartContract/InteropDescriptor.cs

@@ -10,6 +10,8 @@ public class InteropDescriptor
         internal Func<ApplicationEngine, bool> Handler { get; }
         public long Price { get; }
         public Func<EvaluationStack, long> PriceCalculator { get; }
+        public Func<ApplicationEngine, long> StoragePriceCalculator { get; }
+        public bool IsStateful { get; }
         public TriggerType AllowedTriggers { get; }
 
         internal InteropDescriptor(string method, Func<ApplicationEngine, bool> handler, long price, TriggerType allowedTriggers)
@@ -24,6 +26,13 @@ internal InteropDescriptor(string method, Func<ApplicationEngine, bool> handler,
             this.PriceCalculator = priceCalculator;
         }
 
+        public InteropDescriptor(string method, Func<ApplicationEngine, bool> handler, Func<ApplicationEngine, long> priceCalculator, TriggerType allowedTriggers)
+            : this(method, handler, allowedTriggers)
+        {
+            this.StoragePriceCalculator = priceCalculator;
+            this.IsStateful = true;
+        }
+
         private InteropDescriptor(string method, Func<ApplicationEngine, bool> handler, TriggerType allowedTriggers)
         {
             this.Method = method;
@@ -32,8 +41,24 @@ private InteropDescriptor(string method, Func<ApplicationEngine, bool> handler,
             this.AllowedTriggers = allowedTriggers;
         }
 
+        public long GetPrice()
+        {
+            return Price;
+        }
+
+        public long GetPrice(ApplicationEngine applicationEngine)
+        {
+            return StoragePriceCalculator is null ? Price : StoragePriceCalculator(applicationEngine);
+        }
+
         public long GetPrice(EvaluationStack stack)
         {
+#if DEBUG
+            if (IsStateful)
+            {
+                throw new InvalidOperationException();
+            }
+#endif
             return PriceCalculator is null ? Price : PriceCalculator(stack);
         }
 

tests/neo.UnitTests/SmartContract/UT_ApplicationEngine.cs

@@ -7,6 +7,8 @@
 using Neo.Network.P2P.Payloads;
 using Neo.Persistence;
 using Neo.SmartContract;
+using Neo.SmartContract.Native;
+using Neo.VM;
 using Neo.VM.Types;
 using System;