From 7cd7173b71ba80a8698743c6289e6489844376c8 Mon Sep 17 00:00:00 2001
From: Justin Starry <justin@solana.com>
Date: Sat, 25 Jun 2022 17:41:35 +0100
Subject: [PATCH] Refactor: Add get_delegated_stake method to VoteAccounts
 (#26221)

---
 core/src/cluster_info_vote_listener.rs | 9 ++-------
 runtime/src/epoch_stakes.rs            | 4 +---
 runtime/src/vote_account.rs            | 7 +++++++
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/core/src/cluster_info_vote_listener.rs b/core/src/cluster_info_vote_listener.rs
index 9fe7c11086..f069cda1a4 100644
--- a/core/src/cluster_info_vote_listener.rs
+++ b/core/src/cluster_info_vote_listener.rs
@@ -611,10 +611,7 @@ impl ClusterInfoVoteListener {
             // 2) We do not know the hash of the earlier slot
             if slot == last_vote_slot {
                 let vote_accounts = epoch_stakes.stakes().vote_accounts();
-                let stake = vote_accounts
-                    .get(vote_pubkey)
-                    .map(|(stake, _)| *stake)
-                    .unwrap_or_default();
+                let stake = vote_accounts.get_delegated_stake(vote_pubkey);
                 let total_stake = epoch_stakes.total_stake();
 
                 // Fast track processing of the last slot in a vote transactions
@@ -795,9 +792,7 @@ impl ClusterInfoVoteListener {
 
     fn sum_stake(sum: &mut u64, epoch_stakes: Option<&EpochStakes>, pubkey: &Pubkey) {
         if let Some(stakes) = epoch_stakes {
-            if let Some(vote_account) = stakes.stakes().vote_accounts().get(pubkey) {
-                *sum += vote_account.0;
-            }
+            *sum += stakes.stakes().vote_accounts().get_delegated_stake(pubkey)
         }
     }
 }
diff --git a/runtime/src/epoch_stakes.rs b/runtime/src/epoch_stakes.rs
index 11b01165ad..bd878bd06f 100644
--- a/runtime/src/epoch_stakes.rs
+++ b/runtime/src/epoch_stakes.rs
@@ -55,9 +55,7 @@ impl EpochStakes {
     pub fn vote_account_stake(&self, vote_account: &Pubkey) -> u64 {
         self.stakes
             .vote_accounts()
-            .get(vote_account)
-            .map(|(stake, _)| *stake)
-            .unwrap_or(0)
+            .get_delegated_stake(vote_account)
     }
 
     fn parse_epoch_vote_accounts(
diff --git a/runtime/src/vote_account.rs b/runtime/src/vote_account.rs
index cdb09a509b..5214a47273 100644
--- a/runtime/src/vote_account.rs
+++ b/runtime/src/vote_account.rs
@@ -124,6 +124,13 @@ impl VoteAccounts {
         self.vote_accounts.get(pubkey)
     }
 
+    pub fn get_delegated_stake(&self, pubkey: &Pubkey) -> u64 {
+        self.vote_accounts
+            .get(pubkey)
+            .map(|(stake, ..)| *stake)
+            .unwrap_or_default()
+    }
+
     pub(crate) fn iter(&self) -> impl Iterator<Item = (&Pubkey, &(u64, VoteAccount))> {
         self.vote_accounts.iter()
     }