Remove empty stake entries from state

Author: foriequal0

core/src/consensus/tendermint/stake/action_data.rs

@@ -26,7 +26,7 @@ use rlp::{Decodable, Encodable, Rlp, RlpStream};
 
 use super::CUSTOM_ACTION_HANDLER_ID;
 
-fn get_account_key(address: &Address) -> H256 {
+pub fn get_account_key(address: &Address) -> H256 {
     ActionDataKeyBuilder::new(CUSTOM_ACTION_HANDLER_ID, 2).append(&"Account").append(address).into_key()
 }
 
@@ -35,7 +35,7 @@ lazy_static! {
         ActionDataKeyBuilder::new(CUSTOM_ACTION_HANDLER_ID, 1).append(&"StakeholderAddresses").into_key();
 }
 
-fn get_delegation_key(address: &Address) -> H256 {
+pub fn get_delegation_key(address: &Address) -> H256 {
     ActionDataKeyBuilder::new(CUSTOM_ACTION_HANDLER_ID, 2).append(&"Delegation").append(address).into_key()
 }
 
@@ -64,8 +64,12 @@ impl<'a> StakeAccount<'a> {
 
     pub fn save_to_state(&self, state: &mut TopLevelState) -> StateResult<()> {
         let account_key = get_account_key(self.address);
-        let rlp = rlp::encode(&self.balance);
-        state.update_action_data(&account_key, rlp.into_vec())?;
+        if self.balance != 0 {
+            let rlp = rlp::encode(&self.balance);
+            state.update_action_data(&account_key, rlp.into_vec())?;
+        } else {
+            state.remove_action_data(&account_key);
+        }
         Ok(())
     }
 
@@ -97,7 +101,12 @@ impl Stakeholders {
     }
 
     pub fn save_to_state(&self, state: &mut TopLevelState) -> StateResult<()> {
-        state.update_action_data(&*STAKEHOLDER_ADDRESSES_KEY, encode_set(&self.0))?;
+        let key = *STAKEHOLDER_ADDRESSES_KEY;
+        if !self.0.is_empty() {
+            state.update_action_data(&key, encode_set(&self.0))?;
+        } else {
+            state.remove_action_data(&key);
+        }
         Ok(())
     }
 
@@ -106,10 +115,15 @@ impl Stakeholders {
         self.0.contains(address)
     }
 
-    pub fn update(&mut self, account: &StakeAccount) {
+    pub fn update_by_increased_balance(&mut self, account: &StakeAccount) {
         if account.balance > 0 {
             self.0.insert(*account.address);
-        } else {
+        }
+    }
+
+    pub fn update_by_decreased_balance(&mut self, account: &StakeAccount, delegation: &Delegation) {
+        assert!(account.address == delegation.delegator);
+        if account.balance == 0 && delegation.sum() == 0 {
             self.0.remove(account.address);
         }
     }
@@ -138,8 +152,12 @@ impl<'a> Delegation<'a> {
 
     pub fn save_to_state(&self, state: &mut TopLevelState) -> StateResult<()> {
         let key = get_delegation_key(self.delegator);
-        let encoded = encode_map(&self.delegatees);
-        state.update_action_data(&key, encoded)?;
+        if !self.delegatees.is_empty() {
+            let encoded = encode_map(&self.delegatees);
+            state.update_action_data(&key, encoded)?;
+        } else {
+            state.remove_action_data(&key);
+        }
         Ok(())
     }
 
@@ -296,6 +314,24 @@ mod tests {
         assert_eq!(account.balance, 10);
     }
 
+    #[test]
+    fn balance_subtract_all_should_remove_entry_from_db() {
+        let mut state = helpers::get_temp_state();
+        let address = Address::random();
+
+        let mut account = StakeAccount::load_from_state(&state, &address).unwrap();
+        account.add_balance(100).unwrap();
+        account.save_to_state(&mut state).unwrap();
+
+        let mut account = StakeAccount::load_from_state(&state, &address).unwrap();
+        let result = account.subtract_balance(100);
+        assert!(result.is_ok());
+        account.save_to_state(&mut state).unwrap();
+
+        let data = state.action_data(&get_account_key(&address)).unwrap();
+        assert_eq!(data, None);
+    }
+
     #[test]
     fn stakeholders_track() {
         let mut rng = rng();
@@ -311,7 +347,7 @@ mod tests {
 
         let mut stakeholders = Stakeholders::load_from_state(&state).unwrap();
         for account in &accounts {
-            stakeholders.update(account);
+            stakeholders.update_by_increased_balance(account);
         }
         stakeholders.save_to_state(&mut state).unwrap();
 
@@ -334,18 +370,17 @@ mod tests {
 
         let mut stakeholders = Stakeholders::load_from_state(&state).unwrap();
         for account in &accounts {
-            stakeholders.update(account);
+            stakeholders.update_by_increased_balance(account);
         }
         stakeholders.save_to_state(&mut state).unwrap();
 
+        let mut stakeholders = Stakeholders::load_from_state(&state).unwrap();
         for account in &mut accounts {
             if rand::random() {
                 account.balance = 0;
             }
-        }
-        let mut stakeholders = Stakeholders::load_from_state(&state).unwrap();
-        for account in &accounts {
-            stakeholders.update(account);
+            let delegation = Delegation::load_from_state(&state, account.address).unwrap();
+            stakeholders.update_by_decreased_balance(account, &delegation);
         }
         stakeholders.save_to_state(&mut state).unwrap();
 
@@ -357,6 +392,40 @@ mod tests {
         }
     }
 
+    #[test]
+    fn stakeholders_doesnt_untrack_if_delegation_exists() {
+        let mut state = helpers::get_temp_state();
+        let addresses: Vec<_> = (1..100).map(|_| Address::random()).collect();
+        let mut accounts: Vec<_> = addresses
+            .iter()
+            .map(|address| StakeAccount {
+                address,
+                balance: 100,
+            })
+            .collect();
+
+        let mut stakeholders = Stakeholders::load_from_state(&state).unwrap();
+        for account in &accounts {
+            stakeholders.update_by_increased_balance(account);
+        }
+        stakeholders.save_to_state(&mut state).unwrap();
+
+        let mut stakeholders = Stakeholders::load_from_state(&state).unwrap();
+        for account in &mut accounts {
+            // like self-delegate
+            let mut delegation = Delegation::load_from_state(&state, account.address).unwrap();
+            delegation.add_quantity(*account.address, account.balance).unwrap();
+            account.balance = 0;
+            stakeholders.update_by_decreased_balance(account, &delegation);
+        }
+        stakeholders.save_to_state(&mut state).unwrap();
+
+        let stakeholders = Stakeholders::load_from_state(&state).unwrap();
+        for account in &accounts {
+            assert!(stakeholders.contains(account.address));
+        }
+    }
+
     #[test]
     fn initial_delegation_is_empty() {
         let state = helpers::get_temp_state();

core/src/consensus/tendermint/stake/mod.rs

@@ -57,16 +57,14 @@ impl ActionHandler for Stake {
     fn init(&self, state: &mut TopLevelState) -> StateResult<()> {
         let mut stakeholders = Stakeholders::load_from_state(state)?;
         for (address, amount) in self.genesis_stakes.iter() {
-            if *amount > 0 {
-                let account = StakeAccount {
-                    address,
-                    balance: *amount,
-                };
-                stakeholders.update(&account);
-                account.save_to_state(state)?;
-            }
-            stakeholders.save_to_state(state)?;
+            let account = StakeAccount {
+                address,
+                balance: *amount,
+            };
+            account.save_to_state(state)?;
+            stakeholders.update_by_increased_balance(&account);
         }
+        stakeholders.save_to_state(state)?;
         Ok(())
     }
 
@@ -90,12 +88,13 @@ fn transfer_ccs(state: &mut TopLevelState, sender: &Address, receiver: &Address,
     let mut stakeholders = Stakeholders::load_from_state(state)?;
     let mut sender_account = StakeAccount::load_from_state(state, sender)?;
     let mut receiver_account = StakeAccount::load_from_state(state, receiver)?;
+    let sender_delegations = Delegation::load_from_state(state, sender)?;
 
     sender_account.subtract_balance(quantity)?;
     receiver_account.add_balance(quantity)?;
 
-    stakeholders.update(&sender_account);
-    stakeholders.update(&receiver_account);
+    stakeholders.update_by_decreased_balance(&sender_account, &sender_delegations);
+    stakeholders.update_by_increased_balance(&receiver_account);
 
     stakeholders.save_to_state(state)?;
     sender_account.save_to_state(state)?;
@@ -120,6 +119,7 @@ fn delegate_ccs(
 
     delegator.subtract_balance(quantity)?;
     delegation.add_quantity(*delegatee, quantity)?;
+    // delegation does not touch stakeholders
 
     delegation.save_to_state(state)?;
     delegator.save_to_state(state)?;
@@ -139,10 +139,13 @@ pub fn get_stakes(state: &TopLevelState) -> StateResult<HashMap<Address, u64>> {
 
 #[cfg(test)]
 mod tests {
+    use super::action_data::get_account_key;
     use super::*;
+
     use ckey::{public_to_address, Public};
     use consensus::validator_set::new_validator_set;
     use cstate::tests::helpers;
+    use cstate::TopStateView;
     use rlp::Encodable;
 
     #[test]
@@ -210,6 +213,7 @@ mod tests {
         let account1 = StakeAccount::load_from_state(&state, &address1).unwrap();
         let account2 = StakeAccount::load_from_state(&state, &address2).unwrap();
         assert_eq!(account1.balance, 0);
+        assert_eq!(state.action_data(&get_account_key(&address1)).unwrap(), None);
         assert_eq!(account2.balance, 100);
         let stakeholders = Stakeholders::load_from_state(&state).unwrap();
         assert!(!stakeholders.contains(&address1));
@@ -251,6 +255,42 @@ mod tests {
         assert_eq!(delegation_untouched.iter().count(), 0);
     }
 
+    #[test]
+    fn delegate_all() {
+        let delegatee_public = Public::random();
+        let delegatee = public_to_address(&delegatee_public);
+        let delegator = Address::random();
+
+        let mut state = helpers::get_temp_state();
+        let stake = {
+            let mut genesis_stakes = HashMap::new();
+            genesis_stakes.insert(delegatee, 100);
+            genesis_stakes.insert(delegator, 100);
+            Stake::new(genesis_stakes, new_validator_set(vec![delegatee_public]))
+        };
+        assert_eq!(Ok(()), stake.init(&mut state));
+
+        let action = Action::DelegateCCS {
+            address: delegatee,
+            quantity: 100,
+        };
+        let result = stake.execute(&action.rlp_bytes(), &mut state, &delegator);
+        assert_eq!(result, Ok(()));
+
+        let delegator_account = StakeAccount::load_from_state(&state, &delegator).unwrap();
+        let delegation = Delegation::load_from_state(&state, &delegator).unwrap();
+        assert_eq!(delegator_account.balance, 0);
+        assert_eq!(state.action_data(&get_account_key(&delegator)).unwrap(), None);
+        assert_eq!(delegation.iter().count(), 1);
+        assert_eq!(delegation.get_quantity(&delegatee), 100);
+
+        // Should not be touched
+        let delegatee_account = StakeAccount::load_from_state(&state, &delegatee).unwrap();
+        let delegation_untouched = Delegation::load_from_state(&state, &delegatee).unwrap();
+        assert_eq!(delegatee_account.balance, 100);
+        assert_eq!(delegation_untouched.iter().count(), 0);
+    }
+
     #[test]
     fn delegate_only_to_validator() {
         let delegatee_public = Public::random();

state/src/cache/top_cache.rs

@@ -155,7 +155,6 @@ impl TopCache {
         self.action_data.get_mut(a, db)
     }
 
-    #[allow(dead_code)]
     pub fn remove_action_data(&self, address: &H256) {
         self.action_data.remove(address)
     }

state/src/impls/top_level.rs

@@ -967,6 +967,10 @@ impl TopState for TopLevelState {
         *action_data = data.into();
         Ok(())
     }
+
+    fn remove_action_data(&mut self, key: &H256) {
+        self.top_cache.remove_action_data(key)
+    }
 }
 
 fn is_active_account(state: &TopStateView, address: &Address) -> TrieResult<bool> {

state/src/traits.rs

@@ -178,6 +178,7 @@ pub trait TopState {
     fn remove_text(&mut self, key: &H256, sig: &Signature) -> StateResult<()>;
 
     fn update_action_data(&mut self, key: &H256, data: Bytes) -> StateResult<()>;
+    fn remove_action_data(&mut self, key: &H256);
 }
 
 pub trait StateWithCache {