Use less memory when serializing versioned epoch stakes

runtime/src/bank/serde_snapshot.rs

@@ -6,8 +6,7 @@ mod tests {
                 epoch_accounts_hash_utils, test_utils as bank_test_utils, Bank, EpochRewardStatus,
             },
             epoch_stakes::{
-                EpochAuthorizedVoters, EpochStakes, NodeIdToVoteAccounts, StakesSerdeWrapper,
-                VersionedEpochStakes,
+                EpochAuthorizedVoters, EpochStakes, NodeIdToVoteAccounts, VersionedEpochStakes,
             },
             genesis_utils::activate_all_features,
             runtime_config::RuntimeConfig,
@@ -20,7 +19,7 @@ mod tests {
                 create_tmp_accounts_dir_for_tests, get_storages_to_serialize, ArchiveFormat,
                 StorageAndNextAccountsFileId,
             },
-            stakes::{Stakes, StakesEnum},
+            stakes::{SerdeStakesToStakeFormat, Stakes, StakesEnum},
         },
         solana_accounts_db::{
             account_storage::{AccountStorageMap, AccountStorageReference},
@@ -307,7 +306,7 @@ mod tests {
         bank.epoch_stakes.insert(
             42,
             EpochStakes::from(VersionedEpochStakes::Current {
-                stakes: StakesSerdeWrapper::Stake(Stakes::<Stake>::default()),
+                stakes: SerdeStakesToStakeFormat::Stake(Stakes::<Stake>::default()),
                 total_stake: 42,
                 node_id_to_vote_accounts: Arc::<NodeIdToVoteAccounts>::default(),
                 epoch_authorized_voters: Arc::<EpochAuthorizedVoters>::default(),
@@ -536,7 +535,7 @@ mod tests {
         #[cfg_attr(
             feature = "frozen-abi",
             derive(AbiExample),
-            frozen_abi(digest = "HRBDXrGrHMZU4cNebKHT7jEmhrgd3h1c2qUMMywrGPiq")
+            frozen_abi(digest = "J7MnnLU99fYk2hfZPjdqyTYxgHstwRUDk2Yr8fFnXxFp")
         )]
         #[derive(Serialize)]
         pub struct BankAbiTestWrapper {

runtime/src/epoch_stakes.rs

@@ -1,11 +1,7 @@
 use {
-    crate::{
-        stake_account::StakeAccount,
-        stakes::{serde_stakes_to_delegation_format, Stakes, StakesEnum},
-    },
-    serde::{Deserialize, Deserializer, Serialize, Serializer},
-    solana_sdk::{clock::Epoch, pubkey::Pubkey, stake::state::Stake},
-    solana_stake_program::stake_state::Delegation,
+    crate::stakes::{serde_stakes_to_delegation_format, SerdeStakesToStakeFormat, StakesEnum},
+    serde::{Deserialize, Serialize},
+    solana_sdk::{clock::Epoch, pubkey::Pubkey},
     solana_vote::vote_account::VoteAccountsHashMap,
     std::{collections::HashMap, sync::Arc},
 };
@@ -139,74 +135,13 @@ impl EpochStakes {
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum VersionedEpochStakes {
     Current {
-        stakes: StakesSerdeWrapper,
+        stakes: SerdeStakesToStakeFormat,
         total_stake: u64,
         node_id_to_vote_accounts: Arc<NodeIdToVoteAccounts>,
         epoch_authorized_voters: Arc<EpochAuthorizedVoters>,
     },
 }
 
-/// Wrapper struct with custom serialization to support serializing
-/// `Stakes<StakeAccount>` as `Stakes<Stake>` without doing a full deep clone of
-/// the stake data. Serialization works by building a `Stakes<&Stake>` map which
-/// borrows `&Stake` from `StakeAccount` entries in `Stakes<StakeAccount>`. Note
-/// that `Stakes<&Stake>` still copies `Pubkey` keys so the `Stakes<&Stake>`
-/// data structure still allocates a fair amount of memory but the memory only
-/// remains allocated during serialization.
-#[cfg_attr(feature = "frozen-abi", derive(AbiExample, AbiEnumVisitor))]
-#[derive(Debug, Clone)]
-pub enum StakesSerdeWrapper {
-    Stake(Stakes<Stake>),
-    Account(Stakes<StakeAccount<Delegation>>),
-}
-
-#[cfg(feature = "dev-context-only-utils")]
-impl PartialEq<Self> for StakesSerdeWrapper {
-    fn eq(&self, other: &Self) -> bool {
-        match (self, other) {
-            (Self::Stake(stakes), Self::Stake(other)) => stakes == other,
-            (Self::Account(stakes), Self::Account(other)) => stakes == other,
-            (Self::Stake(stakes), Self::Account(other)) => {
-                stakes == &Stakes::<Stake>::from(other.clone())
-            }
-            (Self::Account(stakes), Self::Stake(other)) => {
-                other == &Stakes::<Stake>::from(stakes.clone())
-            }
-        }
-    }
-}
-
-impl From<StakesSerdeWrapper> for StakesEnum {
-    fn from(stakes: StakesSerdeWrapper) -> Self {
-        match stakes {
-            StakesSerdeWrapper::Stake(stakes) => Self::Stakes(stakes),
-            StakesSerdeWrapper::Account(stakes) => Self::Accounts(stakes),
-        }
-    }
-}
-
-impl Serialize for StakesSerdeWrapper {
-    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-    where
-        S: Serializer,
-    {
-        match self {
-            Self::Stake(stakes) => stakes.serialize(serializer),
-            Self::Account(stakes) => Stakes::<&Stake>::from(stakes).serialize(serializer),
-        }
-    }
-}
-
-impl<'de> Deserialize<'de> for StakesSerdeWrapper {
-    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
-    where
-        D: Deserializer<'de>,
-    {
-        let stakes = Stakes::<Stake>::deserialize(deserializer)?;
-        Ok(Self::Stake(stakes))
-    }
-}
-
 impl From<VersionedEpochStakes> for EpochStakes {
     fn from(versioned: VersionedEpochStakes) -> Self {
         let VersionedEpochStakes::Current {
@@ -262,7 +197,7 @@ pub(crate) fn split_epoch_stakes(
                 versioned_epoch_stakes.insert(
                     epoch,
                     VersionedEpochStakes::Current {
-                        stakes: StakesSerdeWrapper::Account(stakes.clone()),
+                        stakes: SerdeStakesToStakeFormat::Account(stakes.clone()),
                         total_stake,
                         node_id_to_vote_accounts,
                         epoch_authorized_voters,
@@ -273,7 +208,7 @@ pub(crate) fn split_epoch_stakes(
                 versioned_epoch_stakes.insert(
                     epoch,
                     VersionedEpochStakes::Current {
-                        stakes: StakesSerdeWrapper::Stake(stakes.clone()),
+                        stakes: SerdeStakesToStakeFormat::Stake(stakes.clone()),
                         total_stake,
                         node_id_to_vote_accounts,
                         epoch_authorized_voters,
@@ -289,10 +224,13 @@ pub(crate) fn split_epoch_stakes(
 pub(crate) mod tests {
     use {
         super::*,
-        crate::{stake_account::StakeAccount, stakes::StakesCache},
+        crate::{
+            stake_account::StakeAccount,
+            stakes::{Stakes, StakesCache},
+        },
         im::HashMap as ImHashMap,
         solana_sdk::{account::AccountSharedData, rent::Rent},
-        solana_stake_program::stake_state::{self, Delegation},
+        solana_stake_program::stake_state::{self, Delegation, Stake},
         solana_vote::vote_account::{VoteAccount, VoteAccounts},
         solana_vote_program::vote_state::{self, create_account_with_authorized},
         std::iter,
@@ -492,7 +430,7 @@ pub(crate) mod tests {
         assert_eq!(
             versioned.get(&epoch),
             Some(&VersionedEpochStakes::Current {
-                stakes: StakesSerdeWrapper::Account(test_stakes),
+                stakes: SerdeStakesToStakeFormat::Account(test_stakes),
                 total_stake: epoch_stakes.total_stake,
                 node_id_to_vote_accounts: epoch_stakes.node_id_to_vote_accounts,
                 epoch_authorized_voters: epoch_stakes.epoch_authorized_voters,
@@ -521,7 +459,7 @@ pub(crate) mod tests {
         assert_eq!(
             versioned.get(&epoch),
             Some(&VersionedEpochStakes::Current {
-                stakes: StakesSerdeWrapper::Stake(test_stakes),
+                stakes: SerdeStakesToStakeFormat::Stake(test_stakes),
                 total_stake: epoch_stakes.total_stake,
                 node_id_to_vote_accounts: epoch_stakes.node_id_to_vote_accounts,
                 epoch_authorized_voters: epoch_stakes.epoch_authorized_voters,
@@ -575,7 +513,7 @@ pub(crate) mod tests {
         assert_eq!(
             versioned.get(&epoch2),
             Some(&VersionedEpochStakes::Current {
-                stakes: StakesSerdeWrapper::Account(Stakes::default()),
+                stakes: SerdeStakesToStakeFormat::Account(Stakes::default()),
                 total_stake: 200,
                 node_id_to_vote_accounts: Arc::default(),
                 epoch_authorized_voters: Arc::default(),
@@ -584,7 +522,7 @@ pub(crate) mod tests {
         assert_eq!(
             versioned.get(&epoch3),
             Some(&VersionedEpochStakes::Current {
-                stakes: StakesSerdeWrapper::Stake(Stakes::default()),
+                stakes: SerdeStakesToStakeFormat::Stake(Stakes::default()),
                 total_stake: 300,
                 node_id_to_vote_accounts: Arc::default(),
                 epoch_authorized_voters: Arc::default(),

runtime/src/stakes.rs

@@ -28,6 +28,7 @@ use {
 
 mod serde_stakes;
 pub(crate) use serde_stakes::serde_stakes_to_delegation_format;
+pub use serde_stakes::SerdeStakesToStakeFormat;
 
 #[derive(Debug, Error)]
 pub enum Error {
@@ -570,23 +571,6 @@ impl From<Stakes<StakeAccount>> for Stakes<Stake> {
     }
 }
 
-impl<'a> From<&'a Stakes<StakeAccount>> for Stakes<&'a Stake> {
-    fn from(stakes: &'a Stakes<StakeAccount>) -> Self {
-        let stake_delegations = stakes
-            .stake_delegations
-            .iter()
-            .map(|(pubkey, stake_account)| (*pubkey, stake_account.stake()))
-            .collect();
-        Self {
-            vote_accounts: stakes.vote_accounts.clone(),
-            stake_delegations,
-            unused: stakes.unused,
-            epoch: stakes.epoch,
-            stake_history: stakes.stake_history.clone(),
-        }
-    }
-}
-
 /// This conversion is memory intensive so should only be used in development
 /// contexts.
 #[cfg(feature = "dev-context-only-utils")]