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#[derive(Clone, Debug)]
32#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
33pub struct BLSPubkeyStakeEntry {
34 pub vote_account_pubkey: Pubkey,
36 pub node_pubkey: Pubkey,
38 pub bls_pubkey: PopVerified<BLSPubkeyAffine>,
40 pub stake: NonZero<u64>,
42}
43
44#[derive(Clone, Debug)]
49#[cfg_attr(feature = "dev-context-only-utils", derive(PartialEq))]
50pub struct BLSPubkeyToRankMap {
51 vote_pubkey_to_rank: HashMap<Pubkey, u16>,
53 sorted_pubkeys: Vec<BLSPubkeyStakeEntry>,
55 node_pubkey_to_rank: HashMap<Pubkey, u16>,
57 total_stake: NonZero<u64>,
59}
60
61#[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 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 #[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#[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#[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: 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 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 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#[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#[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(), 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
468unsafe 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#[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 #[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
532fn 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 }
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 (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 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 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 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 #[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 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()); 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}