Fix candidate (#1935)

neo-project · Sep 15, 2020 · 1ec60fb · 1ec60fb
1 parent 9a160c4
commit 1ec60fb
Showing 1 changed file with 15 additions and 11 deletions.
diff --git a/src/neo/SmartContract/Native/Tokens/NeoToken.cs b/src/neo/SmartContract/Native/Tokens/NeoToken.cs
@@ -47,11 +47,12 @@ protected override void OnBalanceChanging(ApplicationEngine engine, UInt160 acco
         {
             DistributeGas(engine, account, state);
             if (amount.IsZero) return;
-            if (state.VoteTo != null)
-            {
-                engine.Snapshot.Storages.GetAndChange(CreateStorageKey(Prefix_Candidate).Add(state.VoteTo)).GetInteroperable<CandidateState>().Votes += amount;
-                engine.Snapshot.Storages.GetAndChange(CreateStorageKey(Prefix_VotersCount)).Add(amount);
-            }
+            if (state.VoteTo is null) return;
+            engine.Snapshot.Storages.GetAndChange(CreateStorageKey(Prefix_VotersCount)).Add(amount);
+            StorageKey key = CreateStorageKey(Prefix_Candidate).Add(state.VoteTo);
+            CandidateState candidate = engine.Snapshot.Storages.GetAndChange(key).GetInteroperable<CandidateState>();
+            candidate.Votes += amount;
+            CheckCandidate(engine.Snapshot, key, candidate);
         }
 
         private void DistributeGas(ApplicationEngine engine, UInt160 account, NeoAccountState state)
@@ -89,6 +90,12 @@ private BigInteger CalculateBonus(StoreView snapshot, BigInteger value, uint sta
             return value * sum * NeoHolderRewardRatio / 100 / TotalAmount;
         }
 
+        private static void CheckCandidate(StoreView snapshot, StorageKey key, CandidateState candidate)
+        {
+            if (!candidate.Registered && candidate.Votes.IsZero)
+                snapshot.Storages.Delete(key);
+        }
+
         private bool ShouldRefreshCommittee(uint height) => height % (ProtocolSettings.Default.CommitteeMembersCount + ProtocolSettings.Default.ValidatorsCount) == 0;
 
         internal override void Initialize(ApplicationEngine engine)
@@ -189,10 +196,8 @@ private bool UnregisterCandidate(ApplicationEngine engine, ECPoint pubkey)
             if (engine.Snapshot.Storages.TryGet(key) is null) return true;
             StorageItem item = engine.Snapshot.Storages.GetAndChange(key);
             CandidateState state = item.GetInteroperable<CandidateState>();
-            if (state.Votes.IsZero)
-                engine.Snapshot.Storages.Delete(key);
-            else
-                state.Registered = false;
+            state.Registered = false;
+            CheckCandidate(engine.Snapshot, key, state);
             return true;
         }
 
@@ -223,8 +228,7 @@ private bool Vote(ApplicationEngine engine, UInt160 account, ECPoint voteTo)
                 StorageItem storage_validator = engine.Snapshot.Storages.GetAndChange(key);
                 CandidateState state_validator = storage_validator.GetInteroperable<CandidateState>();
                 state_validator.Votes -= state_account.Balance;
-                if (!state_validator.Registered && state_validator.Votes.IsZero)
-                    engine.Snapshot.Storages.Delete(key);
+                CheckCandidate(engine.Snapshot, key, state_validator);
             }
             state_account.VoteTo = voteTo;
             if (validator_new != null)