From d1d53c3fb632b2aa966c61fd27e94551fc2ff3ce Mon Sep 17 00:00:00 2001
From: Rob Walker <rob@solana.com>
Date: Mon, 10 Jun 2019 16:17:29 -0700
Subject: [PATCH] calculate stake from activated amount (#4630)

---
 programs/stake_api/src/stake_state.rs | 15 ++++++++++
 runtime/src/stakes.rs                 | 42 ++++++++++++++++++---------
 2 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/programs/stake_api/src/stake_state.rs b/programs/stake_api/src/stake_state.rs
index 4e718adf8..f70784e68 100644
--- a/programs/stake_api/src/stake_state.rs
+++ b/programs/stake_api/src/stake_state.rs
@@ -51,12 +51,27 @@ impl StakeState {
         Self::from(account).and_then(|state: Self| state.voter_pubkey())
     }
 
+    // utility function, used by Stakes, tests
+    pub fn voter_pubkey_and_stake_from(account: &Account) -> Option<(Pubkey, u64)> {
+        Self::from(account).and_then(|state: Self| state.voter_pubkey_and_stake())
+    }
+
     pub fn voter_pubkey(&self) -> Option<Pubkey> {
         match self {
             StakeState::Stake { voter_pubkey, .. } => Some(*voter_pubkey),
             _ => None,
         }
     }
+    pub fn voter_pubkey_and_stake(&self) -> Option<(Pubkey, u64)> {
+        match self {
+            StakeState::Stake {
+                voter_pubkey,
+                stake,
+                ..
+            } => Some((*voter_pubkey, *stake)),
+            _ => None,
+        }
+    }
 
     pub fn calculate_rewards(
         credits_observed: u64,
diff --git a/runtime/src/stakes.rs b/runtime/src/stakes.rs
index e7772733f..849609647 100644
--- a/runtime/src/stakes.rs
+++ b/runtime/src/stakes.rs
@@ -16,13 +16,17 @@ pub struct Stakes {
 
 impl Stakes {
     // sum the stakes that point to the given voter_pubkey
-    fn calculate_stake(&self, voter_pubkey: &Pubkey) -> u64 {
+    fn calculate_stake(&self, voter: &Pubkey) -> u64 {
         self.stake_accounts
             .iter()
-            .filter(|(_, stake_account)| {
-                Some(*voter_pubkey) == StakeState::voter_pubkey_from(stake_account)
+            .map(|(_, stake_account)| match StakeState::from(stake_account) {
+                Some(StakeState::Stake {
+                    voter_pubkey,
+                    stake,
+                    ..
+                }) if *voter == voter_pubkey => stake,
+                _ => 0,
             })
-            .map(|(_, stake_account)| stake_account.lamports)
             .sum()
     }
 
@@ -45,22 +49,25 @@ impl Stakes {
             }
         } else if solana_stake_api::check_id(&account.owner) {
             //  old_stake is stake lamports and voter_pubkey from the pre-store() version
-            let old_stake = self.stake_accounts.get(pubkey).and_then(|old_account| {
-                StakeState::voter_pubkey_from(old_account)
-                    .map(|old_voter_pubkey| (old_account.lamports, old_voter_pubkey))
-            });
+            let old_stake = self
+                .stake_accounts
+                .get(pubkey)
+                .and_then(|old_account| StakeState::voter_pubkey_and_stake_from(old_account));
 
-            let stake = StakeState::voter_pubkey_from(account)
-                .map(|voter_pubkey| (account.lamports, voter_pubkey));
+            let stake = if account.lamports != 0 {
+                StakeState::voter_pubkey_and_stake_from(account)
+            } else {
+                StakeState::voter_pubkey_from(account).map(|voter_pubkey| (voter_pubkey, 0))
+            };
 
             // if adjustments need to be made...
             if stake != old_stake {
-                if let Some((old_stake, old_voter_pubkey)) = old_stake {
+                if let Some((old_voter_pubkey, old_stake)) = old_stake {
                     self.vote_accounts
                         .entry(old_voter_pubkey)
                         .and_modify(|e| e.0 -= old_stake);
                 }
-                if let Some((stake, voter_pubkey)) = stake {
+                if let Some((voter_pubkey, stake)) = stake {
                     self.vote_accounts
                         .entry(voter_pubkey)
                         .and_modify(|e| e.0 += stake);
@@ -125,7 +132,16 @@ mod tests {
         {
             let vote_accounts = stakes.vote_accounts();
             assert!(vote_accounts.get(&vote_pubkey).is_some());
-            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 42);
+            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 10); // stays old stake, because only 10 is activated
+        }
+
+        // activate more
+        let (_stake_pubkey, mut stake_account) = create_stake_account(42, &vote_pubkey);
+        stakes.store(&stake_pubkey, &stake_account);
+        {
+            let vote_accounts = stakes.vote_accounts();
+            assert!(vote_accounts.get(&vote_pubkey).is_some());
+            assert_eq!(vote_accounts.get(&vote_pubkey).unwrap().0, 42); // now stake of 42 is activated
         }
 
         stake_account.lamports = 0;