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