https://github.com/neo-project/neo/pull/1935

CityOfZion · Jan 18, 2021 · 61e7950 · 61e7950
1 parent 96c9e34
commit 61e7950
Showing 1 changed file with 19 additions and 14 deletions.
diff --git a/neo3/contracts/native/nativecontract.py b/neo3/contracts/native/nativecontract.py
@@ -1119,7 +1119,7 @@ def init(self):
                                        add_engine=True,
                                        safe_method=False)
 
-        self._register_contract_method(self.register_candidate,
+        self._register_contract_method(self.unregister_candidate,
                                        "unregisterCandidate",
                                        5000000,
                                        parameter_types=[cryptography.ECPoint],
@@ -1185,17 +1185,22 @@ def on_balance_changing(self, engine: contracts.ApplicationEngine,
         if amount == vm.BigInteger.zero():
             return
 
-        if not state.vote_to.is_zero():
-            sk_candidate = storage.StorageKey(self.script_hash, self._PREFIX_CANDIDATE + state.vote_to.to_array())
-            si_candidate = engine.snapshot.storages.get(sk_candidate, read_only=False)
-            candidate_state = _CandidateState.from_storage(si_candidate)
-            candidate_state.votes += amount
-            self._candidates_dirty = True
+        if state.vote_to.is_zero():
+            return
 
-            sk_voters_count = storage.StorageKey(self.script_hash, self._PREFIX_VOTERS_COUNT)
-            si_voters_count = engine.snapshot.storages.get(sk_voters_count, read_only=False)
-            new_value = vm.BigInteger(si_voters_count.value) + amount
-            si_voters_count.value = new_value.to_array()
+        sk_voters_count = storage.StorageKey(self.script_hash, self._PREFIX_VOTERS_COUNT)
+        si_voters_count = engine.snapshot.storages.get(sk_voters_count, read_only=False)
+        new_value = vm.BigInteger(si_voters_count.value) + amount
+        si_voters_count.value = new_value.to_array()
+
+        sk_candidate = storage.StorageKey(self.script_hash, self._PREFIX_CANDIDATE + state.vote_to.to_array())
+        si_candidate = engine.snapshot.storages.get(sk_candidate, read_only=False)
+        candidate_state = _CandidateState.from_storage(si_candidate)
+        candidate_state.votes += amount
+        self._candidates_dirty = True
+
+        if not candidate_state.registered and candidate_state.votes == 0:
+            engine.snapshot.storages.delete(sk_candidate)
 
     def on_persist(self, engine: contracts.ApplicationEngine) -> None:
         super(NeoToken, self).on_persist(engine)
@@ -1286,10 +1291,10 @@ def unregister_candidate(self,
             return True
         else:
             state = _CandidateState.from_storage(storage_item)
-            if state.votes == vm.BigInteger.zero():
+            state.registered = False
+            if state.votes == 0:
                 engine.snapshot.storages.delete(storage_key)
-            else:
-                state.registered = False
+
         self._candidates_dirty = True
         return True