auto-claim gas if it's needed

src/neo/Ledger/TransactionVerificationContext.cs

@@ -31,17 +31,18 @@ public void AddTransaction(Transaction tx)
 
         public bool CheckTransaction(Transaction tx, DataCache snapshot)
         {
+            var oracle = tx.GetAttribute<OracleResponse>();
+            if (oracle != null && oracleResponses.ContainsKey(oracle.Id))
+                return false;
+
             BigInteger balance = NativeContract.GAS.BalanceOf(snapshot, tx.Sender);
             senderFee.TryGetValue(tx.Sender, out var totalSenderFeeFromPool);
 
             BigInteger fee = tx.SystemFee + tx.NetworkFee + totalSenderFeeFromPool;
-            if (balance < fee) return false;
-
-            var oracle = tx.GetAttribute<OracleResponse>();
-            if (oracle != null && oracleResponses.ContainsKey(oracle.Id))
-                return false;
+            if (balance >= fee) return true;
 
-            return true;
+            BigInteger unclaimed = NativeContract.NEO.UnclaimedGas(snapshot, tx.Sender, NativeContract.Ledger.CurrentIndex(snapshot));
+            return balance + unclaimed >= fee;
         }
 
         public void RemoveTransaction(Transaction tx)

src/neo/SmartContract/Native/GasToken.cs

@@ -1,5 +1,6 @@
 using Neo.Cryptography.ECC;
 using Neo.Network.P2P.Payloads;
