Skip to main content

solana_runtime/
epoch_stakes.rs

1#[cfg(feature = "dev-context-only-utils")]
2use qualifier_attr::qualifiers;
3use {
4    crate::{stake_history::StakeHistory, stakes::SerdeStakesToStakeFormat},
5    serde::{
6        Deserialize, Deserializer, Serialize, Serializer,
7        de::{SeqAccess, Visitor},
8    },
9    solana_bls_signatures::pubkey::{
10        PopVerified, PubkeyAffine as BLSPubkeyAffine, PubkeyCompressed as BLSPubkeyCompressed,
11    },
12    solana_clock::Epoch,
13    solana_pubkey::Pubkey,
14    solana_stake_interface::state::Stake,
15    solana_vote::vote_account::{VoteAccounts, VoteAccountsHashMap},
16    solana_vote_interface::state::BLS_PUBLIC_KEY_COMPRESSED_SIZE,
17    std::{
18        collections::HashMap,
19        fmt,
20        mem::MaybeUninit,
21        num::NonZero,
22        sync::{Arc, OnceLock},
23    },
24    wincode::{ReadResult, SchemaRead, SchemaWrite, WriteResult, config::Config, io::Reader},
25};
26
27pub type NodeIdToVoteAccounts = HashMap<Pubkey, NodeVoteAccounts>;
28pub type EpochAuthorizedVoters = HashMap<Pubkey, Pubkey>;
29
30/// Entry in the [`BLSPubkeyToRankMap`] associating a validator's identity
31/// pubkey and BLS pubkey with its stake.
32#[derive(Clone, Debug)]
33#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
34pub struct BLSPubkeyStakeEntry {
35    /// The address containing the vote account
36    pub vote_account_pubkey: Pubkey,
37    /// The identity of the validator specified in the vote account
38    pub node_pubkey: Pubkey,
39    /// The bls pubkey of the validator specified in the vote account
40    pub bls_pubkey: PopVerified<BLSPubkeyAffine>,
41    /// The stake of the validator
42    pub stake: NonZero<u64>,
43}
44
45/// Container to store a mapping from validator [`BLSPubkeyAffine`] to rank.
46///
47/// A validator with a smaller rank has a higher stake.
48/// Container also supports lookups from rank to [`BLSPubkeyStakeEntry`].
49#[derive(Clone, Debug)]
50#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
51pub struct BLSPubkeyToRankMap {
52    /// stores a mapping from the vote account pubkey to the node's rank.
53    vote_pubkey_to_rank: HashMap<Pubkey, u16>,
54    /// a mapping from rank to [`BLSPubkeyStakeEntry`].
55    sorted_pubkeys: Vec<BLSPubkeyStakeEntry>,
56    /// a mapping from the node identity pubkey to the node's rank.
57    node_pubkey_to_rank: HashMap<Pubkey, u16>,
58    /// Total stake delegated to this validator.
59    total_stake: NonZero<u64>,
60}
61
62// We cannot auto derive `AbiExample` for `BLSPubkeyToRankMap` because
63// the `BLSPubkeyAffine` type does not implement `AbiExample` or `Default`.
64#[cfg(feature = "frozen-abi")]
65impl solana_frozen_abi::abi_example::AbiExample for BLSPubkeyToRankMap {
66    fn example() -> Self {
67        Self {
68            vote_pubkey_to_rank: HashMap::new(),
69            sorted_pubkeys: Vec::new(),
70            total_stake: NonZero::new(1).unwrap(),
71            node_pubkey_to_rank: HashMap::new(),
72        }
73    }
74}
75
76pub(crate) fn bls_pubkey_compressed_bytes_to_bls_pubkey(
77    bls_pubkey_compressed_bytes: [u8; BLS_PUBLIC_KEY_COMPRESSED_SIZE],
78) -> Option<(BLSPubkeyCompressed, PopVerified<BLSPubkeyAffine>)> {
79    let bls_pubkey_compressed: BLSPubkeyCompressed =
80        wincode::deserialize(&bls_pubkey_compressed_bytes).ok()?;
81    let bls_pubkey_affine = BLSPubkeyAffine::try_from(bls_pubkey_compressed).ok()?;
82    // It is safe to use `new_unchecked` here because data coming from the vote
83    // state has already had its PoP verified.
84    let bls_pubkey_pop_verified = unsafe { PopVerified::new_unchecked(bls_pubkey_affine) };
85    Some((bls_pubkey_compressed, bls_pubkey_pop_verified))
86}
87
88impl BLSPubkeyToRankMap {
89    pub fn new(epoch_vote_accounts_hash_map: &VoteAccountsHashMap) -> Self {
90        let mut candidates = Vec::with_capacity(epoch_vote_accounts_hash_map.len());
91        let mut bls_pubkey_counts = HashMap::new();
92        let mut node_pubkey_counts = HashMap::new();
93        for (&vote_account_pubkey, (stake, account)) in epoch_vote_accounts_hash_map {
94            let Some(stake) = NonZero::new(*stake) else {
95                continue;
96            };
97            let node_pubkey = *account.vote_state_view().node_pubkey();
98            let Some((bls_pubkey_compressed, bls_pubkey)) = account
99                .vote_state_view()
100                .bls_pubkey_compressed()
101                .and_then(bls_pubkey_compressed_bytes_to_bls_pubkey)
102            else {
103                continue;
104            };
105            let entry = BLSPubkeyStakeEntry {
106                vote_account_pubkey,
107                node_pubkey,
108                bls_pubkey,
109                stake,
110            };
111            *bls_pubkey_counts.entry(bls_pubkey_compressed).or_insert(0) += 1;
112            *node_pubkey_counts.entry(node_pubkey).or_insert(0) += 1;
113            candidates.push((entry, bls_pubkey_compressed));
114        }
115        let mut keys_stake_entry_with_compressed: Vec<(BLSPubkeyStakeEntry, BLSPubkeyCompressed)> =
116            candidates
117                .into_iter()
118                .filter_map(|(entry, bls_pubkey_compressed)| {
119                    (bls_pubkey_counts[&bls_pubkey_compressed] == 1
120                        && node_pubkey_counts[&entry.node_pubkey] == 1)
121                        .then_some((entry, bls_pubkey_compressed))
122                })
123                .collect();
124        let total_stake = keys_stake_entry_with_compressed
125            .iter()
126            .fold(0u64, |stake, (entry, _)| {
127                stake.saturating_add(entry.stake.get())
128            });
129        let total_stake = NonZero::new(total_stake).expect("total stakes should not be 0");
130        keys_stake_entry_with_compressed.sort_by(
131            |(a_entry, a_pubkey_compressed), (b_entry, b_pubkey_compressed)| {
132                b_entry
133                    .stake
134                    .cmp(&a_entry.stake)
135                    .then(a_pubkey_compressed.cmp(b_pubkey_compressed))
136            },
137        );
138        let mut sorted_pubkeys = Vec::with_capacity(keys_stake_entry_with_compressed.len());
139        let mut vote_pubkey_to_rank_map =
140            HashMap::with_capacity(keys_stake_entry_with_compressed.len());
141        let mut node_pubkey_to_rank =
142            HashMap::with_capacity(keys_stake_entry_with_compressed.len());
143        for (rank, (entry, _bls_pubkey_compressed)) in
144            keys_stake_entry_with_compressed.into_iter().enumerate()
145        {
146            let rank = u16::try_from(rank).expect("BLS validator rank must fit in u16");
147            vote_pubkey_to_rank_map.insert(entry.vote_account_pubkey, rank);
148            node_pubkey_to_rank.insert(entry.node_pubkey, rank);
149            sorted_pubkeys.push(entry);
150        }
151        Self {
152            vote_pubkey_to_rank: vote_pubkey_to_rank_map,
153            sorted_pubkeys,
154            total_stake,
155            node_pubkey_to_rank,
156        }
157    }
158
159    pub fn is_empty(&self) -> bool {
160        self.sorted_pubkeys.is_empty()
161    }
162
163    pub fn len(&self) -> usize {
164        self.sorted_pubkeys.len()
165    }
166
167    pub fn total_stake(&self) -> NonZero<u64> {
168        self.total_stake
169    }
170
171    pub fn get_rank_for_vote_pubkey(&self, vote_pubkey: &Pubkey) -> Option<&u16> {
172        self.vote_pubkey_to_rank.get(vote_pubkey)
173    }
174
175    pub fn get_pubkey_stake_entry(&self, index: usize) -> Option<&BLSPubkeyStakeEntry> {
176        self.sorted_pubkeys.get(index)
177    }
178
179    /// Returns a node's rank and its canonical stake entry.
180    #[inline]
181    pub fn get_ranked_entry_for_node(
182        &self,
183        node_pubkey: &Pubkey,
184    ) -> Option<(u16, &BLSPubkeyStakeEntry)> {
185        let rank = *self.node_pubkey_to_rank.get(node_pubkey)?;
186        let entry = self.sorted_pubkeys.get(usize::from(rank))?;
187        Some((rank, entry))
188    }
189}
190
191#[cfg_attr(feature = "frozen-abi", derive(AbiExample, StableAbi, StableAbiSample))]
192#[derive(Clone, Serialize, Debug, Deserialize, Default, PartialEq, Eq, SchemaRead, SchemaWrite)]
193pub struct NodeVoteAccounts {
194    pub vote_accounts: Vec<Pubkey>,
195    pub total_stake: u64,
196}
197
198/// Simplified, intermediate representation of [`VersionedEpochStakes`]
199///
200/// Its bincode serializaiton format is identical as `VersionedEpochStakes`, but allows faster
201/// deserialization by ignoring serialized stake delegations entirely.
202#[cfg_attr(
203    feature = "frozen-abi",
204    derive(Serialize, SchemaWrite, AbiEnumVisitor, StableAbi, StableAbiSample)
205)]
206#[derive(Clone, Debug, Deserialize, SchemaRead)]
207#[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
208pub(crate) enum DeserializableVersionedEpochStakes {
209    Current {
210        #[cfg_attr(
211            feature = "frozen-abi",
212            stable_abi_sample(with = "stable_abi_sample_deserializable_epoch_stakes(rng)")
213        )]
214        stakes: DeserializableEpochStakes,
215        total_stake: u64,
216        node_id_to_vote_accounts: NodeIdToVoteAccounts,
217        epoch_authorized_voters: EpochAuthorizedVoters,
218    },
219}
220
221/// Draws in `EpochStakes` declaration order (`epoch` first, not wire order) so the sample matches
222/// the serialize side; the serializer injects empty `stake_delegations`/zero `unused`, mirrored here.
223#[cfg(feature = "frozen-abi")]
224fn stable_abi_sample_deserializable_epoch_stakes(
225    rng: &mut (impl solana_frozen_abi::rand::RngCore + ?Sized),
226) -> DeserializableEpochStakes {
227    use solana_frozen_abi::stable_abi::StableAbi;
228    let epoch = Epoch::random(rng);
229    let vote_accounts = VoteAccounts::random(rng);
230    let stake_history = StakeHistory::random(rng);
231    DeserializableEpochStakes {
232        vote_accounts,
233        _stake_delegations: Vec::new(),
234        _unused: 0,
235        epoch,
236        stake_history,
237    }
238}
239
240#[derive(Clone, Debug, Serialize, SchemaWrite)]
241#[cfg_attr(
242    feature = "frozen-abi",
243    derive(AbiExample, AbiEnumVisitor, StableAbi, StableAbiSample)
244)]
245#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
246pub enum VersionedEpochStakes {
247    Current {
248        stakes: EpochStakes,
249        /// Total stake in Lamports
250        total_stake: u64,
251        node_id_to_vote_accounts: Arc<NodeIdToVoteAccounts>,
252        epoch_authorized_voters: Arc<EpochAuthorizedVoters>,
253        #[cfg_attr(feature = "frozen-abi", stable_abi_sample(with = "Default::default()"))]
254        #[serde(skip)]
255        #[wincode(skip)]
256        bls_pubkey_to_rank_map: OnceLock<Arc<BLSPubkeyToRankMap>>,
257    },
258}
259
260impl From<DeserializableVersionedEpochStakes> for VersionedEpochStakes {
261    fn from(epoch_stakes: DeserializableVersionedEpochStakes) -> Self {
262        let DeserializableVersionedEpochStakes::Current {
263            stakes,
264            total_stake,
265            node_id_to_vote_accounts,
266            epoch_authorized_voters,
267        } = epoch_stakes;
268        Self::Current {
269            stakes: stakes.into(),
270            total_stake,
271            node_id_to_vote_accounts: Arc::new(node_id_to_vote_accounts),
272            epoch_authorized_voters: Arc::new(epoch_authorized_voters),
273            bls_pubkey_to_rank_map: OnceLock::new(),
274        }
275    }
276}
277
278impl VersionedEpochStakes {
279    #[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
280    pub(crate) fn new(stakes: SerdeStakesToStakeFormat, leader_schedule_epoch: Epoch) -> Self {
281        let stakes = EpochStakes::from(stakes);
282        let epoch_vote_accounts = stakes.vote_accounts();
283        let (total_stake, node_id_to_vote_accounts, epoch_authorized_voters) =
284            Self::parse_epoch_vote_accounts(epoch_vote_accounts.as_ref(), leader_schedule_epoch);
285        Self::Current {
286            stakes,
287            total_stake,
288            node_id_to_vote_accounts: Arc::new(node_id_to_vote_accounts),
289            epoch_authorized_voters: Arc::new(epoch_authorized_voters),
290            bls_pubkey_to_rank_map: OnceLock::new(),
291        }
292    }
293
294    #[cfg(feature = "dev-context-only-utils")]
295    pub fn new_for_tests(
296        vote_accounts_hash_map: VoteAccountsHashMap,
297        leader_schedule_epoch: Epoch,
298    ) -> Self {
299        Self::new(
300            SerdeStakesToStakeFormat::Account(crate::stakes::Stakes::new_for_tests(
301                0,
302                solana_vote::vote_account::VoteAccounts::from(Arc::new(vote_accounts_hash_map)),
303                imbl::HashMap::default(),
304            )),
305            leader_schedule_epoch,
306        )
307    }
308
309    pub fn stakes(&self) -> &EpochStakes {
310        match self {
311            Self::Current { stakes, .. } => stakes,
312        }
313    }
314
315    /// Returns the total stake in Lamports.
316    pub fn total_stake(&self) -> u64 {
317        match self {
318            Self::Current { total_stake, .. } => *total_stake,
319        }
320    }
321
322    #[cfg(feature = "dev-context-only-utils")]
323    pub fn set_total_stake(&mut self, total_stake: u64) {
324        match self {
325            Self::Current {
326                total_stake: total_stake_field,
327                ..
328            } => {
329                *total_stake_field = total_stake;
330            }
331        }
332    }
333
334    pub fn node_id_to_vote_accounts(&self) -> &Arc<NodeIdToVoteAccounts> {
335        match self {
336            Self::Current {
337                node_id_to_vote_accounts,
338                ..
339            } => node_id_to_vote_accounts,
340        }
341    }
342
343    pub fn node_id_to_stake(&self, node_id: &Pubkey) -> Option<u64> {
344        self.node_id_to_vote_accounts()
345            .get(node_id)
346            .map(|x| x.total_stake)
347    }
348
349    pub fn epoch_authorized_voters(&self) -> &Arc<EpochAuthorizedVoters> {
350        match self {
351            Self::Current {
352                epoch_authorized_voters,
353                ..
354            } => epoch_authorized_voters,
355        }
356    }
357
358    pub fn bls_pubkey_to_rank_map(&self) -> &Arc<BLSPubkeyToRankMap> {
359        match self {
360            Self::Current {
361                bls_pubkey_to_rank_map,
362                ..
363            } => bls_pubkey_to_rank_map.get_or_init(|| {
364                Arc::new(BLSPubkeyToRankMap::new(
365                    self.stakes().vote_accounts().as_ref(),
366                ))
367            }),
368        }
369    }
370
371    /// Returns the stake in Lamports for the given vote_account.
372    pub fn vote_account_stake(&self, vote_account: &Pubkey) -> u64 {
373        self.stakes()
374            .vote_accounts()
375            .get_delegated_stake(vote_account)
376    }
377
378    fn parse_epoch_vote_accounts(
379        epoch_vote_accounts: &VoteAccountsHashMap,
380        leader_schedule_epoch: Epoch,
381    ) -> (u64, NodeIdToVoteAccounts, EpochAuthorizedVoters) {
382        let mut node_id_to_vote_accounts: NodeIdToVoteAccounts = HashMap::new();
383        let mut epoch_authorized_voters: EpochAuthorizedVoters = HashMap::new();
384        let mut total_stake: u64 = 0;
385
386        for (key, (stake, account)) in epoch_vote_accounts.iter() {
387            total_stake += *stake;
388
389            if *stake == 0 {
390                continue;
391            }
392
393            let vote_state = account.vote_state_view();
394
395            if let Some(authorized_voter) = vote_state.get_authorized_voter(leader_schedule_epoch) {
396                let node_vote_accounts = node_id_to_vote_accounts
397                    .entry(*vote_state.node_pubkey())
398                    .or_default();
399
400                node_vote_accounts.total_stake += stake;
401                node_vote_accounts.vote_accounts.push(*key);
402
403                epoch_authorized_voters.insert(*key, *authorized_voter);
404            }
405        }
406
407        (
408            total_stake,
409            node_id_to_vote_accounts,
410            epoch_authorized_voters,
411        )
412    }
413}
414
415/// The current version of epoch stakes
416#[derive(Clone, Debug, Default)]
417#[cfg_attr(feature = "frozen-abi", derive(AbiExample, StableAbi, StableAbiSample))]
418#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
419pub struct EpochStakes {
420    epoch: Epoch,
421    vote_accounts: VoteAccounts,
422    stake_history: StakeHistory,
423}
424
425impl EpochStakes {
426    pub fn vote_accounts(&self) -> &VoteAccounts {
427        &self.vote_accounts
428    }
429    pub fn staked_nodes(&self) -> Arc<HashMap<Pubkey, u64>> {
430        self.vote_accounts.staked_nodes()
431    }
432}
433
434/// Customization of EpochStakes for snapshot serialization.
435///
436/// Needed because snapshots require additional fields no longer present in EpochStakes: the
437/// fields are reordered relative to `EpochStakes` and an always-empty `stake_delegations` list
438/// plus an unused `u64` are injected to match the historical wire format.
439#[derive(Serialize, SchemaWrite)]
440struct SerializableEpochStakes<'a> {
441    vote_accounts: &'a VoteAccounts,
442    stake_delegations: Vec<(Pubkey, Stake)>,
443    unused: u64,
444    epoch: Epoch,
445    stake_history: &'a StakeHistory,
446}
447
448impl<'a> From<&'a EpochStakes> for SerializableEpochStakes<'a> {
449    fn from(epoch_stakes: &'a EpochStakes) -> Self {
450        Self {
451            vote_accounts: &epoch_stakes.vote_accounts,
452            stake_delegations: Vec::new(), // do not serialize any stake delegations
453            unused: 0,
454            epoch: epoch_stakes.epoch,
455            stake_history: &epoch_stakes.stake_history,
456        }
457    }
458}
459
460impl Serialize for EpochStakes {
461    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
462    where
463        S: Serializer,
464    {
465        SerializableEpochStakes::from(self).serialize(serializer)
466    }
467}
468
469// Mirror the `Serialize` impl above for wincode by delegating to the same `SerializableEpochStakes`
470// wire layout, so the snapshot bytes match bincode.
471unsafe impl<C: wincode::config::Config> SchemaWrite<C> for EpochStakes {
472    type Src = Self;
473
474    fn size_of(src: &Self::Src) -> WriteResult<usize> {
475        <SerializableEpochStakes<'_> as SchemaWrite<C>>::size_of(&src.into())
476    }
477
478    fn write(writer: impl wincode::io::Writer, src: &Self::Src) -> WriteResult<()> {
479        <SerializableEpochStakes<'_> as SchemaWrite<C>>::write(writer, &src.into())
480    }
481}
482
483impl From<SerdeStakesToStakeFormat> for EpochStakes {
484    fn from(stakes: SerdeStakesToStakeFormat) -> Self {
485        let (epoch, vote_accounts, stake_history) = match stakes {
486            SerdeStakesToStakeFormat::Stake(stakes) => stakes.into_epoch_stakes_fields(),
487            SerdeStakesToStakeFormat::Account(stakes) => stakes.into_epoch_stakes_fields(),
488        };
489        Self {
490            epoch,
491            vote_accounts,
492            stake_history,
493        }
494    }
495}
496
497/// Customization of EpochStakes for snapshot deserialization.
498///
499/// Needed because snapshots contain additional fields no longer present in EpochStakes.
500// Sampling is overridden at the parent (`DeserializableVersionedEpochStakes::Current.stakes`), so
501// no StableAbi/StableAbiSample is needed here.
502#[cfg_attr(feature = "frozen-abi", derive(Serialize, SchemaWrite))]
503#[derive(Clone, Debug, Deserialize, SchemaRead)]
504#[cfg_attr(feature = "dev-context-only-utils", qualifiers(pub))]
505pub(crate) struct DeserializableEpochStakes {
506    vote_accounts: VoteAccounts,
507    // Read-and-discarded (always empty); the serialize side writes it empty too.
508    #[serde(deserialize_with = "deserialize_and_ignore_stake_delegations")]
509    #[wincode(with = "IgnoredStakeDelegations")]
510    _stake_delegations: Vec<(Pubkey, Stake)>,
511    _unused: u64,
512    epoch: Epoch,
513    stake_history: StakeHistory,
514}
515
516/// wincode `with` schema for the ignored stake delegations: reads and discards the wire's
517/// length-prefixed `Vec<(Pubkey, Stake)>`, yielding an empty `Vec`; writes it back through `Vec`
518/// (always empty, matching the serialize side).
519struct IgnoredStakeDelegations;
520
521unsafe impl<'de, C: Config> SchemaRead<'de, C> for IgnoredStakeDelegations {
522    type Dst = Vec<(Pubkey, Stake)>;
523
524    fn read(reader: impl Reader<'de>, dst: &mut MaybeUninit<Self::Dst>) -> ReadResult<()> {
525        <Vec<(Pubkey, Stake)> as SchemaRead<'de, C>>::get(reader)?;
526        dst.write(Vec::new());
527        Ok(())
528    }
529}
530
531#[cfg(feature = "frozen-abi")]
532unsafe impl<C: Config> SchemaWrite<C> for IgnoredStakeDelegations {
533    type Src = Vec<(Pubkey, Stake)>;
534
535    const TYPE_META: wincode::TypeMeta = <Vec<(Pubkey, Stake)> as SchemaWrite<C>>::TYPE_META;
536
537    fn size_of(src: &Self::Src) -> WriteResult<usize> {
538        <Vec<(Pubkey, Stake)> as SchemaWrite<C>>::size_of(src)
539    }
540
541    fn write(writer: impl wincode::io::Writer, src: &Self::Src) -> WriteResult<()> {
542        <Vec<(Pubkey, Stake)> as SchemaWrite<C>>::write(writer, src)
543    }
544}
545
546impl From<DeserializableEpochStakes> for EpochStakes {
547    fn from(stakes: DeserializableEpochStakes) -> Self {
548        let DeserializableEpochStakes {
549            vote_accounts,
550            _stake_delegations: _,
551            _unused: _,
552            epoch,
553            stake_history,
554        } = stakes;
555        Self {
556            epoch,
557            vote_accounts,
558            stake_history,
559        }
560    }
561}
562
563/// Snapshot epoch stakes contain delegations, but the main EpochStakes no longer uses them.
564/// This fn does custom deserialization to visit-and-ignore the delegations,
565/// avoiding the need to construct an expensive imbl::HashMap.
566fn deserialize_and_ignore_stake_delegations<'de, D>(
567    deserializer: D,
568) -> Result<Vec<(Pubkey, Stake)>, D::Error>
569where
570    D: Deserializer<'de>,
571{
572    struct IgnoredStakeDelegationsVisitor;
573
574    impl<'de> Visitor<'de> for IgnoredStakeDelegationsVisitor {
575        type Value = Vec<(Pubkey, Stake)>;
576
577        fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
578            formatter.write_str("a sequence of serialized stake delegations")
579        }
580
581        fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
582        where
583            A: SeqAccess<'de>,
584        {
585            while seq.next_element::<(Pubkey, Stake)>()?.is_some() {
586                // nothing to do here, ignore the delegations
587            }
588            Ok(Vec::new())
589        }
590    }
591
592    deserializer.deserialize_seq(IgnoredStakeDelegationsVisitor)
593}
594
595#[cfg(test)]
596pub(crate) mod tests {
597    use {
598        super::*,
599        crate::{
600            serde_snapshot::deserialize_wincode_from, stake_account::StakeAccount, stakes::Stakes,
601        },
602        solana_account::AccountSharedData,
603        solana_bls_signatures::keypair::Keypair as BLSKeypair,
604        solana_rent::Rent,
605        solana_vote::vote_account::VoteAccount,
606        solana_vote_program::vote_state::create_v4_account_with_authorized,
607        std::iter,
608        test_case::test_case,
609    };
610
611    struct VoteAccountInfo {
612        vote_account: Pubkey,
613        account: AccountSharedData,
614        authorized_voter: Pubkey,
615    }
616
617    fn new_vote_accounts(
618        num_nodes: usize,
619        num_vote_accounts_per_node: usize,
620        is_alpenglow: bool,
621    ) -> HashMap<Pubkey, Vec<VoteAccountInfo>> {
622        // Create some vote accounts for each pubkey
623        (0..num_nodes)
624            .map(|_| {
625                let node_id = solana_pubkey::new_rand();
626                (
627                    node_id,
628                    iter::repeat_with(|| {
629                        let authorized_voter = solana_pubkey::new_rand();
630                        let bls_pubkey_compressed: BLSPubkeyCompressed =
631                            (*BLSKeypair::new().public).into();
632                        let bls_pubkey_compressed_serialized =
633                            wincode::serialize(&bls_pubkey_compressed)
634                                .unwrap()
635                                .try_into()
636                                .unwrap();
637
638                        let bls_pubkey = if is_alpenglow {
639                            bls_pubkey_compressed_serialized
640                        } else {
641                            [0u8; BLS_PUBLIC_KEY_COMPRESSED_SIZE]
642                        };
643                        let account = create_v4_account_with_authorized(
644                            &node_id,
645                            &authorized_voter,
646                            bls_pubkey,
647                            &node_id,
648                            0,
649                            &node_id,
650                            0,
651                            &node_id,
652                            100,
653                        );
654                        VoteAccountInfo {
655                            vote_account: solana_pubkey::new_rand(),
656                            account,
657                            authorized_voter,
658                        }
659                    })
660                    .take(num_vote_accounts_per_node)
661                    .collect(),
662                )
663            })
664            .collect()
665    }
666
667    fn new_epoch_vote_accounts(
668        vote_accounts_map: &HashMap<Pubkey, Vec<VoteAccountInfo>>,
669        node_id_to_stake_fn: impl Fn(&Pubkey) -> u64,
670    ) -> VoteAccountsHashMap {
671        // Create and process the vote accounts
672        vote_accounts_map
673            .iter()
674            .flat_map(|(node_id, vote_accounts)| {
675                vote_accounts.iter().map(|v| {
676                    let vote_account = VoteAccount::try_from(v.account.clone()).unwrap();
677                    (v.vote_account, (node_id_to_stake_fn(node_id), vote_account))
678                })
679            })
680            .collect()
681    }
682
683    #[test_case(true; "alpenglow")]
684    #[test_case(false; "towerbft")]
685    fn test_parse_epoch_vote_accounts(is_alpenglow: bool) {
686        let stake_per_account = 100;
687        let num_vote_accounts_per_node = 2;
688        let num_nodes = 10;
689
690        let vote_accounts_map =
691            new_vote_accounts(num_nodes, num_vote_accounts_per_node, is_alpenglow);
692
693        let expected_authorized_voters: HashMap<_, _> = vote_accounts_map
694            .values()
695            .flat_map(|vote_accounts| {
696                vote_accounts
697                    .iter()
698                    .map(|v| (v.vote_account, v.authorized_voter))
699            })
700            .collect();
701
702        let expected_node_id_to_vote_accounts: HashMap<_, _> = vote_accounts_map
703            .iter()
704            .map(|(node_pubkey, vote_accounts)| {
705                let mut vote_accounts = vote_accounts
706                    .iter()
707                    .map(|v| v.vote_account)
708                    .collect::<Vec<_>>();
709                vote_accounts.sort();
710                let node_vote_accounts = NodeVoteAccounts {
711                    vote_accounts,
712                    total_stake: stake_per_account * num_vote_accounts_per_node as u64,
713                };
714                (*node_pubkey, node_vote_accounts)
715            })
716            .collect();
717
718        let epoch_vote_accounts =
719            new_epoch_vote_accounts(&vote_accounts_map, |_| stake_per_account);
720
721        let (total_stake, mut node_id_to_vote_accounts, epoch_authorized_voters) =
722            VersionedEpochStakes::parse_epoch_vote_accounts(&epoch_vote_accounts, 0);
723
724        // Verify the results
725        node_id_to_vote_accounts
726            .iter_mut()
727            .for_each(|(_, node_vote_accounts)| node_vote_accounts.vote_accounts.sort());
728
729        assert!(
730            node_id_to_vote_accounts.len() == expected_node_id_to_vote_accounts.len()
731                && node_id_to_vote_accounts
732                    .iter()
733                    .all(|(k, v)| expected_node_id_to_vote_accounts.get(k).unwrap() == v)
734        );
735        assert!(
736            epoch_authorized_voters.len() == expected_authorized_voters.len()
737                && epoch_authorized_voters
738                    .iter()
739                    .all(|(k, v)| expected_authorized_voters.get(k).unwrap() == v)
740        );
741        assert_eq!(
742            total_stake,
743            num_nodes as u64 * num_vote_accounts_per_node as u64 * 100
744        );
745    }
746
747    #[test_case(true; "alpenglow")]
748    #[test_case(false; "towerbft")]
749    fn test_node_id_to_stake(is_alpenglow: bool) {
750        let num_nodes = 10;
751        let num_vote_accounts_per_node = 2;
752
753        let vote_accounts_map =
754            new_vote_accounts(num_nodes, num_vote_accounts_per_node, is_alpenglow);
755        let node_id_to_stake_map = vote_accounts_map
756            .keys()
757            .enumerate()
758            .map(|(index, node_id)| (*node_id, ((index + 1) * 100) as u64))
759            .collect::<HashMap<_, _>>();
760        let epoch_vote_accounts = new_epoch_vote_accounts(&vote_accounts_map, |node_id| {
761            *node_id_to_stake_map.get(node_id).unwrap()
762        });
763        let epoch_stakes = VersionedEpochStakes::new_for_tests(epoch_vote_accounts, 0);
764
765        assert_eq!(epoch_stakes.total_stake(), 11000);
766        for (node_id, stake) in node_id_to_stake_map.iter() {
767            assert_eq!(
768                epoch_stakes.node_id_to_stake(node_id),
769                Some(*stake * num_vote_accounts_per_node as u64)
770            );
771        }
772    }
773
774    #[test]
775    fn test_bls_pubkey_rank_map() {
776        agave_logger::setup();
777        let num_nodes = 10;
778        let num_vote_accounts = num_nodes;
779
780        let vote_accounts_map = new_vote_accounts(num_nodes, 1, true);
781        let node_id_to_stake_map = vote_accounts_map
782            .keys()
783            .enumerate()
784            .map(|(index, node_id)| (*node_id, ((index + 1) * 100) as u64))
785            .collect::<HashMap<_, _>>();
786        let epoch_vote_accounts = new_epoch_vote_accounts(&vote_accounts_map, |node_id| {
787            *node_id_to_stake_map.get(node_id).unwrap()
788        });
789        let epoch_stakes = VersionedEpochStakes::new_for_tests(epoch_vote_accounts.clone(), 0);
790        let bls_pubkey_to_rank_map = epoch_stakes.bls_pubkey_to_rank_map();
791        let expected_num_vote_accounts = num_vote_accounts;
792        assert_eq!(bls_pubkey_to_rank_map.len(), expected_num_vote_accounts);
793        let expected_total_stake = epoch_stakes.total_stake();
794        assert_eq!(
795            bls_pubkey_to_rank_map.total_stake().get(),
796            expected_total_stake
797        );
798        for expected_rank in 0..bls_pubkey_to_rank_map.len() {
799            let expected_entry = bls_pubkey_to_rank_map
800                .get_pubkey_stake_entry(expected_rank)
801                .unwrap();
802            let (rank, entry) = bls_pubkey_to_rank_map
803                .get_ranked_entry_for_node(&expected_entry.node_pubkey)
804                .unwrap();
805            assert_eq!(usize::from(rank), expected_rank);
806            assert!(std::ptr::eq(entry, expected_entry));
807            assert_eq!(
808                bls_pubkey_to_rank_map
809                    .get_rank_for_vote_pubkey(&entry.vote_account_pubkey)
810                    .copied(),
811                Some(rank)
812            );
813        }
814        assert!(
815            bls_pubkey_to_rank_map
816                .get_ranked_entry_for_node(&Pubkey::new_unique())
817                .is_none()
818        );
819
820        // Convert it to versioned and back, we should get the same rank map
821        let mut bank_epoch_stakes = HashMap::new();
822        bank_epoch_stakes.insert(0, epoch_stakes.clone());
823        let epoch_stakes = bank_epoch_stakes
824            .get(&0)
825            .expect("Epoch stakes should exist");
826        let bls_pubkey_to_rank_map2 = epoch_stakes.bls_pubkey_to_rank_map();
827        assert_eq!(bls_pubkey_to_rank_map2, bls_pubkey_to_rank_map);
828    }
829
830    #[test]
831    #[should_panic(expected = "total stakes should not be 0")]
832    fn test_multiple_vote_accounts_panics() {
833        agave_logger::setup();
834        let num_nodes = 10;
835
836        let vote_accounts_map = new_vote_accounts(num_nodes, 2, true);
837        let node_id_to_stake_map = vote_accounts_map
838            .keys()
839            .enumerate()
840            .map(|(index, node_id)| (*node_id, ((index + 1) * 100) as u64))
841            .collect::<HashMap<_, _>>();
842        let epoch_vote_accounts = new_epoch_vote_accounts(&vote_accounts_map, |node_id| {
843            *node_id_to_stake_map.get(node_id).unwrap()
844        });
845        let epoch_stakes = VersionedEpochStakes::new_for_tests(epoch_vote_accounts.clone(), 0);
846        epoch_stakes.bls_pubkey_to_rank_map();
847    }
848
849    #[test]
850    fn test_bls_pubkey_rank_map_excludes_duplicate_bls_and_identity() {
851        let new_bls_pubkey = || {
852            let compressed: BLSPubkeyCompressed = (*BLSKeypair::new().public).into();
853            wincode::serialize(&compressed).unwrap().try_into().unwrap()
854        };
855
856        let duplicate_bls_pubkey_serialized = new_bls_pubkey();
857        let duplicate_node_bls_pubkey_serialized = new_bls_pubkey();
858        let duplicate_node_bls_pubkey_serialized_2 = new_bls_pubkey();
859        let shared_voter_bls_pubkey_serialized = new_bls_pubkey();
860        let shared_voter_bls_pubkey_serialized_2 = new_bls_pubkey();
861        let unique_bls_pubkey_serialized = new_bls_pubkey();
862
863        let duplicate_bls_vote_pubkey = Pubkey::new_unique();
864        let duplicate_bls_vote_pubkey_2 = Pubkey::new_unique();
865        let duplicate_node_vote_pubkey = Pubkey::new_unique();
866        let duplicate_node_vote_pubkey_2 = Pubkey::new_unique();
867        let shared_voter_vote_pubkey = Pubkey::new_unique();
868        let shared_voter_vote_pubkey_2 = Pubkey::new_unique();
869        let unique_vote_pubkey = Pubkey::new_unique();
870
871        let duplicate_node_pubkey = Pubkey::new_unique();
872        let shared_authorized_voter = Pubkey::new_unique();
873        let shared_voter_node_pubkey = Pubkey::new_unique();
874        let shared_voter_node_pubkey_2 = Pubkey::new_unique();
875        let unique_node_pubkey = Pubkey::new_unique();
876        let unique_voter = Pubkey::new_unique();
877
878        let account = |node_pubkey, authorized_voter, bls_pubkey| {
879            VoteAccount::try_from(create_v4_account_with_authorized(
880                &node_pubkey,
881                &authorized_voter,
882                bls_pubkey,
883                &node_pubkey,
884                0,
885                &node_pubkey,
886                0,
887                &node_pubkey,
888                100,
889            ))
890            .unwrap()
891        };
892        let epoch_vote_accounts = VoteAccountsHashMap::from([
893            (
894                duplicate_bls_vote_pubkey,
895                (
896                    100,
897                    account(
898                        Pubkey::new_unique(),
899                        Pubkey::new_unique(),
900                        duplicate_bls_pubkey_serialized,
901                    ),
902                ),
903            ),
904            (
905                duplicate_bls_vote_pubkey_2,
906                (
907                    100,
908                    account(
909                        Pubkey::new_unique(),
910                        Pubkey::new_unique(),
911                        duplicate_bls_pubkey_serialized,
912                    ),
913                ),
914            ),
915            (
916                duplicate_node_vote_pubkey,
917                (
918                    100,
919                    account(
920                        duplicate_node_pubkey,
921                        Pubkey::new_unique(),
922                        duplicate_node_bls_pubkey_serialized,
923                    ),
924                ),
925            ),
926            (
927                duplicate_node_vote_pubkey_2,
928                (
929                    100,
930                    account(
931                        duplicate_node_pubkey,
932                        Pubkey::new_unique(),
933                        duplicate_node_bls_pubkey_serialized_2,
934                    ),
935                ),
936            ),
937            (
938                shared_voter_vote_pubkey,
939                (
940                    100,
941                    account(
942                        shared_voter_node_pubkey,
943                        shared_authorized_voter,
944                        shared_voter_bls_pubkey_serialized,
945                    ),
946                ),
947            ),
948            (
949                shared_voter_vote_pubkey_2,
950                (
951                    100,
952                    account(
953                        shared_voter_node_pubkey_2,
954                        shared_authorized_voter,
955                        shared_voter_bls_pubkey_serialized_2,
956                    ),
957                ),
958            ),
959            (
960                unique_vote_pubkey,
961                (
962                    50,
963                    account(
964                        unique_node_pubkey,
965                        unique_voter,
966                        unique_bls_pubkey_serialized,
967                    ),
968                ),
969            ),
970        ]);
971
972        let rank_map = BLSPubkeyToRankMap::new(&epoch_vote_accounts);
973
974        assert_eq!(rank_map.len(), 3);
975        assert_eq!(rank_map.total_stake().get(), 250);
976        for vote_pubkey in [
977            duplicate_bls_vote_pubkey,
978            duplicate_bls_vote_pubkey_2,
979            duplicate_node_vote_pubkey,
980            duplicate_node_vote_pubkey_2,
981        ] {
982            assert!(rank_map.get_rank_for_vote_pubkey(&vote_pubkey).is_none());
983        }
984        assert!(
985            rank_map
986                .get_ranked_entry_for_node(&duplicate_node_pubkey)
987                .is_none()
988        );
989        for node_pubkey in [
990            shared_voter_node_pubkey,
991            shared_voter_node_pubkey_2,
992            unique_node_pubkey,
993        ] {
994            let (rank, entry) = rank_map.get_ranked_entry_for_node(&node_pubkey).unwrap();
995            assert_eq!(entry.node_pubkey, node_pubkey);
996            assert_eq!(
997                rank_map
998                    .get_rank_for_vote_pubkey(&entry.vote_account_pubkey)
999                    .copied(),
1000                Some(rank)
1001            );
1002        }
1003        assert!(rank_map.get_pubkey_stake_entry(rank_map.len()).is_none());
1004    }
1005
1006    #[test]
1007    fn test_versioned_epoch_stakes_does_not_serialize_delegations() {
1008        // test-only types to get the serialized EpochStakes that still have stake delegations
1009        #[derive(Deserialize)]
1010        enum SerializedVersionedEpochStakes {
1011            Current {
1012                stakes: SerializedEpochStakes,
1013                total_stake: u64,
1014                node_id_to_vote_accounts: NodeIdToVoteAccounts,
1015                epoch_authorized_voters: EpochAuthorizedVoters,
1016            },
1017        }
1018        #[derive(Deserialize)]
1019        struct SerializedEpochStakes {
1020            vote_accounts: VoteAccounts,
1021            stake_delegations: Vec<(Pubkey, Stake)>,
1022            unused: u64,
1023            epoch: Epoch,
1024            stake_history: StakeHistory,
1025        }
1026
1027        let epoch = 42;
1028        let delegated_amount = 456_789;
1029        let rent = Rent::default();
1030        let ((vote_pubkey, vote_account), (stake_pubkey, stake_account)) =
1031            crate::stakes::tests::create_staked_node_accounts(123, &rent);
1032        let vote_account = VoteAccount::try_from(vote_account).unwrap();
1033        let vote_accounts = VoteAccounts::from(Arc::new(HashMap::from([(
1034            vote_pubkey,
1035            (delegated_amount, vote_account),
1036        )])));
1037        let stake_account = StakeAccount::try_from(stake_account).unwrap();
1038        let stakes = Stakes::new_for_tests(
1039            epoch,
1040            vote_accounts,
1041            imbl::HashMap::from_iter([(stake_pubkey, stake_account)]),
1042        );
1043
1044        // ensure stake delegations start off *not* empty
1045        assert!(!stakes.stake_delegations().is_empty());
1046
1047        let epoch_stakes = VersionedEpochStakes::new(SerdeStakesToStakeFormat::Account(stakes), 0);
1048
1049        assert_eq!(
1050            epoch_stakes
1051                .stakes()
1052                .vote_accounts()
1053                .get_delegated_stake(&vote_pubkey),
1054            delegated_amount,
1055        );
1056
1057        let serialized_bytes = bincode::serialize(&epoch_stakes).unwrap();
1058        let serialized_epoch_stakes: SerializedVersionedEpochStakes =
1059            bincode::deserialize(&serialized_bytes).unwrap();
1060        match serialized_epoch_stakes {
1061            SerializedVersionedEpochStakes::Current {
1062                stakes,
1063                total_stake,
1064                node_id_to_vote_accounts,
1065                epoch_authorized_voters,
1066            } => {
1067                assert_eq!(
1068                    stakes.vote_accounts.get_delegated_stake(&vote_pubkey),
1069                    delegated_amount,
1070                );
1071                assert!(stakes.stake_delegations.is_empty()); // delegations are *not* serialized
1072                assert_eq!(stakes.unused, 0);
1073                assert_eq!(stakes.epoch, epoch);
1074                assert_eq!(stakes.stake_history, StakeHistory::default());
1075                assert_eq!(total_stake, delegated_amount);
1076                assert_eq!(node_id_to_vote_accounts.len(), 1);
1077                assert_eq!(epoch_authorized_voters.len(), 1);
1078            }
1079        }
1080
1081        let deserialized_epoch_stakes: VersionedEpochStakes =
1082            deserialize_wincode_from::<_, DeserializableVersionedEpochStakes>(
1083                std::io::Cursor::new(&serialized_bytes),
1084            )
1085            .unwrap()
1086            .into();
1087        assert_eq!(
1088            deserialized_epoch_stakes
1089                .stakes()
1090                .vote_accounts()
1091                .get_delegated_stake(&vote_pubkey),
1092            delegated_amount,
1093        );
1094    }
1095}