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
@@ -73,7 +98,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);
         }
@@ -106,6 +131,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 @@ internal class InteropDescriptor
         public 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; }
 
         public InteropDescriptor(string method, Func<ApplicationEngine, bool> handler, long price, TriggerType allowedTriggers)
@@ -24,6 +26,13 @@ public InteropDescriptor(string method, Func<ApplicationEngine, bool> handler, F
             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);
         }
     }

src/neo/SmartContract/InteropService.cs

@@ -19,6 +19,8 @@ namespace Neo.SmartContract
     public static partial class InteropService
     {
         public const long GasPerByte = 100000;
+        public const long GasPerReusedByte = -GasPerByte;
+        public const long GasPerReleasedByte = -GasPerByte;
         public const int MaxStorageKeySize = 64;
         public const int MaxStorageValueSize = ushort.MaxValue;
         public const int MaxNotificationSize = 1024;
@@ -52,7 +54,7 @@ public static partial class InteropService
         public static readonly uint System_Storage_Get = Register("System.Storage.Get", Storage_Get, 0_01000000, TriggerType.Application);
         public static readonly uint System_Storage_Put = Register("System.Storage.Put", Storage_Put, GetStoragePrice, TriggerType.Application);
         public static readonly uint System_Storage_PutEx = Register("System.Storage.PutEx", Storage_PutEx, GetStoragePrice, TriggerType.Application);
-        public static readonly uint System_Storage_Delete = Register("System.Storage.Delete", Storage_Delete, 0_01000000, TriggerType.Application);
+        public static readonly uint System_Storage_Delete = Register("System.Storage.Delete", Storage_Delete, GetDeleteStoragePrice, TriggerType.Application);
         public static readonly uint System_StorageContext_AsReadOnly = Register("System.StorageContext.AsReadOnly", StorageContext_AsReadOnly, 0_00000400, TriggerType.Application);
 
         private static bool CheckItemForNotification(StackItem state)
@@ -109,19 +111,73 @@ private static bool CheckStorageContext(ApplicationEngine engine, StorageContext
             return true;
         }
 
+        public static long GetPrice(uint hash)
+        {
+            return methods[hash].GetPrice();
+        }
+
         public static long GetPrice(uint hash, EvaluationStack stack)
         {
             return methods[hash].GetPrice(stack);
         }
 
+        public static long GetPrice(uint hash, ApplicationEngine applicationEngine)
+        {
+            var interopDescriptor = methods[hash];
+            return interopDescriptor.IsStateful ? interopDescriptor.GetPrice(applicationEngine) : interopDescriptor.GetPrice(applicationEngine.CurrentContext.EvaluationStack);
+        }
+
         public static Dictionary<uint, string> SupportedMethods()
         {
             return methods.ToDictionary(p => p.Key, p => p.Value.Method);
         }
 
-        private static long GetStoragePrice(EvaluationStack stack)
+        private static long GetDeleteStoragePrice(ApplicationEngine engine)
         {
-            return (stack.Peek(1).GetByteLength() + stack.Peek(2).GetByteLength()) * GasPerByte;
+            var stack = engine.CurrentContext.EvaluationStack;
+            var key = stack.Peek(1);
+            StorageKey skey = new StorageKey
+            {
+                ScriptHash = engine.CurrentScriptHash,
+                Key = key.GetSpan().ToArray()
+            };
+
+            var skeyValue = engine.Snapshot.Storages.TryGet(skey);
+
+            if (skey == null || !engine.TryAddUpdatedKey(key.GetSpan().ToArray()))
+                return 0_01000000;
+
+            return skeyValue.Value.Length * GasPerReleasedByte;
+        }
+
+        private static long GetStoragePrice(ApplicationEngine engine)
+        {
+            var stack = engine.CurrentContext.EvaluationStack;
+            var key = stack.Peek(1);
+            var newDataSize = stack.Peek(2).GetByteLength();
+            StorageKey skey = new StorageKey
+            {
+                ScriptHash = engine.CurrentScriptHash,
+                Key = key.GetSpan().ToArray()
+            };
+
+            var skeyValue = engine.Snapshot.Storages.TryGet(skey);
+
+            if (skeyValue == null || !engine.TryAddUpdatedKey(skey.Key))
+                return (key.GetByteLength() + newDataSize) * GasPerByte;
+
+            var currentOccupiedBytes = skeyValue.Value.Length;
+
+            if (newDataSize <= currentOccupiedBytes)
+            {
+                var releasedBytes = currentOccupiedBytes - newDataSize;
+                return releasedBytes * GasPerReleasedByte;
+            }
+            else
+            {
+                var reusedBytes = currentOccupiedBytes;
+                return (newDataSize - reusedBytes) * GasPerByte;
+            }
         }
 
         internal static bool Invoke(ApplicationEngine engine, uint method)
@@ -147,6 +203,13 @@ private static uint Register(string method, Func<ApplicationEngine, bool> handle
             return descriptor.Hash;
         }
 
+        private static uint Register(string method, Func<ApplicationEngine, bool> handler, Func<ApplicationEngine, long> priceCalculator, TriggerType allowedTriggers)
+        {
+            InteropDescriptor descriptor = new InteropDescriptor(method, handler, priceCalculator, allowedTriggers);
+            methods.Add(descriptor.Hash, descriptor);
+            return descriptor.Hash;
+        }
+
         private static bool ExecutionEngine_GetScriptContainer(ApplicationEngine engine)
         {
             engine.CurrentContext.EvaluationStack.Push(
@@ -515,7 +578,8 @@ private static bool PutEx(ApplicationEngine engine, StorageContext context, byte
                 Key = key
             };
 
-            if (engine.Snapshot.Storages.TryGet(skey)?.IsConstant == true) return false;
+            var skeyValue = engine.Snapshot.Storages.TryGet(skey);
+            if (skeyValue?.IsConstant == true) return false;
 
             if (value.Length == 0 && !flags.HasFlag(StorageFlags.Constant))
             {
@@ -528,6 +592,7 @@ private static bool PutEx(ApplicationEngine engine, StorageContext context, byte
                 item.Value = value;
                 item.IsConstant = flags.HasFlag(StorageFlags.Constant);
             }
+
             return true;
         }
 
@@ -564,7 +629,9 @@ private static bool Storage_Delete(ApplicationEngine engine)
                     ScriptHash = context.ScriptHash,
                     Key = engine.CurrentContext.EvaluationStack.Pop().GetSpan().ToArray()
                 };
-                if (engine.Snapshot.Storages.TryGet(key)?.IsConstant == true) return false;
+
+                var value = engine.Snapshot.Storages.TryGet(key);
+                if (value?.IsConstant == true) return false;
                 engine.Snapshot.Storages.Delete(key);
                 return true;
             }

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;
 

tests/neo.UnitTests/SmartContract/UT_ContractParameterContext.cs

@@ -3,6 +3,7 @@
 using Neo.Cryptography.ECC;
 using Neo.Network.P2P.Payloads;
 using Neo.SmartContract;
+using Neo.VM;
 using Neo.Wallets;
 using System;