+using System.Collections.Generic;
 
 namespace Neo.SmartContract.Native
 {
@@ -21,8 +22,10 @@ internal override void Initialize(ApplicationEngine engine)
         internal override void OnPersist(ApplicationEngine engine)
         {
             long totalNetworkFee = 0;
+            HashSet<UInt160> distributed = new HashSet<UInt160>();
             foreach (Transaction tx in engine.PersistingBlock.Transactions)
             {
+                if (distributed.Add(tx.Sender)) NEO.DistributeGas(engine, tx.Sender, false);
                 Burn(engine, tx.Sender, tx.SystemFee + tx.NetworkFee);
                 totalNetworkFee += tx.NetworkFee;
             }

src/neo/SmartContract/Native/NeoToken.cs

@@ -44,7 +44,7 @@ public override BigInteger TotalSupply(DataCache snapshot)
 
         protected override void OnBalanceChanging(ApplicationEngine engine, UInt160 account, NeoAccountState state, BigInteger amount)
         {
-            DistributeGas(engine, account, state);
+            DistributeGas(engine, account, state, true);
             if (amount.IsZero) return;
             if (state.VoteTo is null) return;
             engine.Snapshot.GetAndChange(CreateStorageKey(Prefix_VotersCount)).Add(amount);
@@ -54,14 +54,22 @@ protected override void OnBalanceChanging(ApplicationEngine engine, UInt160 acco
             CheckCandidate(engine.Snapshot, state.VoteTo, candidate);
         }
 
-        private void DistributeGas(ApplicationEngine engine, UInt160 account, NeoAccountState state)
+        internal void DistributeGas(ApplicationEngine engine, UInt160 account, bool callOnPayment)
         {
+            StorageItem storage = engine.Snapshot.TryGet(CreateStorageKey(Prefix_Account).Add(account));
+            if (storage is null) return;
+
             // PersistingBlock is null when running under the debugger
             if (engine.PersistingBlock is null) return;
 
+            DistributeGas(engine, account, storage.GetInteroperable<NeoAccountState>(), callOnPayment);
+        }
+
+        private void DistributeGas(ApplicationEngine engine, UInt160 account, NeoAccountState state, bool callOnPayment)
+        {
             BigInteger gas = CalculateBonus(engine.Snapshot, state.VoteTo, state.Balance, state.BalanceHeight, engine.PersistingBlock.Index);
             state.BalanceHeight = engine.PersistingBlock.Index;
-            GAS.Mint(engine, account, gas, true);
+            GAS.Mint(engine, account, gas, callOnPayment);
         }
 
         private BigInteger CalculateBonus(DataCache snapshot, ECPoint vote, BigInteger value, uint start, uint end)
@@ -270,7 +278,7 @@ private bool Vote(ApplicationEngine engine, UInt160 account, ECPoint voteTo)
                 else
                     item.Add(-state_account.Balance);
             }
-            DistributeGas(engine, account, state_account);
+            DistributeGas(engine, account, state_account, true);
             if (state_account.VoteTo != null)
             {
                 StorageKey key = CreateStorageKey(Prefix_Candidate).Add(state_account.VoteTo);

src/neo/Wallets/Wallet.cs

@@ -130,6 +130,7 @@ public BigDecimal GetAvailable(DataCache snapshot, UInt160 asset_id)
 
         public BigDecimal GetBalance(DataCache snapshot, UInt160 asset_id, params UInt160[] accounts)
         {
+            bool isGas = asset_id == NativeContract.GAS.Hash;
             byte[] script;
             using (ScriptBuilder sb = new ScriptBuilder())
             {
@@ -138,6 +139,12 @@ public BigDecimal GetBalance(DataCache snapshot, UInt160 asset_id, params UInt16
                 {
                     sb.EmitDynamicCall(asset_id, "balanceOf", CallFlags.ReadOnly, account);
                     sb.Emit(OpCode.ADD);
+                    if (isGas)
+                    {
+                        // add unclaimed gas
+                        sb.EmitDynamicCall(NativeContract.NEO.Hash, "unclaimedGas", account, NativeContract.Ledger.CurrentIndex(snapshot));
+                        sb.Emit(OpCode.ADD);
+                    }
                 }
                 sb.EmitDynamicCall(asset_id, "decimals", CallFlags.ReadOnly);
                 script = sb.ToArray();
@@ -257,11 +264,19 @@ public Transaction MakeTransaction(DataCache snapshot, TransferOutput[] outputs,
             {
                 foreach (var (assetId, group, sum) in outputs.GroupBy(p => p.AssetId, (k, g) => (k, g, g.Select(p => p.Value.Value).Sum())))
                 {
+                    var isGas = assetId == NativeContract.GAS.Hash;
                     var balances = new List<(UInt160 Account, BigInteger Value)>();
                     foreach (UInt160 account in accounts)
                         using (ScriptBuilder sb2 = new ScriptBuilder())
                         {
                             sb2.EmitDynamicCall(assetId, "balanceOf", CallFlags.ReadOnly, account);
+                            if (isGas)
+                            {
+                                // add unclaimed gas
+                                sb2.EmitDynamicCall(NativeContract.NEO.Hash, "unclaimedGas", account, NativeContract.Ledger.CurrentIndex(snapshot));
+                                sb2.Emit(OpCode.ADD);
+                            }
+
                             using (ApplicationEngine engine = ApplicationEngine.Run(sb2.ToArray(), snapshot, settings: ProtocolSettings))
                             {
                                 if (engine.State != VMState.HALT)
@@ -302,7 +317,10 @@ public Transaction MakeTransaction(DataCache snapshot, TransferOutput[] outputs,
                 script = sb.ToArray();
             }
             if (balances_gas is null)
-                balances_gas = accounts.Select(p => (Account: p, Value: NativeContract.GAS.BalanceOf(snapshot, p))).Where(p => p.Value.Sign > 0).ToList();
+                balances_gas = accounts.Select(p => (
+                    Account: p,
+                    Value: NativeContract.GAS.BalanceOf(snapshot, p) + NativeContract.NEO.UnclaimedGas(snapshot, p, NativeContract.Ledger.CurrentIndex(snapshot))
+                )).Where(p => p.Value.Sign > 0).ToList();
 
             return MakeTransaction(snapshot, script, cosignerList.Values.ToArray(), Array.Empty<TransactionAttribute>(), balances_gas);
         }
@@ -318,7 +336,10 @@ public Transaction MakeTransaction(DataCache snapshot, byte[] script, UInt160 se
             {
                 accounts = new[] { sender };
             }
-            var balances_gas = accounts.Select(p => (Account: p, Value: NativeContract.GAS.BalanceOf(snapshot, p))).Where(p => p.Value.Sign > 0).ToList();
+            var balances_gas = accounts.Select(p => (
+                Account: p,
+                Value: NativeContract.GAS.BalanceOf(snapshot, p) + NativeContract.NEO.UnclaimedGas(snapshot, p, NativeContract.Ledger.CurrentIndex(snapshot))
+                )).Where(p => p.Value.Sign > 0).ToList();
             return MakeTransaction(snapshot, script, cosigners ?? Array.Empty<Signer>(), attributes ?? Array.Empty<TransactionAttribute>(), balances_gas, maxGas);
         }